physicskit 0.1.0__py3-none-any.whl

physicskit/plasma/tests/test_new_time_domain_features.py

@@ -0,0 +1,222 @@
+"""Physics-correctness and animation-smoke tests for the newly added time-domain
+plasma simulations: magnetic reconnection and Weibel filamentation
+(:mod:`physicskit.plasma.instabilities`), Langmuir waves
+(:mod:`physicskit.plasma.kinetic`), ion-acoustic solitons and Alfven waves
+(:mod:`physicskit.plasma.waves`), drift-wave turbulence
+(:mod:`physicskit.plasma.turbulence`), wakefield particle acceleration
+(:mod:`physicskit.plasma.acceleration`), and the two-stream phase-space
+animation (:mod:`physicskit.plasma.visualizers`).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from matplotlib.animation import PillowWriter
+from scipy.signal import find_peaks
+
+from physicskit.plasma.acceleration import simulate_wakefield_acceleration
+from physicskit.plasma.instabilities import (
+    reconnection_field_from_flux,
+    reconnection_harris_ic,
+    simulate_reconnection,
+    simulate_weibel_filamentation,
+    weibel_fastest_growing_mode,
+    weibel_growth_rate,
+)
+from physicskit.plasma.kinetic import langmuir_wave_ic, pic_simulate, two_stream_ic
+from physicskit.plasma.turbulence import drift_wave_noise_ic, simulate_hasegawa_mima
+from physicskit.plasma.visualizers import (
+    animate_alfven_wave,
+    animate_drift_wave_turbulence,
+    animate_ion_acoustic_soliton,
+    animate_langmuir_wave,
+    animate_reconnection,
+    animate_two_stream_phase_space,
+    animate_wakefield_acceleration,
+    animate_weibel_filamentation,
+)
+from physicskit.plasma.waves import alfven_wave_pulse_ic, ion_acoustic_soliton_profile, simulate_alfven_wave
+
+
+class TestReconnection:
+    def test_resistivity_dissipates_magnetic_energy_faster_at_higher_eta(self):
+        nx, ny, Lx, Ly = 48, 48, 20.0, 20.0
+        dx, dy = Lx / nx, Ly / ny
+        psi0 = reconnection_harris_ic(nx, ny, Lx, Ly, sheet_width=1.0, perturbation_amplitude=0.2)
+
+        def magnetic_energy(psi):
+            Bx, By = reconnection_field_from_flux(psi, dx, dy)
+            return np.sum(Bx**2 + By**2) * dx * dy
+
+        e0 = magnetic_energy(psi0)
+        r_low = simulate_reconnection(psi0, eta=0.01, v0=0.05, dt=0.01, steps=200, Lx=Lx, Ly=Ly)
+        r_high = simulate_reconnection(psi0, eta=0.05, v0=0.05, dt=0.01, steps=200, Lx=Lx, Ly=Ly)
+        e_low = magnetic_energy(r_low["psi"])
+        e_high = magnetic_energy(r_high["psi"])
+        assert e_high < e_low < e0
+
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        psi0 = reconnection_harris_ic(32, 32, Lx=20.0, Ly=20.0, sheet_width=1.0, perturbation_amplitude=0.2)
+        anim = animate_reconnection(psi0, eta=0.02, v0=0.05, dt=0.02, steps_per_frame=5, n_frames=5, Lx=20.0, Ly=20.0)
+        out = tmp_path / "reconnection.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0
+
+
+class TestWeibelFilamentation:
+    def test_growth_rate_matches_zero_k_limit(self):
+        wpe, aniso = 1.0, 4.0
+        gamma0 = weibel_growth_rate(0.0, wpe, aniso)
+        assert gamma0 == pytest.approx(wpe * np.sqrt(aniso - 1.0))
+
+    def test_growth_rate_vanishes_above_cutoff(self):
+        wpe, aniso, c = 1.0, 4.0, 3e8
+        k_max = (wpe / c) * np.sqrt(aniso - 1.0)
+        assert weibel_growth_rate(2.0 * k_max, wpe, aniso, c=c) == 0.0
+
+    def test_fastest_growing_mode_matches_zero_k_formula(self):
+        k0, gamma_max = weibel_fastest_growing_mode(wpe=1.0, temperature_anisotropy=4.0)
+        assert k0 == 0.0
+        assert gamma_max == pytest.approx(np.sqrt(3.0))
+
+    def test_fastest_growing_mode_is_zero_for_isotropic_temperature(self):
+        assert weibel_fastest_growing_mode(wpe=1.0, temperature_anisotropy=1.0) == (0.0, 0.0)
+
+    def test_current_filaments_grow_in_time(self):
+        x = np.linspace(0, 20.0, 256, endpoint=False)
+        t = np.array([0.0, 5.0])
+        J = simulate_weibel_filamentation(x, t, wpe=1.0, temperature_anisotropy=4.0, n_modes=8, seed=0)
+        assert np.std(J[1]) > 100 * np.std(J[0])
+
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        x = np.linspace(0, 20.0, 128, endpoint=False)
+        t = np.linspace(0, 5.0, 6)
+        anim = animate_weibel_filamentation(x, t, wpe=1.0, temperature_anisotropy=4.0, n_modes=6)
+        out = tmp_path / "weibel.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0
+
+
+class TestLangmuirWave:
+    def test_field_energy_oscillates_without_substantial_decay(self):
+        k, L = 2 * np.pi / 4.0, 4.0
+        x0, v0 = langmuir_wave_ic(20000, L=L, k_mode=k, alpha=0.05, v_th=0.05, seed=1)
+        result = pic_simulate(x0, v0, L=L, ng=64, dt=0.05, steps=400)
+        field_energy = result["field_energy"]
+        peaks, _ = find_peaks(field_energy)
+        assert len(peaks) >= 4
+        peak_heights = field_energy[peaks]
+        assert peak_heights[-1] > 0.5 * peak_heights[0]
+
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        k, L = 2 * np.pi / 4.0, 4.0
+        x0, v0 = langmuir_wave_ic(4000, L=L, k_mode=k, alpha=0.05, v_th=0.05, seed=0)
+        anim = animate_langmuir_wave(x0, v0, L=L, ng=32, dt=0.05, steps_per_frame=4, n_frames=5)
+        out = tmp_path / "langmuir.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0
+
+
+class TestTwoStreamAnimation:
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        x0, v0 = two_stream_ic(2000, L=10.0, v_drift=3.0, v_th=0.5, seed=0)
+        anim = animate_two_stream_phase_space(x0, v0, L=10.0, ng=32, dt=0.05, steps_per_frame=5, n_frames=5)
+        out = tmp_path / "two_stream.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0
+
+
+class TestIonAcousticSoliton:
+    def test_soliton_propagates_at_its_speed_and_keeps_amplitude(self):
+        N, L = 512, 60.0
+        x = np.linspace(-L / 2, L / 2, N, endpoint=False)
+        u0 = ion_acoustic_soliton_profile(x, speed=4.0, x0=-15.0)
+        from physicskit.plasma.waves import ion_acoustic_soliton_evolve
+
+        dt, steps = 0.0005, 4000
+        u = ion_acoustic_soliton_evolve(u0, x, dt, steps)
+        shift = x[np.argmax(u)] - x[np.argmax(u0)]
+        assert shift == pytest.approx(4.0 * steps * dt, abs=0.5)
+        assert u.max() == pytest.approx(u0.max(), abs=0.05)
+
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        x = np.linspace(-30, 30, 256, endpoint=False)
+        u0 = ion_acoustic_soliton_profile(x, speed=4.0, x0=-15.0)
+        anim = animate_ion_acoustic_soliton(u0, x, dt=0.0005, steps_per_frame=200, n_frames=5)
+        out = tmp_path / "ion_acoustic.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0
+
+
+class TestAlfvenWave:
+    def test_pulse_splits_and_propagates_at_alfven_speed(self):
+        x = np.linspace(-20, 20, 256, endpoint=False)
+        By0, vy0 = alfven_wave_pulse_ic(x, x0=0.0, width=1.0, amplitude=0.1)
+        B0, rho0, mu0 = 1.0, 1.0, 1.0
+        vA = B0 / np.sqrt(mu0 * rho0)
+        dt, steps = 0.002, 2000
+        result = simulate_alfven_wave(By0, vy0, x, dt, steps, B0, rho0, mu0=mu0)
+        By = result["By"]
+        right_half = By[x > 0]
+        left_half = By[x < 0]
+        peak_right = x[x > 0][np.argmax(right_half)]
+        peak_left = x[x < 0][np.argmax(left_half)]
+        expected_shift = vA * steps * dt
+        assert peak_right == pytest.approx(expected_shift, abs=1.0)
+        assert peak_left == pytest.approx(-expected_shift, abs=1.0)
+
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        x = np.linspace(-20, 20, 128, endpoint=False)
+        By0, vy0 = alfven_wave_pulse_ic(x, x0=0.0, width=1.0, amplitude=0.1)
+        anim = animate_alfven_wave(By0, vy0, x, dt=0.002, steps_per_frame=100, n_frames=5, B0=1.0, rho0=1.0)
+        out = tmp_path / "alfven.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0
+
+
+class TestDriftWaveTurbulence:
+    def test_higher_dissipation_decays_faster(self):
+        n, length = 48, 2 * np.pi
+        phi0 = drift_wave_noise_ic(n, length, amplitude=0.05, seed=2)
+        dt, steps = 0.02, 300
+        r_low = simulate_hasegawa_mima(phi0, dt, steps, length, nu=0.02)
+        r_high = simulate_hasegawa_mima(phi0, dt, steps, length, nu=0.06)
+        rms_low = np.sqrt(np.mean(r_low["phi"] ** 2))
+        rms_high = np.sqrt(np.mean(r_high["phi"] ** 2))
+        assert np.isfinite(rms_low) and np.isfinite(rms_high)
+        assert rms_high < rms_low
+
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        phi0 = drift_wave_noise_ic(48, 2 * np.pi, amplitude=0.05, seed=0)
+        anim = animate_drift_wave_turbulence(phi0, dt=0.02, steps_per_frame=10, n_frames=5, length=2 * np.pi)
+        out = tmp_path / "drift_wave.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0
+
+
+class TestWakefieldAcceleration:
+    def test_phase_locked_particle_gains_energy_matching_work_energy_theorem(self):
+        result = simulate_wakefield_acceleration(x0=0.0, v0=0.9, q=1.0, m=1.0, E0=0.05, k=1.0, v_phase=1.0, dt=0.001, steps=20000)
+        ke = result["kinetic_energy"]
+        assert ke[-1] > ke[0]
+        v, t, x = result["v"], result["t"], result["x"]
+        dt = t[1] - t[0]
+        from physicskit.plasma.acceleration import wakefield_e_field
+
+        work = np.sum(wakefield_e_field(x[:-1], t[:-1], 0.05, 1.0, 1.0) * v[:-1] * dt)
+        assert (ke[-1] - ke[0]) == pytest.approx(work, rel=0.1)
+
+    @pytest.mark.slow
+    def test_animation_saves_to_gif(self, tmp_path):
+        anim = animate_wakefield_acceleration(x0=0.0, v0=0.9, q=1.0, m=1.0, E0=0.05, k=1.0, v_phase=1.0, dt=0.01, steps=400, frame_stride=40)
+        out = tmp_path / "wakefield.gif"
+        anim.save(out, writer=PillowWriter(fps=10))
+        assert out.exists() and out.stat().st_size > 0

physicskit/plasma/tests/test_plasma.py

@@ -0,0 +1,232 @@
+"""Physics-correctness tests for physicskit.plasma."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from scipy.signal import find_peaks
+
+from physicskit.plasma.kinetic import (
+    deposit_number_density,
+    landau_damping_ic,
+    landau_damping_rate,
+    pic_simulate,
+    solve_poisson_1d,
+    two_stream_ic,
+)
+from physicskit.plasma.mhd import (
+    alfven_speed,
+    grad_shafranov_rhs_solovev,
+    lundquist_number,
+    magnetosonic_speeds,
+    petschek_rate,
+    resistive_diffusion_time,
+    solovev_particular_solution,
+    solve_grad_shafranov,
+    sound_speed,
+    sweet_parker_layer_width,
+    sweet_parker_rate,
+)
+from physicskit.plasma.single_particle import (
+    ME,
+    MP,
+    QE,
+    boris_integrate,
+    curvature_drift,
+    cyclotron_frequency,
+    exb_drift,
+    grad_b_drift,
+    larmor_radius,
+    magnetic_mirror_bounce,
+    magnetic_moment,
+)
+from physicskit.plasma.waves import EPS0, cold_plasma_dispersion, plasma_frequency, rl_parameters, stix_parameters, whistler_dispersion
+
+
+class TestBorisPusher:
+    def test_conserves_kinetic_energy_in_pure_b_field(self):
+        B = np.array([0.0, 0.0, 1.0])
+        E = np.zeros(3)
+        omega_c = cyclotron_frequency(QE, MP, 1.0)
+        dt = (2 * np.pi / omega_c) / 500
+        pos_hist, vel_hist = boris_integrate(np.zeros(3), np.array([1e5, 0.0, 0.0]), QE, MP, E, B, dt, steps=5000)
+        speeds = np.linalg.norm(vel_hist, axis=1)
+        assert speeds == pytest.approx(1e5, rel=1e-9)
+
+    def test_gyroradius_matches_larmor_formula(self):
+        B = np.array([0.0, 0.0, 1.0])
+        E = np.zeros(3)
+        v0 = 1e5
+        omega_c = cyclotron_frequency(QE, MP, 1.0)
+        dt = (2 * np.pi / omega_c) / 2000
+        pos_hist, _ = boris_integrate(np.zeros(3), np.array([v0, 0.0, 0.0]), QE, MP, E, B, dt, steps=2000)
+        # the particle starts on its own gyro-circle, so its farthest excursion
+        # from the starting point is the orbit *diameter*, 2 * r_L
+        diameter_from_orbit = np.max(np.linalg.norm(pos_hist, axis=1))
+        assert diameter_from_orbit == pytest.approx(2 * larmor_radius(v0, QE, MP, 1.0), rel=1e-3)
+
+    def test_exb_drift_matches_analytic_average_velocity(self):
+        E = np.array([0.0, 1e3, 0.0])
+        B = np.array([0.0, 0.0, 1.0])
+        omega_c = cyclotron_frequency(QE, MP, 1.0)
+        dt = (2 * np.pi / omega_c) / 1000
+        pos_hist, vel_hist = boris_integrate(np.zeros(3), np.zeros(3), QE, MP, E, B, dt, steps=10000)
+        drift = exb_drift(E, B)
+        measured = np.mean(vel_hist[-1000:], axis=0)
+        assert measured == pytest.approx(drift, abs=5.0)
+
+    def test_magnetic_mirror_reflects_confined_particle(self):
+        B_func = lambda z: 1.0 + 4.0 * (z / 0.05) ** 2
+        z_hist, v_par_hist = magnetic_mirror_bounce(z0=0.0, v_par0=2e4, v_perp0=8e4, m=MP, B_func=B_func, steps=6000)
+        assert v_par_hist[0] > 0
+        assert v_par_hist[-1] < 0
+
+
+class TestGuidingCenterDrifts:
+    def test_grad_b_and_curvature_drift_are_perpendicular_to_b(self):
+        B = np.array([0.0, 0.0, 1.0])
+        grad_B = np.array([0.1, 0.0, 0.0])
+        drift = grad_b_drift(1e5, QE, MP, B, grad_B)
+        assert drift @ B == pytest.approx(0.0, abs=1e-12)
+
+        R_c = np.array([1.0, 0.0, 0.0])
+        drift_c = curvature_drift(1e5, QE, MP, B, R_c)
+        assert drift_c @ B == pytest.approx(0.0, abs=1e-12)
+
+    def test_electron_and_ion_grad_b_drifts_are_antiparallel(self):
+        B = np.array([0.0, 0.0, 1.0])
+        grad_B = np.array([0.1, 0.0, 0.0])
+        drift_ion = grad_b_drift(1e5, QE, MP, B, grad_B)
+        drift_electron = grad_b_drift(1e5, -QE, ME, B, grad_B)
+        assert drift_ion @ drift_electron < 0
+
+    def test_magnetic_moment_matches_definition(self):
+        assert magnetic_moment(v_perp=2e5, m=MP, B=2.0) == pytest.approx(0.5 * MP * 2e5**2 / 2.0)
+
+
+class TestMHDWaveSpeeds:
+    def test_alfven_speed_matches_formula(self):
+        B, rho = 1.5, 2e-6
+        mu0 = 4 * np.pi * 1e-7
+        assert alfven_speed(B, rho) == pytest.approx(B / np.sqrt(mu0 * rho))
+
+    def test_magnetosonic_parallel_limits_reduce_to_max_and_min(self):
+        vA, cs = 8e5, 1e5
+        v_fast, v_slow = magnetosonic_speeds(vA, cs, theta=0.0)
+        assert v_fast == pytest.approx(max(vA, cs))
+        assert v_slow == pytest.approx(min(vA, cs))
+
+    def test_magnetosonic_perpendicular_limit(self):
+        vA, cs = 8e5, 1e5
+        v_fast, v_slow = magnetosonic_speeds(vA, cs, theta=np.pi / 2)
+        assert v_fast == pytest.approx(np.sqrt(vA**2 + cs**2))
+        assert v_slow == pytest.approx(0.0, abs=1e-6)
+
+
+class TestGradShafranov:
+    def test_sor_solution_matches_solovev_analytic(self):
+        R = np.linspace(0.5, 1.5, 61)
+        Z = np.linspace(-0.5, 0.5, 61)
+        c1, c2 = 1.0, -2.0
+        psi = solve_grad_shafranov(R, Z, c1, c2)
+        RR, ZZ = np.meshgrid(R, Z, indexing="ij")
+        psi_exact = solovev_particular_solution(RR, ZZ, c1, c2)
+        assert np.max(np.abs(psi - psi_exact)) < 1e-3
+
+
+class TestReconnection:
+    def test_petschek_faster_than_sweet_parker_at_high_lundquist_number(self):
+        S = 1e10
+        assert petschek_rate(S) > sweet_parker_rate(S)
+
+
+class TestColdPlasmaWaves:
+    def test_parallel_propagation_gives_r_and_l_waves(self):
+        electrons = (1e19, -QE, ME)
+        ions = (1e19, QE, MP)
+        S, D, P = stix_parameters(omega=2e9, B=1.0, species=[electrons, ions])
+        R, L = rl_parameters(S, D)
+        n2_plus, n2_minus = cold_plasma_dispersion(theta=0.0, S=S, D=D, P=P)
+        assert sorted([n2_plus, n2_minus]) == pytest.approx(sorted([R, L]))
+
+    def test_perpendicular_propagation_gives_o_and_x_modes(self):
+        electrons = (1e19, -QE, ME)
+        ions = (1e19, QE, MP)
+        S, D, P = stix_parameters(omega=2e9, B=1.0, species=[electrons, ions])
+        R, L = rl_parameters(S, D)
+        n2_plus, n2_minus = cold_plasma_dispersion(theta=np.pi / 2, S=S, D=D, P=P)
+        assert sorted([n2_plus, n2_minus]) == pytest.approx(sorted([P, R * L / S]))
+
+    def test_plasma_frequency_matches_formula(self):
+        n = 1e19
+        assert plasma_frequency(n, q=QE, m=ME) == pytest.approx(np.sqrt(n * QE**2 / (EPS0 * ME)))
+
+    def test_whistler_dispersion_matches_electron_only_r_wave_root(self):
+        """See the docstring's own worked example: for omega << wce, the
+        whistler approximation should closely match the exact electron-only
+        R-wave root of cold_plasma_dispersion."""
+        wpe, wce = 1.784e11, 1.759e11
+        omega = 0.001 * wce
+        S, D, P = stix_parameters(omega, B=1.0, species=[(wpe**2 * EPS0 * ME / QE**2, -QE, ME)])
+        R, _ = rl_parameters(S, D)
+        n2_approx = whistler_dispersion(omega, wpe, wce)
+        assert n2_approx == pytest.approx(R, rel=0.01)
+
+
+class TestKineticPIC:
+    def test_cic_deposit_conserves_total_charge(self):
+        x = np.array([0.1, 2.4, 4.9, 7.7, 9.99])
+        rho = deposit_number_density(x, L=10.0, ng=20, n0=3.0)
+        assert np.sum(rho) * (10.0 / 20) == pytest.approx(3.0 * 10.0)
+
+    def test_poisson_solver_matches_analytic_sinusoid(self):
+        ng = 64
+        L = 2 * np.pi
+        x_grid = np.linspace(0, L, ng, endpoint=False)
+        rho = np.sin(x_grid)
+        E = solve_poisson_1d(rho, L)
+        assert E == pytest.approx(-np.cos(x_grid), abs=1e-10)
+
+    def test_two_stream_ic_shapes_and_symmetric_mean_velocity(self):
+        x, v = two_stream_ic(2000, L=10.0, v_drift=3.0, v_th=0.5, seed=0)
+        assert x.shape == (2000,)
+        assert v.shape == (2000,)
+        assert np.mean(v) == pytest.approx(0.0, abs=0.5)
+
+    def test_landau_damping_matches_analytic_decay_rate(self):
+        k, v_th = 0.5, 1.0
+        L = 2 * np.pi / k
+        x0, v0 = landau_damping_ic(100000, L=L, k_mode=k, alpha=0.05, v_th=v_th, seed=1)
+        result = pic_simulate(x0, v0, L=L, ng=64, dt=0.05, steps=250)
+        field_energy, t = result["field_energy"], result["t"]
+
+        peak_idx, _ = find_peaks(field_energy)
+        t_peaks, e_peaks = t[peak_idx], field_energy[peak_idx]
+        above_noise_floor = e_peaks > 1e-4
+        t_peaks, e_peaks = t_peaks[above_noise_floor][:6], e_peaks[above_noise_floor][:6]
+
+        gamma_fit = 0.5 * np.polyfit(t_peaks, np.log(e_peaks), 1)[0]
+        gamma_analytic = landau_damping_rate(k, v_th)
+        assert gamma_fit == pytest.approx(gamma_analytic, rel=0.5)
+
+
+def test_sound_speed_matches_formula():
+    assert sound_speed(gamma=5 / 3, p=1.0, rho=1e-6) == pytest.approx(np.sqrt(5 / 3 * 1.0 / 1e-6))
+
+
+def test_grad_shafranov_rhs_solovev_matches_formula():
+    R = np.array([1.0, 2.0])
+    np.testing.assert_allclose(grad_shafranov_rhs_solovev(R, c1=1.0, c2=-2.0), 1.0 * R**2 - 2.0)
+
+
+def test_lundquist_number_matches_formula():
+    assert lundquist_number(L=1.0, vA=1e6, eta=1.0) == pytest.approx(1e6)
+
+
+def test_resistive_diffusion_time_matches_formula():
+    assert resistive_diffusion_time(L=1.0, eta=1.0) == pytest.approx(1.0)
+    assert resistive_diffusion_time(L=2.0, eta=1.0) == pytest.approx(4.0)
+
+
+def test_sweet_parker_layer_width_matches_formula():
+    assert sweet_parker_layer_width(L=1e7, S=1e6) == pytest.approx(1e7 / np.sqrt(1e6))

physicskit/plasma/tests/test_visualizers.py

@@ -0,0 +1,79 @@
+"""Coverage for physicskit.plasma.visualizers's static plotting functions
+(previously untested): the ax=None/ax=<given> branch of each. All pure
+matplotlib, small inputs -- fast."""
+
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from physicskit.plasma.kinetic import landau_damping_ic
+from physicskit.plasma.mhd import safety_factor_large_aspect_ratio, solve_grad_shafranov
+from physicskit.plasma.single_particle import MP, QE, boris_integrate
+from physicskit.plasma.visualizers import (
+    plot_cma_diagram,
+    plot_drift_trajectory,
+    plot_field_energy_history,
+    plot_flux_surfaces,
+    plot_particle_orbit_3d,
+    plot_phase_space,
+    plot_q_profile,
+)
+from physicskit.plasma.waves import cma_coordinates
+
+
+def test_plot_particle_orbit_3d_given_ax():
+    pos_hist, _ = boris_integrate(np.zeros(3), np.array([1e5, 0.0, 0.0]), QE, MP, np.zeros(3), np.array([0.0, 0.0, 1.0]), 1e-10, steps=20)
+    fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
+    _, ax_given = plot_particle_orbit_3d(pos_hist, ax=ax)
+    assert ax_given is ax
+    assert fig is ax.figure
+
+
+def test_plot_drift_trajectory_given_ax():
+    E = np.array([0.0, 1e3, 0.0])
+    B = np.array([0.0, 0.0, 1.0])
+    pos_hist, _ = boris_integrate(np.zeros(3), np.zeros(3), QE, MP, E, B, 1e-10, steps=20)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_drift_trajectory(pos_hist, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_flux_surfaces_given_ax():
+    R = np.linspace(0.5, 1.5, 11)
+    Z = np.linspace(-0.5, 0.5, 11)
+    psi = solve_grad_shafranov(R, Z, c1=1.0, c2=-2.0)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_flux_surfaces(R, Z, psi, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_q_profile_given_ax():
+    r = np.linspace(0.05, 0.5, 10)
+    q = np.array([safety_factor_large_aspect_ratio(ri, R0=1.0, Bt=2.0, Bp=0.2) for ri in r])
+    fig, ax = plt.subplots()
+    _, ax_given = plot_q_profile(r, q, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_phase_space_given_ax():
+    x, v = landau_damping_ic(200, L=4 * np.pi, k_mode=0.5, alpha=0.1, v_th=1.0, seed=0)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_phase_space(x, v, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_field_energy_history_given_ax():
+    t = np.linspace(0, 10, 20)
+    field_energy = np.exp(-0.1 * t)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_field_energy_history(t, field_energy, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_cma_diagram_given_ax():
+    omega = np.linspace(0.1, 5.0, 10)
+    X, Y = cma_coordinates(omega, wpe=2.0, wce=3.0)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_cma_diagram(X, Y, ax=ax)
+    assert ax_given is ax