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

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




@dataclass
class Params:
    N: int = 77
    m: float = 0.01
    r: float = 0.005
    q_val: float = 100.0
    k_val: float = 0.0

    A: float = 50.0
    T_period: float = 0.3

    T_max: float = 1.0
    dt: float = 1e-4

    n_snapshots: int = 14
    front_eps: float = 1e-4


def build_q_array(N: int, q_val: float | np.ndarray) -> np.ndarray:
    if np.isscalar(q_val):
        q = np.full(N + 1, float(q_val), dtype=float)
    else:
        q = np.asarray(q_val, dtype=float)
        if q.shape != (N + 1,):
            raise ValueError("q array must have shape (N+1,)")
    return q


def compute_link_forces(x: np.ndarray, q: np.ndarray) -> np.ndarray:
    N = x.size
    x_pad = np.empty(N + 2, dtype=float)
    x_pad[0] = 0.0
    x_pad[1:-1] = x
    x_pad[-1] = 0.0

    left = x_pad[0:N]
    cen = x_pad[1:N+1]
    right = x_pad[2:N+2]

    qL = q[0:N]
    qR = q[1:N+1]

    F = qL * (left - cen) + qR * (right - cen)
    return F


def external_force(t: float, p: Params) -> np.ndarray:
    f = np.zeros(p.N, dtype=float)
    omega = 2.0 * np.pi / p.T_period
    if t <= 0.5 * p.T_period:
        f[0] = p.A * np.sin(omega * t)
    return f


def total_energy(x: np.ndarray, v: np.ndarray, p: Params, q: np.ndarray) -> float:
    T = 0.5 * p.m * float(np.dot(v, v))
    x_pad = np.empty(p.N + 2, dtype=float)
    x_pad[0] = 0.0
    x_pad[1:-1] = x
    x_pad[-1] = 0.0
    diffs = x_pad[1:] - x_pad[:-1]
    V_links = 0.5 * float(np.dot(q, diffs * diffs))
    V_self = 0.5 * p.k_val * float(np.dot(x, x))
    return T + V_links + V_self




def simulate(p: Params):
    N = p.N
    q = build_q_array(N, p.q_val)

    steps = int(np.floor(p.T_max / p.dt + 1e-15))
    t_grid = np.arange(steps + 1, dtype=float) * p.dt

    x = np.zeros(N, dtype=float)
    v = np.zeros(N, dtype=float)

    E_log = np.empty(steps + 1, dtype=float)
    E_log[0] = total_energy(x, v, p, q)

    t_front: Optional[float] = None



    snap_times = np.linspace(0.0, p.T_max, p.n_snapshots + 1)[1:]
    snap_dict: dict[float, np.ndarray] = {}

    for n in range(steps):
        t = t_grid[n]
        F = compute_link_forces(x, q)
        f_ext = external_force(t, p)
        a = (F - p.r * v - p.k_val * x + f_ext) / p.m

        v += a * p.dt
        x += v * p.dt

        E_log[n + 1] = total_energy(x, v, p, q)

        if t_front is None and abs(x[-1]) > p.front_eps:
            t_front = t


        for ts in snap_times:
            if ts not in snap_dict and (t <= ts < t + p.dt):
                snap_dict[ts] = x.copy()


    if p.T_max not in snap_dict:
        snap_dict[p.T_max] = x.copy()

    print(f"T_max={p.T_max:g} с, dt={p.dt:g}, кроків={steps}")
    print(f"N={p.N}, m={p.m}, r={p.r}, q={p.q_val}, k={p.k_val}, A={p.A}, T={p.T_period}")
    if t_front is not None:
        print(f"Час проходження фронту до останнього вузла: ≈ {t_front:.4f} с")
    else:
        print("Фронт не досяг правого краю за відведений час.")


    fig1, ax1 = plt.subplots(figsize=(7.6, 4.6))
    ax1.plot(t_grid, E_log)
    ax1.set_xlabel('Час, с')
    ax1.set_ylabel('Повна енергія E(t)')
    ax1.set_title('Зміна повної енергії системи у часі')
    ax1.grid(True, ls=':')
    fig1.tight_layout()


    if snap_dict:
        snaps_sorted = sorted(snap_dict.items(), key=lambda kv: kv[0])
        n = len(snaps_sorted)
        cols = 4 if n >= 13 else 3
        rows = int(np.ceil(n / cols))
        fig2, axes = plt.subplots(rows, cols, figsize=(4.9 * cols, 3.4 * rows), sharey=True)
        axes = np.atleast_1d(axes).ravel()
        xi = np.arange(1, N + 1)
        for ax, (ts, xs) in zip(axes, snaps_sorted):
            ax.plot(xi, xs, marker='o', lw=1)
            ax.axhline(0.0, color='0.6', ls='--', lw=0.8)
            ax.set_xlim(1, N)
            ax.set_title(f"t = {ts:.3f} с")
            ax.set_xlabel('Номер осцилятора')
            ax.set_ylabel('Зсув')
            ax.grid(True, ls=':')
            ax.set_xticks(np.arange(1, N + 1, 5))
        for ax in axes[n:]:
            ax.axis('off')
        fig2.suptitle('Знімки профілю зсувів уздовж ланцюжка')
        fig2.tight_layout(rect=[0, 0, 1, 0.96])

    plt.show()


if __name__ == '__main__':
    params = Params()
    simulate(params)