Fix multispecies perturbation bookkeeping

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}

tests/test_simulation_runs.py

@@ -2,6 +2,35 @@
 import jax.numpy as jnp
 from spectrax import simulation
 
+@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, 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.01,
+    }
+    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."""
     result = simulation()