physicskit 0.1.0__py3-none-any.whl

physicskit/classical/tests/test_visualizers.py

@@ -0,0 +1,172 @@
+"""Tests for physicskit.classical.visualizers, previously untested (the whole package
+had only manual, un-persisted smoke checks during development). These
+run headlessly (see conftest.py's Agg backend) and check both "does it
+run without raising" and, where practical, that the output is actually
+correct (e.g. Poincare-section crossings at known locations), not just
+that a figure object came back.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from matplotlib.animation import PillowWriter
+
+from physicskit.classical.systems.chains import FPUTChain
+from physicskit.classical.systems.hamiltonian import HenonHeilesSystem
+from physicskit.classical.systems.lagrangian import DoublePendulum, ElasticPendulum
+from physicskit.classical.systems.newtonian import KeplerSystem
+from physicskit.classical.systems.rotations import EulersDisk, EulerTop, Rattleback
+from physicskit.classical.visualizers.animations import (
+    SideBySideAnimator,
+    animate_elastic_pendulum,
+    animate_eulers_disk,
+    animate_rattleback,
+    animate_rigid_body_tumble,
+    orbit_trace_animation,
+    pendulum_animation,
+)
+from physicskit.classical.visualizers.interactive import interactive_orbit, interactive_so3_momentum_sphere
+from physicskit.classical.visualizers.modal_analysis import animate_modal_energies, plot_modal_energy_bars
+from physicskit.classical.visualizers.phase_space import (
+    plot_phase_portrait,
+    plot_phase_swarm,
+    plot_poincare_section,
+    plot_so3_momentum_sphere,
+    poincare_section,
+)
+
+
+def test_poincare_section_finds_known_crossings():
+    """Synthetic q(t) = sin(t), p(t) = cos(t): q crosses zero going
+    upward (direction=+1, p>0) exactly at t = 0, 2*pi, 4*pi, ..."""
+    t = np.linspace(0, 4 * np.pi, 20_000)
+    q = np.sin(t)
+    p = np.cos(t)
+
+    cq, cp = poincare_section(q, p, section_q=q, section_p=p, value=0.0, direction=1)
+
+    assert len(cq) == 2  # crossings at t=0 (endpoint, not detected by the interior scan) is excluded; t=2pi, 4pi found
+    assert np.allclose(cq, 0.0, atol=1e-3)
+    assert np.allclose(cp, 1.0, atol=1e-3)
+
+
+def test_plot_phase_portrait_1d_and_2d():
+    q1 = np.linspace(0, 1, 50)
+    p1 = np.cos(q1)
+    ax = plot_phase_portrait(q1, p1)
+    assert ax is not None
+
+    q2 = np.column_stack([q1, q1 * 2])
+    p2 = np.column_stack([p1, -p1])
+    ax2 = plot_phase_portrait(q2, p2, label="test")
+    assert len(ax2.lines) >= 2
+
+
+def test_plot_phase_swarm_and_poincare_wrapper():
+    q = np.random.default_rng(0).normal(size=200)
+    p = np.random.default_rng(1).normal(size=200)
+    ax = plot_phase_swarm(q, p)
+    assert ax is not None
+
+    hh = HenonHeilesSystem(np.array([0.0, 0.3]), np.array([0.3, 0.0]))
+    result = hh.integrate((0, 100), dt=0.01, method="yoshida4")
+    ax2 = plot_poincare_section(result.q[:, 0], result.p[:, 0], result.q[:, 1], result.p[:, 1], value=0.0, direction=1)
+    assert ax2 is not None
+
+
+def test_plot_so3_momentum_sphere():
+    top = EulerTop([0.01, 1.0, 0.01], I1=1.0, I2=2.0, I3=3.0)
+    result = top.integrate((0, 10), dt=1e-3, method="implicit_midpoint")
+    ax = plot_so3_momentum_sphere(result.y[:, :3], I1=1.0, I2=2.0, I3=3.0)
+    assert ax is not None
+
+
+@pytest.mark.slow
+def test_modal_energy_bars_and_animation(tmp_path):
+    chain = FPUTChain(n=16, beta=0.5, mode=1, amplitude=0.5)
+    result = chain.integrate((0, 5), dt=0.05, method="yoshida4")  # short: this only needs a few frames
+
+    ax = plot_modal_energy_bars(chain, result.q[0], result.p[0])
+    assert len(ax.patches) == 16
+
+    anim = animate_modal_energies(chain, result, stride=10)
+    out = tmp_path / "modal_energy.gif"
+    anim.save(out, writer=PillowWriter(fps=10))
+    assert out.exists() and out.stat().st_size > 0
+
+
+@pytest.mark.slow
+def test_side_by_side_animator_pendulum_and_orbit(tmp_path):
+    dp = DoublePendulum([2.0, 1.0], [0.5, -0.3])
+    result = dp.integrate((0, 0.02), dt=1e-4, method="implicit_midpoint")  # short: just needs a few frames
+    anim = pendulum_animation(result, dp.positions, mode="phase", stride=10)
+    assert isinstance(anim, SideBySideAnimator)
+    fig = anim.build()
+    assert fig is not None
+    out1 = tmp_path / "pendulum.gif"
+    anim.save(out1, writer=PillowWriter(fps=10))
+    assert out1.exists() and out1.stat().st_size > 0
+
+    kep = KeplerSystem.from_orbital_elements(a=1.0, e=0.4)
+    result2 = kep.integrate((0, 0.5), dt=1e-3, method="yoshida4")
+    anim2 = orbit_trace_animation(result2, mode="energy", stride=50)
+    out2 = tmp_path / "orbit.gif"
+    anim2.save(out2, writer=PillowWriter(fps=10))
+    assert out2.exists() and out2.stat().st_size > 0
+
+
+@pytest.mark.slow
+def test_animate_elastic_pendulum_saves_gif(tmp_path):
+    system = ElasticPendulum([0.1, 0.05], [0.0, 0.0])
+    # short: animate_elastic_pendulum has no stride and renders one frame
+    # per sample, so this only needs enough samples for a valid animation
+    result = system.integrate((0, 0.1), dt=1e-3, method="implicit_midpoint")
+    anim = animate_elastic_pendulum(system, result)
+    out = tmp_path / "elastic_pendulum.gif"
+    anim.save(out, writer=PillowWriter(fps=10))
+    assert out.exists() and out.stat().st_size > 0
+
+
+@pytest.mark.slow
+def test_animate_rigid_body_tumble_saves_gif(tmp_path):
+    top = EulerTop([0.01, 1.0, 0.01], I1=1.0, I2=2.0, I3=3.0)
+    result = top.integrate((0, 2.0), dt=1e-3, method="implicit_midpoint")
+    anim = animate_rigid_body_tumble(top, result, stride=20)
+    out = tmp_path / "rigid_body_tumble.gif"
+    anim.save(out, writer=PillowWriter(fps=10))
+    assert out.exists() and out.stat().st_size > 0
+
+
+@pytest.mark.slow
+def test_animate_eulers_disk_saves_gif(tmp_path):
+    disk = EulersDisk(0.5, decay_rate=0.02, precession_const=1.0)
+    # short: animate_eulers_disk has no stride and renders one frame per
+    # sample (collapse time here is t_f=6.25, so 0.3 stays far from it)
+    result = disk.integrate((0.0, 0.15), dt=1e-3, method="rk4")
+    anim = animate_eulers_disk(disk, result)
+    out = tmp_path / "eulers_disk.gif"
+    anim.save(out, writer=PillowWriter(fps=10))
+    assert out.exists() and out.stat().st_size > 0
+
+
+@pytest.mark.slow
+def test_animate_rattleback_saves_gif(tmp_path):
+    system = Rattleback([0.01, 0.01, 3.0])
+    result = system.integrate((0.0, 8.0), dt=5e-3, method="rk4")
+    anim = animate_rattleback(system, result, stride=20)
+    out = tmp_path / "rattleback.gif"
+    anim.save(out, writer=PillowWriter(fps=10))
+    assert out.exists() and out.stat().st_size > 0
+
+
+def test_interactive_visualizers_return_populated_figures():
+    top = EulerTop([0.01, 1.0, 0.01], I1=1.0, I2=2.0, I3=3.0)
+    result = top.integrate((0, 10), dt=1e-3, method="implicit_midpoint")
+    fig = interactive_so3_momentum_sphere(result.y[:, :3], I1=1.0, I2=2.0, I3=3.0)
+    assert len(fig.data) == 3  # sphere surface + trajectory + start marker
+
+    kep = KeplerSystem.from_orbital_elements(a=1.0, e=0.3, c_pn=0.01)
+    result2 = kep.integrate((0, 50), dt=1e-3, method="yoshida4")
+    fig2 = interactive_orbit(result2, system=kep)
+    assert len(fig2.data) == 4  # orbit + focus + 2 LRL arrows

physicskit/classical/utils/__init__.py

@@ -0,0 +1,20 @@
+from physicskit.classical.utils.conservation import (
+    angular_momentum_2d,
+    angular_momentum_drift,
+    energy_drift,
+    lrl_drift,
+    relative_energy_drift,
+)
+from physicskit.classical.utils.stepsize import StepSizeEstimate, estimate_dt
+from physicskit.classical.utils.symbolic import LagrangianEngine
+
+__all__ = [
+    "LagrangianEngine",
+    "energy_drift",
+    "relative_energy_drift",
+    "angular_momentum_2d",
+    "angular_momentum_drift",
+    "lrl_drift",
+    "StepSizeEstimate",
+    "estimate_dt",
+]

physicskit/classical/utils/conservation.py

@@ -0,0 +1,126 @@
+"""Conservation-law diagnostics: energy drift, angular momentum, and the
+Laplace-Runge-Lenz (LRL) vector.
+
+These are thin, system-agnostic helpers that operate on the arrays
+already produced by ``system.integrate(...)`` (a
+:class:`physicskit.classical.core.base_system.SimulationResult`), used both by the
+test suite (``tests/test_conservation.py``) and by the visualizers to
+annotate plots with drift metrics.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+__all__ = [
+    "energy_drift",
+    "relative_energy_drift",
+    "angular_momentum_2d",
+    "angular_momentum_drift",
+    "lrl_drift",
+]
+
+
+def energy_drift(energy: np.ndarray) -> np.ndarray:
+    """Absolute energy deviation from the initial value.
+
+    Parameters
+    ----------
+    energy : ndarray, shape (n_steps + 1,)
+        Energy trajectory, e.g. ``result.energy``.
+
+    Returns
+    -------
+    ndarray
+        ``|H(t) - H(0)|`` at each sample.
+    """
+    energy = np.asarray(energy, dtype=np.float64)
+    return np.abs(energy - energy[0])
+
+
+def relative_energy_drift(energy: np.ndarray) -> np.ndarray:
+    """Relative energy deviation from the initial value.
+
+    This is the quantity conventionally checked against a tolerance
+    (e.g. < 1e-6) to certify a symplectic integrator is behaving
+    correctly over long integrations.
+
+    Parameters
+    ----------
+    energy : ndarray, shape (n_steps + 1,)
+        Energy trajectory, e.g. ``result.energy``.
+
+    Returns
+    -------
+    ndarray
+        ``|H(t) - H(0)| / |H(0)|`` at each sample; ``inf`` everywhere
+        if ``H(0)`` is exactly zero.
+    """
+    energy = np.asarray(energy, dtype=np.float64)
+    e0 = energy[0]
+    if e0 == 0.0:
+        return np.full_like(energy, np.inf)
+    return np.abs(energy - e0) / np.abs(e0)
+
+
+def angular_momentum_2d(q: np.ndarray, p: np.ndarray) -> np.ndarray:
+    """Out-of-plane angular momentum for a planar trajectory.
+
+    Parameters
+    ----------
+    q, p : ndarray, shape (n_steps, 2)
+        Position and momentum trajectories.
+
+    Returns
+    -------
+    ndarray, shape (n_steps,)
+        ``L = x*py - y*px`` at each sample.
+    """
+    q = np.asarray(q, dtype=np.float64)
+    p = np.asarray(p, dtype=np.float64)
+    return q[:, 0] * p[:, 1] - q[:, 1] * p[:, 0]
+
+
+def angular_momentum_drift(q: np.ndarray, p: np.ndarray) -> np.ndarray:
+    """Absolute drift of the planar angular momentum from its initial value.
+
+    Parameters
+    ----------
+    q, p : ndarray, shape (n_steps, 2)
+        Position and momentum trajectories.
+
+    Returns
+    -------
+    ndarray, shape (n_steps,)
+        ``|L(t) - L(0)|`` at each sample.
+    """
+    L = angular_momentum_2d(q, p)
+    return np.abs(L - L[0])
+
+
+def lrl_drift(system, q: np.ndarray, p: np.ndarray) -> np.ndarray:
+    """Drift of the Laplace-Runge-Lenz vector's magnitude across a trajectory.
+
+    For the unperturbed 1/r potential the LRL vector is exactly
+    conserved (both magnitude and direction); its direction sweeping
+    out an angle over time is the visual signature of apsidal
+    precession once a perturbation is present, so magnitude drift
+    remaining small while direction changes is the expected, physically
+    correct behavior for a perturbed orbit.
+
+    Parameters
+    ----------
+    system : object
+        A :class:`physicskit.classical.systems.newtonian.KeplerSystem`-like object
+        exposing ``lrl_vector(q, p) -> array``.
+    q, p : ndarray, shape (n_steps, 2)
+        Position and momentum trajectories.
+
+    Returns
+    -------
+    ndarray, shape (n_steps,)
+        ``||A(t)| - |A(0)||`` at each sample.
+    """
+    vecs = np.array([system.lrl_vector(qi, pi) for qi, pi in zip(q, p)])
+    mags = np.linalg.norm(vecs, axis=1)
+    return np.abs(mags - mags[0])

physicskit/classical/utils/stepsize.py

@@ -0,0 +1,187 @@
+"""Automated timestep selection.
+
+None of physicskit.classical's integrators are adaptive -- Verlet, Yoshida4, and
+implicit midpoint all take a single fixed ``dt`` for the whole run.
+Choosing that ``dt`` by hand (integrate a trial run, check the energy
+drift, halve or double, repeat) is exactly the kind of "black art" a
+numerical library should not leave to its users. :func:`estimate_dt`
+automates that search.
+
+For a quasi-periodic system (e.g. an unperturbed Kepler orbit), the
+characteristic energy-drift amplitude at a given ``dt`` is essentially
+independent of how many periods you run for, so a short probe
+integration is enough. **That is not true for chaotic systems**: a
+short probe can find a ``dt`` that looks fine over a few hundred steps
+but fails badly once the trajectory (e.g. a double pendulum) later
+swings through a more demanding configuration than the probe window
+happened to sample. To guard against exactly that,
+:func:`estimate_dt` always re-verifies its candidate at the full
+intended step count (``target_steps``) and shrinks ``dt`` further if
+that longer run doesn't also meet the tolerance -- it does not just
+trust the fast probe.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+
+__all__ = ["StepSizeEstimate", "estimate_dt"]
+
+
+@dataclass
+class StepSizeEstimate:
+    """Result of :func:`estimate_dt`.
+
+    Attributes
+    ----------
+    dt : float
+        The recommended, verified step size.
+    achieved_drift : float
+        Relative energy drift measured at ``dt`` over ``verified_at_steps``.
+    verified_at_steps : int
+        Step count the verification stage actually confirmed.
+    n_evaluations : int
+        Total number of trial integrations run during the search.
+    """
+
+    dt: float
+    achieved_drift: float
+    verified_at_steps: int
+    n_evaluations: int
+
+
+def _drift_for(system_factory, method, dt, n_steps) -> float:
+    system = system_factory()
+    e0 = system.energy()
+    result = system.integrate((0.0, n_steps * dt), dt=dt, method=method)
+    denom = abs(e0) if e0 != 0.0 else 1.0
+    return float(np.max(np.abs(result.energy - e0)) / denom)
+
+
+def estimate_dt(
+    system_factory: Callable[[], object],
+    method: str,
+    tol: float = 1e-6,
+    target_steps: int = 100_000,
+    probe_steps: int = 2000,
+    dt_initial: float = 1e-2,
+    growth: float = 2.0,
+    rel_precision: float = 0.05,
+    max_iter: int = 40,
+) -> StepSizeEstimate:
+    """Find the largest ``dt`` for which ``target_steps`` of ``method``
+    keeps the relative energy drift below ``tol``.
+
+    A fast bracket-and-bisect search over a short ``probe_steps``
+    window finds an initial candidate; that candidate is then always
+    re-checked at the full ``target_steps`` and shrunk further (by
+    ``growth``) if the longer run doesn't also clear ``tol`` -- see the
+    module docstring for why this verification stage matters for
+    chaotic systems.
+
+    Parameters
+    ----------
+    system_factory : callable
+        Zero-argument callable returning a *fresh* system instance in
+        its initial state (a new instance each call, since
+        ``integrate`` mutates state) -- e.g.
+        ``lambda: KeplerSystem.from_orbital_elements(a=1.0, e=0.5)``.
+    method : str
+        Integration method to pass to ``system.integrate`` (e.g.
+        ``"yoshida4"`` or ``"implicit_midpoint"``).
+    tol : float
+        Target relative energy-drift tolerance,
+        ``max|H(t) - H(0)| / |H(0)|``.
+    target_steps : int
+        The step count you actually intend to integrate for; the
+        returned ``dt`` is verified to meet ``tol`` over this many
+        steps, not just over the fast probe.
+    probe_steps : int
+        Length of each *fast* bracketing/bisection trial, in steps.
+        Should cover at least a few characteristic oscillation periods.
+    dt_initial : float
+        Starting guess for the fast probe stage.
+    growth : float
+        Factor used both to bracket a pass/fail boundary during the
+        fast probe stage, and to shrink ``dt`` during verification.
+    rel_precision : float
+        Stop bisecting once the probe-stage bracket
+        ``[dt_good, dt_bad]`` has ``(dt_bad - dt_good) / dt_good``
+        below this.
+    max_iter : int
+        Safety cap on iterations, applied separately to the
+        bracketing, bisection, and verification stages.
+
+    Returns
+    -------
+    StepSizeEstimate
+        ``.dt`` is the recommended, *verified* step size;
+        ``.achieved_drift`` is the relative drift measured at that
+        ``dt`` over ``target_steps`` (i.e. what verification actually
+        confirmed, not the fast-probe estimate).
+    """
+    n_evals = 0
+
+    def probe_drift(dt: float) -> float:
+        nonlocal n_evals
+        n_evals += 1
+        return _drift_for(system_factory, method, dt, probe_steps)
+
+    def target_drift(dt: float) -> float:
+        nonlocal n_evals
+        n_evals += 1
+        return _drift_for(system_factory, method, dt, target_steps)
+
+    # Stage 1: fast bracket + bisect over the short probe window.
+    dt = dt_initial
+    drift = probe_drift(dt)
+
+    if drift < tol:
+        dt_good, dt_bad = dt, None
+        for _ in range(max_iter):
+            dt *= growth
+            drift = probe_drift(dt)
+            if drift >= tol:
+                dt_bad = dt
+                break
+            dt_good = dt
+        else:
+            dt_bad = dt  # never failed within max_iter growth steps; use the last value as the bracket edge
+    else:
+        dt_bad = dt
+        dt_good = None
+        for _ in range(max_iter):
+            dt /= growth
+            drift = probe_drift(dt)
+            if drift < tol:
+                dt_good = dt
+                break
+            dt_bad = dt
+        else:
+            raise RuntimeError(f"Could not find a dt meeting tol={tol:.1e} within {max_iter} shrink steps")
+
+    lo, hi = dt_good, dt_bad
+    for _ in range(max_iter):
+        if (hi - lo) / lo < rel_precision:
+            break
+        mid = np.sqrt(lo * hi)
+        if probe_drift(mid) < tol:
+            lo = mid
+        else:
+            hi = mid
+
+    # Stage 2: verify (and if necessary shrink further) at the full target_steps.
+    candidate = lo
+    for _ in range(max_iter):
+        drift = target_drift(candidate)
+        if drift < tol:
+            return StepSizeEstimate(dt=candidate, achieved_drift=drift, verified_at_steps=target_steps, n_evaluations=n_evals)
+        candidate /= growth
+    raise RuntimeError(
+        f"Could not find a dt meeting tol={tol:.1e} over target_steps={target_steps} within {max_iter} "
+        "verification shrink steps -- this system's worst-case error may need a much smaller dt than "
+        "usual (common for chaotic systems); try a smaller dt_initial."
+    )