Field-scaled closure rate for the asymmetric-Hermite top-mode instability

spectrax/_model.py

@@ -10,7 +10,8 @@
 from jax import jit
 from functools import partial
 
-__all__ = ['plasma_current', 'Hermite_Fourier_system', 'basis_rate_terms', 'uniform_acceleration']
+__all__ = ['plasma_current', 'Hermite_Fourier_system', 'basis_rate_terms', 'uniform_acceleration',
+           'field_scaled_closure_rate']
 
 
 @partial(jit, static_argnames=['Nn', 'Nm', 'Np', 'Ns'])
@@ -125,8 +126,9 @@ def Hermite_Fourier_system(Ck, C, F, kx_grid, ky_grid, kz_grid, k2_grid, col,
         Square-root ladder coefficients for the Hermite recurrences along each axis.
     Lx, Ly, Lz : float
         Domain lengths in each spatial direction.
-    nu : float
-        Collision frequency.
+    nu : float or jnp.ndarray
+        Collision frequency: a scalar, or one rate per species with shape ``(Ns, 1, 1, 1, 1, 1, 1)``
+        (see :func:`field_scaled_closure_rate`).
     D : float
         Hyper-diffusion coefficient.
     alpha_s, u_s : jnp.ndarray
@@ -293,3 +295,40 @@ def uniform_acceleration(F, u_s, qs, Omega_cs, Ns):
     u = u_s.reshape(Ns, 3)
     accel = (qs * Omega_cs)[:, None] * (F0[None, :3] + jnp.cross(u, F0[None, 3:]))
     return accel.reshape(3 * Ns), F0
+
+
+def field_scaled_closure_rate(F, alpha_s, qs, Omega_cs, Nn, Nm, Np, Ns, c):
+    """
+    Per-species closure rate that follows the field, for the asymmetric-Hermite top-mode instability.
+
+    The truncated asymmetrically weighted (AW) Fourier-Hermite operator ``-v d/dx + (q/m) E d/dv`` is not
+    anti-self-adjoint: in a non-uniform field it has spurious growing modes that live in the top third of the
+    Hermite ladder at the largest retained ``|k|``, with growth rate of order ``0.3 sqrt(2N) |q/m| |E| / a``
+    (rising with N and with the Fourier cut-off; the symmetrically weighted basis gives exactly zero). A fixed
+    ``nu`` is overtaken once the self-consistent field grows. This returns
+
+        nu_s = c |q_s/m_s| max_i sqrt(2 N_i) max_x |E_i - <E_i>| / a_{s,i},
+
+    over the velocity axes with ``N_i > 3``, shaped ``(Ns, 1, 1, 1, 1, 1, 1)`` so it can be passed as ``nu`` to
+    :func:`Hermite_Fourier_system` together with ``hypercollision_spectrum(order>=2)``, which leaves density,
+    momentum and energy untouched. The uniform part ``<E>`` is excluded: it only shifts the distribution.
+    ``c = 0`` gives exactly zero.
+
+    Parameters
+    ----------
+    F : jnp.ndarray, shape (6, Ny, Nx, Nz)
+        Real-space fields (de-aliased), as used by the kinetic RHS.
+    alpha_s : jnp.ndarray, shape (3 * Ns,)
+    qs, Omega_cs : jnp.ndarray, shape (Ns,)
+        ``|q/m|_s = |qs * Omega_cs|`` in SPECTRAX units.
+    c : float
+        Dimensionless strength.
+    """
+    E = F[:3]
+    dE = jnp.max(jnp.abs(E - jnp.mean(E, axis=(-3, -2, -1), keepdims=True)), axis=(-3, -2, -1))  # (3,)
+    N = jnp.array([Nn, Nm, Np], dtype=float)
+    a = alpha_s.reshape(Ns, 3)
+    active = (N > 3)[None, :]
+    per_axis = jnp.where(active, jnp.sqrt(2 * N)[None, :] * dE[None, :] / a, 0.0)
+    rate = c * jnp.abs(qs * Omega_cs) * jnp.max(per_axis, axis=1)
+    return rate.reshape(Ns, 1, 1, 1, 1, 1, 1)

tests/test_field_closure.py

@@ -0,0 +1,51 @@
+"""Field-scaled closure rate: exact zero at c = 0, analytic value, per-species nu in the kinetic RHS, low moments."""
+import numpy as np
+import jax
+import jax.numpy as jnp
+from spectrax import field_scaled_closure_rate, hypercollision_spectrum
+from spectrax._remap import low_moments
+from tests.test_moving_basis import ORDERS, random_state, ode_args, rhs
+
+jax.config.update("jax_enable_x64", True)
+
+
+def _fields(Nx=16, E1=3e-3, E0=0.7):
+    x = np.arange(Nx) * 2 * np.pi / Nx
+    F = np.zeros((6, 1, Nx, 1))
+    F[0, 0, :, 0] = E0 + E1 * np.cos(x)
+    return jnp.asarray(F)
+
+
+def test_zero_strength_is_exactly_zero():
+    a = jnp.full(6, 0.05)
+    r = field_scaled_closure_rate(_fields(), a, jnp.array([-1.0, 1.0]), jnp.array([1.0, 1 / 1836]), 64, 1, 1, 2, 0.0)
+    assert r.shape == (2, 1, 1, 1, 1, 1, 1) and float(jnp.abs(r).max()) == 0.0
+
+
+def test_analytic_value_excludes_uniform_field_and_short_axes():
+    a = jnp.array([0.05, 0.02, 0.02, 0.001, 0.001, 0.001])
+    qs, om = jnp.array([-1.0, 1.0]), jnp.array([1.0, 1 / 1836])
+    r = np.asarray(field_scaled_closure_rate(_fields(), a, qs, om, 64, 1, 1, 2, 0.5)).ravel()
+    expect = 0.5 * np.array([1.0, 1 / 1836]) * np.sqrt(128) * 3e-3 / np.array([0.05, 0.001])
+    np.testing.assert_allclose(r, expect, rtol=1e-14)  # E0 = 0.7 does not enter; Nm = Np = 1 axes are inactive
+
+
+def test_jit_and_per_species_nu_in_rhs_keeps_low_moments():
+    """Passing the rate as nu gives dC - dC(nu=0) = -nu_s col C per species, with zero rate of n, M_i, M_ii."""
+    Ns = 2
+    Nn, Nm, Np = ORDERS
+    Ck, Fk, u, a = random_state(Ns, seed=3)
+    nu_s = jnp.array([0.3, 1.7]).reshape(Ns, 1, 1, 1, 1, 1, 1)
+    base = list(ode_args(Ns, u, a, nu=0.0))
+    col = hypercollision_spectrum(Nn, Nm, Np, 2)
+    base[7 + 14] = col
+    with_nu = list(base)
+    with_nu[7 + 1] = nu_s
+    dC0, _ = rhs(Ns, (Ck, Fk), tuple(base))
+    dC1, _ = rhs(Ns, (Ck, Fk), tuple(with_nu))
+    np.testing.assert_allclose(dC1 - dC0, -nu_s * col[None, :, :, :, None, None, None] * Ck, rtol=1e-12, atol=1e-15)
+    B = jnp.stack([u, a])
+    assert float(jnp.abs(low_moments(dC1 - dC0, B, Nn, Nm, Np, Ns)).max()) < 1e-14 * float(jnp.abs(Ck).max())
+    f = jax.jit(field_scaled_closure_rate, static_argnames=("Nn", "Nm", "Np", "Ns"))
+    r = f(_fields(), jnp.full(6, 0.05), jnp.array([-1.0, 1.0]), jnp.array([1.0, 0.1]), Nn=64, Nm=1, Np=1, Ns=2, c=1.0)
+    assert np.all(np.isfinite(np.asarray(r)))