diff --git a/spectrax/_model.py b/spectrax/_model.py index cdadf55..7fd721f 100644 --- a/spectrax/_model.py +++ b/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 - | / 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 ```` 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) diff --git a/tests/test_field_closure.py b/tests/test_field_closure.py new file mode 100644 index 0000000..c14d999 --- /dev/null +++ b/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)))