diff --git a/spectrax/_diagnostics.py b/spectrax/_diagnostics.py index 8fdb788..3d18bc6 100644 --- a/spectrax/_diagnostics.py +++ b/spectrax/_diagnostics.py @@ -19,7 +19,13 @@ import jax.numpy as jnp -__all__ = ["diagnostics"] +__all__ = ["diagnostics", "hermite_tail_fraction"] + + +def _rfft_weights(Nx: int, Nx_kept: int) -> jnp.ndarray: + """Return Parseval weights for an RFFT of length ``Nx``.""" + weights = jnp.full(Nx_kept, 2.0).at[0].set(1.0) + return weights.at[-1].set(1.0 + (Nx % 2) * (Nx_kept > 1)) def _infer_Ns(output: Mapping[str, Any]) -> int: @@ -68,6 +74,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(int(output.get("Nx", 2 * (Nx_kept - 1))), 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). @@ -198,9 +248,8 @@ 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) + Nx = output.get("Nx", 2 * (Nx_kept - 1)) + rfft_weights = _rfft_weights(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/tests/test_diagnostics.py b/tests/test_diagnostics.py new file mode 100644 index 0000000..cac4234 --- /dev/null +++ b/tests/test_diagnostics.py @@ -0,0 +1,63 @@ +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)