diff --git a/spectrax/_diagnostics.py b/spectrax/_diagnostics.py index 8fdb788..0c994cc 100644 --- a/spectrax/_diagnostics.py +++ b/spectrax/_diagnostics.py @@ -19,7 +19,24 @@ import jax.numpy as jnp -__all__ = ["diagnostics"] +__all__ = ["diagnostics", "hermite_tail_fraction"] + + +def _rfft_weights(Nx: int, Nx_kept: int) -> jnp.ndarray: + """Parseval weights for the stored half of a real FFT of original length ``Nx``. + + ``k = 0`` has weight 1; the Nyquist coefficient (present only for even ``Nx``) has + weight 1; every other stored coefficient stands for itself and its omitted conjugate, + weight 2. ``Nx`` must be the original grid length: the stored array must be the full + real FFT, ``Nx_kept == Nx // 2 + 1``, not a cropped or dealiased subset. + """ + Nx, Nx_kept = int(Nx), int(Nx_kept) + if Nx < 1 or Nx_kept != Nx // 2 + 1: + raise ValueError(f"Stored real-FFT axis has {Nx_kept} modes; a full rfft of Nx={Nx} has {Nx // 2 + 1}.") + weights = jnp.full(Nx_kept, 2.0).at[0].set(1.0) + if Nx % 2 == 0 and Nx > 1: + weights = weights.at[-1].set(1.0) + return weights def _infer_Ns(output: Mapping[str, Any]) -> int: @@ -68,6 +85,50 @@ def _infer_masses(output: Mapping[str, Any], Ns: int) -> jnp.ndarray: return jnp.ones((Ns,)) +def hermite_tail_fraction(output: Mapping[str, Any], width: int = 1) -> jnp.ndarray: + """Return the perturbation norm fraction at each Hermite boundary. + + The tail is the union of the outer ``width`` modes on every Hermite axis + with more than one retained mode. Squared ``dCk`` coefficients are summed + with Parseval weights over Fourier space. For SPECTRAX's matched-Maxwellian + basis this norm is proportional to perturbed-distribution free energy, but + it is not the total plasma energy. + + Returns an array with shape ``(Nt, Ns)``. A zero perturbation, or a layout + with no active Hermite axis, has zero tail fraction. + """ + if not isinstance(width, int) or width < 1: + raise ValueError("width must be a positive integer.") + + dCk = jnp.asarray(output["dCk"]) + Ns = int(output["Ns"]) if "Ns" in output else _infer_Ns(output) + Nn, Nm, Np = (int(output[key]) for key in ("Nn", "Nm", "Np")) + expected_modes = Ns * Nn * Nm * Np + if dCk.ndim != 5 or dCk.shape[1] != expected_modes: + raise ValueError(f"dCk must have shape (Nt, {expected_modes}, Ny, Nx//2+1, Nz).") + + n = jnp.arange(Nn)[None, None, :] + m = jnp.arange(Nm)[None, :, None] + p = jnp.arange(Np)[:, None, None] + tail = jnp.zeros((Np, Nm, Nn), dtype=bool) + if Nn > 1: + tail |= n >= max(Nn - width, 0) + if Nm > 1: + tail |= m >= max(Nm - width, 0) + if Np > 1: + tail |= p >= max(Np - width, 0) + + Nx_kept = dCk.shape[-2] + weights = _rfft_weights(output["Nx"], Nx_kept) + power = jnp.abs(dCk.reshape(dCk.shape[0], Ns, Np, Nm, Nn, *dCk.shape[-3:])) ** 2 + power *= weights.reshape(1, 1, 1, 1, 1, 1, -1, 1) + total = jnp.sum(power, axis=(2, 3, 4, 5, 6, 7)) + boundary = jnp.sum(power * tail.reshape(1, 1, Np, Nm, Nn, 1, 1, 1), + axis=(2, 3, 4, 5, 6, 7)) + denominator = jnp.where(total > 0, total, 1.0) + return jnp.where(total > 0, boundary / denominator, 0.0) + + def diagnostics(output: dict) -> None: """Compute and attach common diagnostics to ``output`` (mutates in-place). @@ -83,6 +144,7 @@ def diagnostics(output: dict) -> None: - ``Fk``: field Fourier coefficients of shape ``(Nt, 6, Ny, Nx//2+1, Nz)`` - ``Omega_cs``: array of shape ``(Ns,)`` (cyclotron frequencies) - ``Lx``: domain length in x (used for ``k_norm``) + - ``Nx``: original x grid length; ``Fk``/``Ck`` must hold its full real FFT Optionally, it may use: - ``Nn``, ``Nm``, ``Np`` to decode flattened Hermite indices @@ -198,9 +260,7 @@ def _take_mode(n: int, m: int, p: int) -> jnp.ndarray: kinetic_energy_species = pref[None, :] * (term0[None, :] * C000 + term1 + term2) # (Nt, Ns) kinetic_energy = jnp.sum(kinetic_energy_species, axis=1) # (Nt,) # Field energy - rfft_weights = jnp.full(Nx_kept, 2.0) - rfft_weights = rfft_weights.at[0].set(1.0) - rfft_weights = rfft_weights.at[-1].set(1.0) + rfft_weights = _rfft_weights(output["Nx"], Nx_kept) weight_grid = rfft_weights.reshape(1, 1, -1, 1) EM_energy = 0.5 * jnp.sum((jnp.abs(Fk) ** 2) * weight_grid, axis=(-4, -3, -2, -1)) * Omega_cs[0] ** 2 diff --git a/spectrax/_simulation.py b/spectrax/_simulation.py index eba5de0..19ae472 100644 --- a/spectrax/_simulation.py +++ b/spectrax/_simulation.py @@ -186,6 +186,6 @@ def simulation(input_parameters={}, Nx=33, Ny=1, Nz=1, Nn=20, Nm=1, Np=1, Ns=2, # Output results temporary_output = {"Ck": Ck, "Fk": Fk, "time": time, "dCk": dCk, "solver_stats": sol.stats} - output = {**temporary_output, **parameters} + output = {**temporary_output, **parameters, "Nx": Nx} # static grid length for the rFFT weights diagnostics(output) return output diff --git a/tests/test_diagnostics.py b/tests/test_diagnostics.py new file mode 100644 index 0000000..6fd221c --- /dev/null +++ b/tests/test_diagnostics.py @@ -0,0 +1,112 @@ +import jax.numpy as jnp +import pytest + +from spectrax import diagnostics, hermite_tail_fraction + + +@pytest.mark.parametrize(("Nx", "expected_energy"), [(1, 0.5), (4, 0.5), (5, 1.0)]) +def test_rfft_field_energy_weights_nyquist_only_for_even_grid(Nx, expected_energy): + Nx_kept = Nx // 2 + 1 + output = { + "Nx": Nx, + "Nn": 3, + "Nm": 1, + "Np": 1, + "alpha_s": jnp.ones(6), + "u_s": jnp.zeros(6), + "Omega_cs": jnp.ones(2), + "mi_me": 1.0, + "Lx": 1.0, + "Ck": jnp.zeros((1, 6, 1, Nx_kept, 1), dtype=jnp.complex128), + "Fk": jnp.zeros((1, 6, 1, Nx_kept, 1), dtype=jnp.complex128) + .at[0, 0, 0, -1, 0] + .set(1.0), + } + + diagnostics(output) + + assert output["EM_energy"][0] == pytest.approx(expected_energy) + + +def test_hermite_tail_fraction_uses_odd_rfft_weights(): + dCk = jnp.zeros((1, 4, 1, 3, 1), dtype=jnp.complex128) + dCk = dCk.at[0, 0, 0, -1, 0].set(1.0) + dCk = dCk.at[0, 3, 0, 0, 0].set(1.0) + output = {"dCk": dCk, "Nx": 5, "Nn": 4, "Nm": 1, "Np": 1, + "Ns": 1, "u_s": jnp.zeros(3)} + + assert hermite_tail_fraction(output)[0, 0] == pytest.approx(1 / 3) + + +def test_hermite_tail_fraction_resolves_species_and_active_axes(): + Nn, Nm, Np, Ns = 3, 2, 2, 2 + modes = Nn * Nm * Np + dCk = jnp.zeros((1, Ns * modes, 1, 1, 1), dtype=jnp.complex128) + dCk = dCk.at[0, 0, 0, 0, 0].set(2.0) + dCk = dCk.at[0, Nn, 0, 0, 0].set(1.0) + dCk = dCk.at[0, modes, 0, 0, 0].set(1.0) + output = {"dCk": dCk, "Nx": 1, "Nn": Nn, "Nm": Nm, "Np": Np, + "Ns": Ns, "u_s": jnp.zeros(3 * Ns)} + + assert hermite_tail_fraction(output)[0] == pytest.approx(jnp.array([0.2, 0.0])) + + +def test_hermite_tail_fraction_ignores_inactive_axes_and_zero_norm(): + output = {"dCk": jnp.zeros((1, 1, 1, 1, 1)), "Nx": 1, + "Nn": 1, "Nm": 1, "Np": 1, "Ns": 1} + + assert hermite_tail_fraction(output)[0, 0] == 0.0 + + +def test_hermite_tail_fraction_rejects_invalid_width(): + with pytest.raises(ValueError, match="positive integer"): + hermite_tail_fraction({}, width=0) + + +def _field_output(Fk, Nx, Omega=1.3): + Nt, _, Ny, Nx_kept, Nz = Fk.shape + return {"Nx": Nx, "Nn": 1, "Nm": 1, "Np": 1, "alpha_s": jnp.ones(3), "u_s": jnp.zeros(3), + "Omega_cs": jnp.array([Omega]), "Lx": 1.0, + "Ck": jnp.zeros((Nt, 1, Ny, Nx_kept, Nz), dtype=jnp.complex128), "Fk": Fk} + + +@pytest.mark.parametrize("Ny, Nx, Nz", [(1, 1, 1), (1, 2, 1), (1, 7, 1), (1, 8, 1), (5, 6, 1), + (4, 9, 1), (1, 6, 5), (3, 5, 4), (4, 4, 4), (2, 3, 7)]) +def test_field_energy_matches_independent_full_fft_parseval(Ny, Nx, Nz): + """EM_energy from the stored half spectrum equals numpy's full-FFT sum and the real-space mean, + for odd and even Nx, full-FFT y/z axes, and singleton axes.""" + import numpy as np + F = np.random.default_rng(Nx + 10 * Ny + 100 * Nz).standard_normal((6, Ny, Nx, Nz)) + Fk = jnp.fft.rfftn(F, axes=(-1, -3, -2), norm="forward")[None] + output = _field_output(Fk, Nx) + diagnostics(output) + full = np.sum(np.abs(np.fft.fftn(F, axes=(1, 2, 3), norm="forward")) ** 2) + np.testing.assert_allclose(output["EM_energy"][0], 0.5 * 1.3**2 * full, rtol=1e-12) + np.testing.assert_allclose(output["EM_energy"][0], 0.5 * 1.3**2 * np.sum(np.mean(F**2, axis=(1, 2, 3))), + rtol=1e-12) + + +@pytest.mark.parametrize("Nx, Nx_kept", [(9, 3), (8, 4), (9, 6)]) +def test_field_energy_rejects_cropped_or_mismatched_rfft_axis(Nx, Nx_kept): + """A cropped/dealiased array is not a full rfft: the weights cannot be inferred and must be refused.""" + output = _field_output(jnp.zeros((1, 6, 1, Nx_kept, 1), dtype=jnp.complex128), Nx) + with pytest.raises(ValueError, match="full rfft"): + diagnostics(output) + + +def test_field_energy_requires_original_Nx(): + output = _field_output(jnp.zeros((1, 6, 1, 5, 1), dtype=jnp.complex128), 9) + del output["Nx"] + with pytest.raises(KeyError): + diagnostics(output) + + +def test_hermite_tail_fraction_is_per_species_for_three_species(): + """Each species' tail fraction uses only its own coefficients (3 species, unequal tails).""" + Nn, Ns = 4, 3 + dCk = jnp.zeros((1, Ns * Nn, 1, 2, 1), dtype=jnp.complex128) + for s, tail in enumerate([0.0, 1.0, 3.0]): + dCk = dCk.at[0, s * Nn + 1, 0, 1, 0].set(1.0) # bulk, kx=1 (weight 2 at Nx=3) + dCk = dCk.at[0, s * Nn + Nn - 1, 0, 1, 0].set(tail) # outermost Hermite mode + out = hermite_tail_fraction({"dCk": dCk, "Nx": 3, "Nn": Nn, "Nm": 1, "Np": 1, "Ns": Ns}) + assert out[0] == pytest.approx(jnp.array([0.0, 0.5, 0.9]))