Add Hermite boundary-resolution diagnostic

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
 

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)