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..4b9f17e 100644 --- a/tests/test_simulation_runs.py +++ b/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()