Fix odd-grid RFFT energy weights; add a Hermite boundary diagnostic

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
 

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

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]))