physicskit 0.1.0__py3-none-any.whl

physicskit/astro/tests/test_stellar_dynamo.py

@@ -0,0 +1,198 @@
+import numpy as np
+import pytest
+
+from physicskit.astro.stellar_dynamo import (
+    _spectral_grid_2d,
+    alpha_omega_growth_rate,
+    alpha_omega_wave_frequency,
+    convective_roll_rhs,
+    dominant_mode_growth_rate,
+    kinetic_energy,
+    simulate_alpha_omega_dynamo,
+    simulate_stellar_convection,
+)
+
+# ---------------------------------------------------------------------------
+# Part 2: alpha-omega mean-field dynamo -- the core dispersion-relation check
+# ---------------------------------------------------------------------------
+
+
+def test_alpha_omega_growth_rate_matches_simulation_for_unstable_mode():
+    """The single most important correctness guarantee in this module:
+    seed a pure Fourier mode, evolve it, and check the numerically
+    observed growth rate against the closed-form dispersion relation."""
+    Lx = 2 * np.pi
+    nx = 64
+    x = np.linspace(0, Lx, nx, endpoint=False)
+    k0 = 1.0  # m = 1: k0 = 2*pi*m/Lx with Lx = 2*pi
+    alpha, shear, eta = 1.0, 5.0, 0.05
+
+    A0 = 1e-3 * np.cos(k0 * x)
+    B0 = np.zeros_like(A0)
+
+    sigma_predicted = alpha_omega_growth_rate(alpha, shear, k0, eta)
+    assert sigma_predicted > 0  # this test is specifically the unstable/growing case
+
+    dt = 0.01
+    n_steps = 600
+    save_every = 5
+    times, A_snapshots, B_snapshots = simulate_alpha_omega_dynamo(A0, B0, alpha, shear, eta, Lx, dt, n_steps, save_every)
+
+    sigma_observed = dominant_mode_growth_rate(B_snapshots, dt_save=dt * save_every)
+    assert sigma_observed == pytest.approx(sigma_predicted, rel=0.02)
+
+    # Cross-check against A too, and against the module's independent
+    # frequency formula (same magnitude as the growth rate, by construction
+    # of the dispersion relation s^2 = i*alpha*shear*k).
+    sigma_observed_A = dominant_mode_growth_rate(A_snapshots, dt_save=dt * save_every)
+    assert sigma_observed_A == pytest.approx(sigma_predicted, rel=0.02)
+    omega_freq = alpha_omega_wave_frequency(alpha, shear, k0)
+    assert omega_freq == pytest.approx(np.sqrt(alpha * shear * k0 / 2.0), rel=1e-10)
+
+
+def test_alpha_omega_dispersion_relation_holds_for_several_wavenumbers():
+    """The closed-form growth rate and wave-frequency magnitude must agree
+    (they're the real and imaginary parts of the same complex root)."""
+    alpha, shear, eta = 0.8, 3.0, 0.1
+    k = np.array([0.5, 1.0, 2.0, 3.0])
+    sigma = alpha_omega_growth_rate(alpha, shear, k, eta)
+    freq = alpha_omega_wave_frequency(alpha, shear, k)
+    assert freq == pytest.approx(sigma + eta * k**2, rel=1e-10)
+
+
+def test_alpha_omega_dynamo_decays_in_stable_regime():
+    """Large diffusivity relative to alpha*shear*k should make the
+    dispersion relation predict decay, and the simulated field energy
+    should actually decay."""
+    Lx = 2 * np.pi
+    nx = 64
+    x = np.linspace(0, Lx, nx, endpoint=False)
+    k0 = 1.0
+    alpha, shear, eta = 1.0, 1.0, 2.0
+
+    sigma_predicted = alpha_omega_growth_rate(alpha, shear, k0, eta)
+    assert sigma_predicted < 0  # this test is specifically the stable/decaying case
+
+    A0 = 1e-3 * np.cos(k0 * x)
+    B0 = np.zeros_like(A0)
+    dt = 0.01
+    n_steps = 600
+    save_every = 5
+    times, A_snapshots, B_snapshots = simulate_alpha_omega_dynamo(A0, B0, alpha, shear, eta, Lx, dt, n_steps, save_every)
+
+    initial_energy = np.sum(A_snapshots[0] ** 2 + B_snapshots[0] ** 2)
+    final_energy = np.sum(A_snapshots[-1] ** 2 + B_snapshots[-1] ** 2)
+    assert final_energy < initial_energy
+
+
+def test_simulate_alpha_omega_dynamo_shapes_and_finiteness():
+    Lx = 2 * np.pi
+    nx = 40
+    x = np.linspace(0, Lx, nx, endpoint=False)
+    A0 = 1e-3 * np.cos(2 * x)
+    B0 = np.zeros_like(A0)
+    times, A_snapshots, B_snapshots = simulate_alpha_omega_dynamo(A0, B0, alpha=1.0, shear=2.0, eta=0.1, Lx=Lx, dt=0.02, n_steps=30, save_every=3)
+    n_saved = 30 // 3 + 1
+    assert times.shape == (n_saved,)
+    assert A_snapshots.shape == (n_saved, nx)
+    assert B_snapshots.shape == (n_saved, nx)
+    assert np.all(np.isfinite(A_snapshots))
+    assert np.all(np.isfinite(B_snapshots))
+    assert A_snapshots[0] == pytest.approx(A0)
+    assert B_snapshots[0] == pytest.approx(B0)
+
+
+def test_alpha_omega_growth_rate_vectorizes_over_k():
+    k = np.array([0.0, 1.0, 2.0, 4.0])
+    sigma = alpha_omega_growth_rate(alpha=1.0, shear=2.0, k=k, eta=0.1)
+    assert sigma.shape == (4,)
+    # k=0: no alpha-omega drive at all, growth rate is exactly zero (no diffusion either, at k=0).
+    assert sigma[0] == pytest.approx(0.0)
+
+
+def test_dominant_mode_growth_rate_recovers_known_slope():
+    t = np.arange(20) * 0.1
+    sigma_true = 1.7
+    A_snapshots = np.exp(sigma_true * t)[:, None] * np.ones((20, 6))
+    sigma_fit = dominant_mode_growth_rate(A_snapshots, dt_save=0.1)
+    assert sigma_fit == pytest.approx(sigma_true, rel=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# Part 1: 2D convective rolls
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("seed", [0, 1])
+def test_convection_grows_and_stays_finite(seed):
+    """Small random noise should grow into sustained convective motion
+    (kinetic energy at the end of the run meaningfully larger than at the
+    start) without ever producing NaN/Inf, for at least two random seeds."""
+    nx = ny = 32
+    Lx = Ly = 2 * np.pi
+    times, omega_snapshots, T_snapshots = simulate_stellar_convection(nx, ny, Lx, Ly, n_steps=2000, seed=seed)
+
+    assert np.all(np.isfinite(omega_snapshots))
+    assert np.all(np.isfinite(T_snapshots))
+
+    _, _, KX, KY, K2 = _spectral_grid_2d(nx, ny, Lx, Ly)
+    KE_initial = kinetic_energy(omega_snapshots[0], KX, KY, K2)
+    KE_final = kinetic_energy(omega_snapshots[-1], KX, KY, K2)
+    assert KE_final > 100.0 * KE_initial
+
+
+def test_convective_roll_rhs_shapes():
+    nx = ny = 32
+    Lx = Ly = 2 * np.pi
+    _, Y, KX, KY, K2 = _spectral_grid_2d(nx, ny, Lx, Ly)
+    rng = np.random.default_rng(0)
+    omega = 1e-3 * rng.standard_normal((ny, nx))
+    T = 1e-3 * rng.standard_normal((ny, nx))
+    domega_dt, dT_dt = convective_roll_rhs(omega, T, Y, KX, KY, K2, nu=0.05, kappa=0.05, g_alpha=1.0, beta=1.0, Ly=Ly, tau_relax=0.5)
+    assert domega_dt.shape == (ny, nx)
+    assert dT_dt.shape == (ny, nx)
+    assert np.all(np.isfinite(domega_dt))
+    assert np.all(np.isfinite(dT_dt))
+
+
+def test_convective_roll_rhs_relaxation_pulls_uniform_zero_field_toward_teq():
+    """With zero velocity and zero perturbation temperature, the only
+    surviving term in dT'/dt is the Newtonian-cooling relaxation toward
+    Teq(y); check its sign and magnitude directly against the formula."""
+    nx = ny = 16
+    Lx = Ly = 2 * np.pi
+    _, Y, KX, KY, K2 = _spectral_grid_2d(nx, ny, Lx, Ly)
+    omega = np.zeros((ny, nx))
+    T = np.zeros((ny, nx))
+    beta, tau_relax = 2.0, 0.5
+    _, dT_dt = convective_roll_rhs(omega, T, Y, KX, KY, K2, nu=0.05, kappa=0.05, g_alpha=1.0, beta=beta, Ly=Ly, tau_relax=tau_relax)
+    Teq = -beta * np.cos(2 * np.pi * Y / Ly)
+    expected = Teq / tau_relax
+    assert dT_dt == pytest.approx(expected, abs=1e-10)
+
+
+def test_simulate_stellar_convection_shapes():
+    nx, ny = 24, 32
+    Lx, Ly = 2 * np.pi, 3.0
+    n_steps, save_every = 20, 5
+    times, omega_snapshots, T_snapshots = simulate_stellar_convection(nx, ny, Lx, Ly, n_steps=n_steps, save_every=save_every, seed=0)
+    n_saved = n_steps // save_every + 1
+    assert times.shape == (n_saved,)
+    assert omega_snapshots.shape == (n_saved, ny, nx)
+    assert T_snapshots.shape == (n_saved, ny, nx)
+
+
+def test_kinetic_energy_of_zero_vorticity_is_zero():
+    nx = ny = 32
+    Lx = Ly = 2 * np.pi
+    _, _, KX, KY, K2 = _spectral_grid_2d(nx, ny, Lx, Ly)
+    assert kinetic_energy(np.zeros((ny, nx)), KX, KY, K2) == 0.0
+
+
+def test_kinetic_energy_positive_for_nonzero_vorticity():
+    nx = ny = 32
+    Lx = Ly = 2 * np.pi
+    X, Y, KX, KY, K2 = _spectral_grid_2d(nx, ny, Lx, Ly)
+    omega = np.sin(X) * np.cos(Y)
+    assert kinetic_energy(omega, KX, KY, K2) > 0.0

physicskit/astro/tests/test_stellar_structure.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pytest
+
+from physicskit.astro.stellar_structure import (
+    PolytropicStar,
+    chandrasekhar_mass,
+    lane_emden,
+    main_sequence_luminosity,
+)
+
+
+def test_lane_emden_n0_matches_analytic_solution():
+    xi, theta = lane_emden(0.0)
+    assert xi[-1] == pytest.approx(np.sqrt(6.0), rel=1e-3)
+    # analytic: theta(xi) = 1 - xi^2/6
+    analytic = 1.0 - xi**2 / 6.0
+    assert theta == pytest.approx(analytic, abs=2e-3)
+
+
+def test_lane_emden_n1_matches_analytic_solution():
+    xi, theta = lane_emden(1.0)
+    assert xi[-1] == pytest.approx(np.pi, rel=1e-3)
+    # analytic: theta(xi) = sin(xi)/xi
+    analytic = np.sin(xi) / xi
+    assert theta == pytest.approx(analytic, abs=2e-3)
+
+
+def test_polytropic_star_n0_matches_uniform_sphere_mass():
+    star = PolytropicStar(n=0.0, K=2.5, rho_c=3.0)
+    uniform_mass = 4.0 / 3.0 * np.pi * star.radius**3 * star.rho_c
+    assert star.mass == pytest.approx(uniform_mass, rel=1e-2)
+
+
+def test_polytropic_star_positive_physical_quantities():
+    star = PolytropicStar(n=1.5, K=1.0, rho_c=1.0)
+    assert star.radius > 0
+    assert star.mass > 0
+    assert star.xi1 > 0
+
+
+def test_chandrasekhar_mass_near_known_value():
+    M = chandrasekhar_mass(2.0)
+    assert 1.2 < M < 1.5
+
+
+def test_chandrasekhar_mass_scales_as_inverse_mu_e_squared():
+    assert chandrasekhar_mass(4.0) == pytest.approx(chandrasekhar_mass(2.0) / 4.0)
+
+
+def test_main_sequence_luminosity_power_law():
+    assert main_sequence_luminosity(2.0) == pytest.approx(2.0**3.5)
+    assert main_sequence_luminosity(1.0) == pytest.approx(1.0)

physicskit/astro/tests/test_visualizers.py

@@ -0,0 +1,159 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+from matplotlib.animation import PillowWriter
+from matplotlib.collections import LineCollection
+
+from physicskit.astro.cosmic_web import lagrangian_grid, random_displacement_potential
+from physicskit.astro.galactic_dynamics import circular_velocity
+from physicskit.astro.nbody import NBodySystem, figure_eight_initial_conditions
+from physicskit.astro.stellar_dynamo import simulate_alpha_omega_dynamo, simulate_stellar_convection
+from physicskit.astro.stellar_structure import lane_emden
+from physicskit.astro.visualizers import (
+    animate_dynamo_wave,
+    animate_nbody_trajectories,
+    animate_stellar_convection,
+    animate_zeldovich_collapse,
+    plot_dynamo_butterfly_diagram,
+    plot_lane_emden,
+    plot_nbody_trajectories,
+    plot_rotation_curve,
+    plot_zeldovich_snapshot,
+)
+
+
+@pytest.mark.slow
+def test_animate_nbody_trajectories_figure_eight_saves_gif(tmp_path):
+    positions, velocities, masses = figure_eight_initial_conditions()
+    system = NBodySystem(positions, velocities, masses)
+    history = system.simulate(dt=0.002, n_steps=300)
+
+    anim = animate_nbody_trajectories(history, skip=5)
+    out = tmp_path / "nbody_figure_eight.gif"
+    anim.save(out, writer=PillowWriter(fps=10))
+    assert out.exists() and out.stat().st_size > 0
+
+
+def test_animate_nbody_trajectories_uses_one_collection_and_marker_per_body():
+    history = np.random.default_rng(1).normal(size=(50, 3, 3))
+
+    anim = animate_nbody_trajectories(history, title="N-body test")
+    ax = anim._fig.axes[0]
+    collections = [c for c in ax.collections if isinstance(c, LineCollection)]
+    assert len(collections) == 3
+    assert len(ax.lines) == 3
+    assert ax.get_title() == "N-body test"
+
+    anim._draw_frame(0)  # a single point per body: the <2-point trail branch
+    assert collections[0].get_segments() == []
+
+
+def test_animate_nbody_trajectories_trail_fades_with_alpha():
+    """The trail must actually fade: older segments more transparent than
+    the most recent one."""
+    history = np.cumsum(np.random.default_rng(2).normal(size=(100, 3, 3)), axis=0)
+
+    anim = animate_nbody_trajectories(history, trail=100)
+    anim._draw_frame(99)
+    ax = anim._fig.axes[0]
+    lc = next(c for c in ax.collections if isinstance(c, LineCollection))
+    alphas = np.asarray(lc.get_colors())[:, 3]
+    assert alphas[0] < alphas[-1]
+
+
+def test_plot_nbody_trajectories_returns_one_line_per_body():
+    # A synthetic array is enough here -- plot_nbody_trajectories only
+    # cares about (n_steps+1, N, 3) shape, not real dynamics.
+    history = np.random.default_rng(0).normal(size=(10, 3, 3))
+
+    fig, ax = plot_nbody_trajectories(history)
+    assert isinstance(fig, plt.Figure)
+    assert len(ax.lines) == 3
+
+    _, ax_given = plot_nbody_trajectories(history, ax=ax)
+    assert ax_given is ax  # reuses the supplied axes rather than creating a new figure
+
+
+def test_plot_lane_emden_returns_figure():
+    xi, theta = lane_emden(1.5, n_points=50)
+    fig, ax = plot_lane_emden(xi, theta)
+    assert isinstance(fig, plt.Figure)
+    assert len(ax.lines) == 2  # theta(xi) curve + the axhline(0) reference line
+    assert len(ax.collections) == 1  # the surface-point scatter
+
+    _, ax_given = plot_lane_emden(xi, theta, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_rotation_curve_with_and_without_observed_data():
+    r = np.linspace(1.0, 10.0, 10)
+    v_model = circular_velocity(r, lambda rr: 1.0)
+
+    fig1, ax1 = plot_rotation_curve(r, v_model)
+    assert isinstance(fig1, plt.Figure)
+    assert len(ax1.collections) == 0  # no scatter without observed data
+
+    fig2, ax2 = plot_rotation_curve(r, v_model, v_observed=v_model * 1.1)
+    assert isinstance(fig2, plt.Figure)
+    assert len(ax2.collections) == 1  # scatter added for observed data
+
+    _, ax_given = plot_rotation_curve(r, v_model, ax=ax1)
+    assert ax_given is ax1
+
+
+def _small_zeldovich_field():
+    q = lagrangian_grid(5, 1.0)
+    k_vectors, amplitudes, phases = random_displacement_potential(4, k_min=2 * np.pi, k_max=6 * np.pi, amplitude_scale=0.02, seed=0)
+    return q, k_vectors, amplitudes, phases
+
+
+def test_animate_zeldovich_collapse_builds_and_draws_a_frame():
+    q, k_vectors, amplitudes, phases = _small_zeldovich_field()
+    D_values = np.linspace(0.0, 0.1, 3)
+    anim = animate_zeldovich_collapse(q, D_values, k_vectors, amplitudes, phases)
+    anim._draw_frame(1)
+    assert anim._fig.axes[0].collections  # the density-colored scatter exists
+
+
+def test_plot_zeldovich_snapshot_returns_figure():
+    q, k_vectors, amplitudes, phases = _small_zeldovich_field()
+    fig, ax = plot_zeldovich_snapshot(q, 0.05, k_vectors, amplitudes, phases)
+    assert isinstance(fig, plt.Figure)
+    assert ax.collections
+
+    _, ax_given = plot_zeldovich_snapshot(q, 0.05, k_vectors, amplitudes, phases, ax=ax)
+    assert ax_given is ax
+
+
+@pytest.mark.slow
+def test_animate_stellar_convection_builds_and_draws_a_frame():
+    times, omega_snaps, T_snaps = simulate_stellar_convection(16, 16, 2 * np.pi, 2 * np.pi, n_steps=20, save_every=10, seed=0)
+    anim = animate_stellar_convection(times, omega_snaps, T_snaps)
+    anim._draw_frame(1)
+    ax_omega, ax_T = anim._fig.axes[0], anim._fig.axes[1]
+    assert len(ax_omega.images) == 1 and len(ax_T.images) == 1
+
+
+def test_animate_dynamo_wave_builds_and_draws_a_frame():
+    Lx = 2 * np.pi
+    x = np.linspace(0, Lx, 16, endpoint=False)
+    A0 = 1e-3 * np.cos(x)
+    B0 = np.zeros_like(A0)
+    times, A_snaps, B_snaps = simulate_alpha_omega_dynamo(A0, B0, alpha=1.0, shear=5.0, eta=0.05, Lx=Lx, dt=0.01, n_steps=20, save_every=5)
+    anim = animate_dynamo_wave(times, A_snaps, B_snaps, x)
+    anim._draw_frame(1)
+    assert len(anim._fig.axes[0].lines) == 2  # A(x) and B(x)
+
+
+def test_plot_dynamo_butterfly_diagram_returns_figure():
+    Lx = 2 * np.pi
+    x = np.linspace(0, Lx, 16, endpoint=False)
+    A0 = 1e-3 * np.cos(x)
+    B0 = np.zeros_like(A0)
+    times, A_snaps, B_snaps = simulate_alpha_omega_dynamo(A0, B0, alpha=1.0, shear=5.0, eta=0.05, Lx=Lx, dt=0.01, n_steps=20, save_every=5)
+    fig, ax = plot_dynamo_butterfly_diagram(times, B_snaps, x)
+    assert isinstance(fig, plt.Figure)
+    assert ax.collections
+
+    _, ax_given = plot_dynamo_butterfly_diagram(times, B_snaps, x, ax=ax)
+    assert ax_given is ax