Revise paste Мірошниченко_КМЗПМ_ЛР№1

from dataclasses import dataclass 
import numpy as np
import matplotlib.pyplot as plt

@dataclass
class TimeGrid:
    t_min: float = 1e-3  # мінімальний час
    t_max: float = 1e4  # максимальний час
    dt: float = 0.002  # крок дискретизації часу

def values(self) -> np.ndarray: 
        span = self.t_max - self.t_min
        num = int(np.floor(span / self.dt)) + 1
        return np.logspace(np.log10(self.t_min), np.log10(self.t_max), num)

def msd_free_times(D: float, N: int, t: np.ndarray, seed: int | None = 0):
    rng = np.random.default_rng(seed)
    x = np.zeros(N)
    msd = np.zeros_like(t)
    t_prev = 0.0
    for k, tk in enumerate(t):
        dt = tk - t_prev
        x += np.sqrt(D * dt) * rng.standard_normal(N)
        msd[k] = np.mean(x * x) - np.mean(x) ** 2
        t_prev = tk
    return t, msd

@dataclass
class CompositeParams:
    D: float = 1.0
    gamma0: float = 1.0
    gamma: float = 1.0
    k: float = 1.0
    M: int = 0
    N: int = 1000
    seed: int | None = 0

def msd_composite_times(params: CompositeParams, t: np.ndarray):
    rng = np.random.default_rng(params.seed)
    N, M = params.N, params.M
    x = np.zeros(N)
    sat = np.zeros((N, M)) if M > 0 else np.zeros((N, 0))
    msd = np.zeros_like(t)
    t_prev = 0.0
    for k, tk in enumerate(t):
        dt = tk - t_prev
        if M > 0:
            diff = x[:, None] - sat
            force_on_x = -params.k * np.sum(diff, axis=1)
            x += (dt / params.gamma0) * force_on_x + (np.sqrt(params.D * dt) / params.gamma0) * rng.standard_normal(N)
            sat += (dt / params.gamma) * (params.k * diff) + (np.sqrt(params.D * dt) / params.gamma) * rng.standard_normal((N, M))
        else:
            x += (np.sqrt(params.D * dt) / params.gamma0) * rng.standard_normal(N)
        msd[k] = np.mean(x * x) - np.mean(x) ** 2
        t_prev = tk
    return t, msd

def style_log_axes(ax: plt.Axes, xlabel: str = "time", ylabel: str = r"⟨(δx)²⟩") -> None:
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.grid(True, which="both", ls=":", lw=0.7, alpha=0.6)
    ax.set_xlim(1e-3, 1e4)
    ax.set_ylim(1e-4, 1e4)

def show_fig37() -> None:
    grid = TimeGrid(dt=0.002)
    t_line = grid.values()
    t1, msd1 = msd_free_times(D=0.1, N=200, t=t_line, seed=1)
    _, msd2 = msd_free_times(D=1.0, N=200, t=t_line, seed=2)
    idx = np.linspace(0, len(t_line) - 1, 60).astype(int)
    t_mark = t_line[idx]
    msd1_mark = msd1[idx]
    msd2_mark = msd2[idx]
    fig, ax = plt.subplots(figsize=(8, 6), dpi=140)
    style_log_axes(ax)

ax.plot(t_line, msd1, '-', lw=1.5, color='royalblue', label='D = 0.1')
    ax.plot(t_line, msd2, '-', lw=1.5, color='red', label='D = 1.0')

ax.plot(t_mark, msd1_mark, 'o', ms=3, color='royalblue')
    ax.plot(t_mark, msd2_mark, 'o', ms=3, color='red')

tt = np.logspace(-3, 4, 50)
    ax.plot(tt, tt, color='0.3', lw=1.0)
    ax.plot(tt, 0.1 * tt, color='0.3', lw=1.0)
    ax.text(3, 2, 't', fontsize=9)
    ax.text(3, 0.2, '0.1t', fontsize=9)
    ax.legend(frameon=False, loc='upper left')
    ax.set_title('Середньоквадратичне зміщення: вільна частинка')
    plt.tight_layout()
    plt.show()

def show_fig() -> None:
    grid = TimeGrid(dt=0.002)
    base = CompositeParams(D=1.0, gamma0=1.0, gamma=1.0, k=1.0, N=1000)
    p0 = CompositeParams(**{**base.__dict__, 'M': 0, 'seed': 3})
    p10 = CompositeParams(**{**base.__dict__, 'M': 2, 'seed': 4})
    t_line = grid.values()
    t0, msd0 = msd_composite_times(p0, t_line)
    _, msd10 = msd_composite_times(p10, t_line)
    idx = np.linspace(0, len(t_line) - 1, 60).astype(int)
    t_mark = t_line[idx]
    msd0_mark = msd0[idx]
    msd10_mark = msd10[idx]
    fig, ax = plt.subplots(figsize=(8, 6), dpi=140)
    style_log_axes(ax)

ax.plot(t_line, msd0, '-', lw=1.5, color='crimson', label='M = 0, N = 1000')
    ax.plot(t_line, msd10, '-', lw=1.5, color='darkgreen', label='M = 2, N = 1000')

ax.scatter(t_mark, msd0_mark, s=5, color='crimson')
    ax.scatter(t_mark, msd10_mark, s=5,color='darkgreen')
    tt = np.logspace(-3, 4, 50)
    ax.plot(tt, tt, color='0.3', lw=1.0)
    ax.set_title('MSD композитної частинки')
    ax.legend(frameon=False, loc='upper left', title='$\\gamma_0 = 1$\n$\\gamma = 1$\n$k = 1$')
    plt.tight_layout()
    plt.show()

def show_fig41() -> None:
    grid = TimeGrid(dt=0.002)
    base = CompositeParams(D=25.0, gamma0=5.0, gamma=1.0, k=1.0, N=1000)
    p0 = CompositeParams(**{**base.__dict__, 'M': 0, 'seed': 3})
    p10 = CompositeParams(**{**base.__dict__, 'M': 10, 'seed': 4})
    t_line = grid.values()
    t0, msd0 = msd_composite_times(p0, t_line)
    _, msd10 = msd_composite_times(p10, t_line)
    idx = np.linspace(0, len(t_line) - 1, 60).astype(int)
    t_mark = t_line[idx]
    msd0_mark = msd0[idx]
    msd10_mark = msd10[idx]
    fig, ax = plt.subplots(figsize=(8, 6), dpi=140)
    style_log_axes(ax)

ax.plot(t_line, msd0, '-', lw=1.5, color='crimson', label='M = 0, N = 1000')
    ax.plot(t_line, msd10, '-', lw=1.5, color='darkgreen', label='M = 10, N = 1000')

ax.scatter(t_mark, msd0_mark, s=5, color='crimson')
    ax.scatter(t_mark, msd10_mark, s=5,color='darkgreen')
    tt = np.logspace(-3, 4, 50)
    ax.plot(tt, tt, color='0.3', lw=1.0)
    ax.set_title('MSD композитної частинки')
    ax.legend(frameon=False, loc='upper left', title='$\\gamma_0 = 5$\n$\\gamma = 1$\n$k = 1$')
    plt.tight_layout()
    plt.show()

def main() -> None:
    show_fig37()
    show_fig()
    show_fig41()

if __name__ == "__main__":
    main()