From f939a3e69c9b21ac48432c8fc4b7745c155cc4b5 Mon Sep 17 00:00:00 2001 From: rogeriojorge Date: Sat, 26 Sep 2026 16:00:05 -0500 Subject: [PATCH 1/2] Fix multispecies perturbation bookkeeping --- spectrax/_simulation.py | 3 +-- tests/test_simulation_runs.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/spectrax/_simulation.py b/spectrax/_simulation.py index eba5de0..4a7ea81 100644 --- a/spectrax/_simulation.py +++ b/spectrax/_simulation.py @@ -181,8 +181,7 @@ def simulation(input_parameters={}, Nx=33, Ny=1, Nz=1, Nn=20, Nm=1, Np=1, Ns=2, Fk = sol.ys[:,(-6 * (Nx//2+1) * Ny * Nz):].reshape(len(sol.ts), 6, Ny, Nx//2+1, Nz) # Set n = 0, k = 0 mode to zero to get array with time evolution of perturbation. - dCk = Ck.at[:, 0, 0, 0, 0].set(0) - dCk = dCk.at[:, Nn * Nm * Np, 0, 0, 0].set(0) + dCk = Ck.at[:, jnp.arange(Ns) * Nn * Nm * Np, 0, 0, 0].set(0) # Output results temporary_output = {"Ck": Ck, "Fk": Fk, "time": time, "dCk": dCk, "solver_stats": sol.stats} diff --git a/tests/test_simulation_runs.py b/tests/test_simulation_runs.py index 2d05a1e..739575b 100644 --- a/tests/test_simulation_runs.py +++ b/tests/test_simulation_runs.py @@ -2,6 +2,25 @@ import jax.numpy as jnp from spectrax import simulation +def test_multispecies_background_removed_from_perturbation(): + """Remove every species' homogeneous density from dCk.""" + Ns = 3 + Ck_0 = jnp.zeros((Ns, 1, 2, 1), dtype=jnp.complex128) + Ck_0 = Ck_0.at[:, 0, 0, 0].set(jnp.arange(1, Ns + 1)) + parameters = { + "Ck_0": Ck_0, + "Fk_0": jnp.zeros((6, 1, 2, 1), dtype=jnp.complex128), + "qs": jnp.array([-1.0, 1.0, 1.0]), + "alpha_s": jnp.ones(3 * Ns), + "u_s": jnp.zeros(3 * Ns), + "Omega_cs": jnp.ones(Ns), + "t_max": 0.001, + } + + result = simulation(parameters, Nx=3, Nn=1, Ns=Ns, timesteps=2, + dt=0.001, adaptive_time_step=False) + assert jnp.all(result["dCk"][:, :, 0, 0, 0] == 0) + def test_simulation_runs(): """Test if the simulation runs without errors with default parameters.""" result = simulation() From 3f89a099d59b17ac9154f4abb232d891dcdad3a8 Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Sat, 3 Oct 2026 09:29:22 -0500 Subject: [PATCH 2/2] Test perturbation bookkeeping for 1-4 populations with a mean current dCk must differ from Ck only in each species' uniform density slot; the k=0 first Hermite moment (mean current) and all k != 0 coefficients are kept, and the net mean current still drives the uniform field. --- tests/test_simulation_runs.py | 34 ++++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/tests/test_simulation_runs.py b/tests/test_simulation_runs.py index 739575b..4b9f17e 100644 --- a/tests/test_simulation_runs.py +++ b/tests/test_simulation_runs.py @@ -2,24 +2,34 @@ import jax.numpy as jnp from spectrax import simulation -def test_multispecies_background_removed_from_perturbation(): - """Remove every species' homogeneous density from dCk.""" - Ns = 3 - Ck_0 = jnp.zeros((Ns, 1, 2, 1), dtype=jnp.complex128) - Ck_0 = Ck_0.at[:, 0, 0, 0].set(jnp.arange(1, Ns + 1)) +@pytest.mark.parametrize("Ns", [1, 2, 3, 4]) +def test_perturbation_removes_every_background_density_and_keeps_mean_current(Ns): + """dCk differs from Ck only in each species' (n=0, k=0) density slot; the k=0 first Hermite + moment (the physical mean current) and every k != 0 coefficient are kept unchanged.""" + Nn, Nx = 2, 4 + Ck_0 = jnp.zeros((Ns * Nn, 1, Nx // 2 + 1, 1), dtype=jnp.complex128) + for s in range(Ns): + Ck_0 = Ck_0.at[s * Nn, 0, 0, 0].set(1.0 + s) # background density + Ck_0 = Ck_0.at[s * Nn + 1, 0, 0, 0].set(0.1 * (s + 1)) # uniform drift: mean current + Ck_0 = Ck_0.at[s * Nn, 0, 1, 0].set(1e-3j * (s + 1)) # density perturbation at kx=1 parameters = { "Ck_0": Ck_0, - "Fk_0": jnp.zeros((6, 1, 2, 1), dtype=jnp.complex128), - "qs": jnp.array([-1.0, 1.0, 1.0]), + "Fk_0": jnp.zeros((6, 1, Nx // 2 + 1, 1), dtype=jnp.complex128), + "qs": jnp.array([(-1.0) ** (s + 1) for s in range(Ns)]), "alpha_s": jnp.ones(3 * Ns), "u_s": jnp.zeros(3 * Ns), "Omega_cs": jnp.ones(Ns), - "t_max": 0.001, + "t_max": 0.01, } - - result = simulation(parameters, Nx=3, Nn=1, Ns=Ns, timesteps=2, - dt=0.001, adaptive_time_step=False) - assert jnp.all(result["dCk"][:, :, 0, 0, 0] == 0) + result = simulation(parameters, Nx=Nx, Nn=Nn, Ns=Ns, timesteps=2, dt=0.001, adaptive_time_step=False) + Ck, dCk = result["Ck"], result["dCk"] + density = jnp.arange(Ns) * Nn + assert jnp.all(dCk[:, density, 0, 0, 0] == 0) + assert jnp.all(Ck[:, density, 0, 0, 0] != 0) + keep = jnp.ones(Ck.shape[1:], dtype=bool).at[density, 0, 0, 0].set(False) + assert jnp.array_equal(jnp.where(keep, dCk, 0), jnp.where(keep, Ck, 0)) + assert jnp.all(dCk[:, density + 1, 0, 0, 0] != 0) # mean current survives in the perturbation + assert result["Fk"][-1, 0, 0, 0, 0] != 0 # the net mean current drives the uniform Ex def test_simulation_runs(): """Test if the simulation runs without errors with default parameters."""