Use SOLVAX Newton-Krylov with a Hermite line preconditioner in implicit midpoint

Examples/1D_bump_on_tail.py

@@ -0,0 +1,59 @@
+"""Example: 1D bump-on-tail instability, with the beam expanded in the bulk's Hermite basis.
+
+A bulk Maxwellian (density 0.9, thermal speed 1) and a warm beam (density 0.1, thermal speed 0.5,
+drift 4.5) are shifted by -0.45 to the frame with no net current, so that SPECTRAX's Vlasov-Ampere
+system has no uniform plasma oscillation. Each population is its own species.
+
+The beam is expanded in the bulk's Hermite width (alpha = sqrt(2)), not its own (sqrt(2) * 0.5):
+`compute_C_nmp(..., vth_s=...)` gives the closed-form coefficients of a Maxwellian narrower than its
+basis. In its own narrow basis the trapped beam, which spreads over 2 < v < 6, has a divergent
+expansion: its coefficients grow by 1e2-1e7 and the adaptive step collapses (near t = 28 with nu = 3).
+Set BEAM_BASIS_VT = 0.5 to see it. Linear theory: omega = 0.866 + 0.198 i at k = 0.3.
+"""
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+
+from spectrax import compute_C_nmp, inverse_HF_transform, simulation
+
+jax.config.update("jax_enable_x64", True)
+BEAM_BASIS_VT = 1.0  # 0.5: the beam in its own basis (fails)
+species = [(0.9, -0.45, 1.0), (0.1, 4.05, 0.5)]  # (density, drift, thermal speed)
+k, eps, t_max, Nx, Nn = 0.3, 0.04, 100.0, 33, 128
+alpha = np.array([np.sqrt(2.0), np.sqrt(2.0) * BEAM_BASIS_VT])
+
+Ck_0 = []
+for (n, u, vth), a in zip(species, alpha):
+    C = n * compute_C_nmp(jnp.full((1, 3, 1, Nx, 1), u), [a] * 3, [0.0] * 3, Nn, 1, 1, 1,
+                          vth_s=[vth, a / np.sqrt(2), a / np.sqrt(2)])[0, 0, 0]  # (Nn, 1, Nx//2+1, 1)
+    Ck_0.append(C.at[:, :, 1].set(0.5 * eps * C[:, :, 0]))  # density n (1 + eps cos kx)
+parameters = {
+    "Lx": 2 * np.pi / k, "Ly": 1.0, "Lz": 1.0, "mi_me": 1.0, "Ti_Te": 1.0, "nx": 1, "t_max": t_max,
+    "qs": jnp.array([-1.0, -1.0]), "Omega_cs": jnp.array([1.0, 1.0]), "nu": 3.0, "D": 0.0, "ode_tolerance": 1e-8,
+    "alpha_s": jnp.repeat(jnp.asarray(alpha), 3), "u_s": jnp.array([0.0] * 6), "Ck_0": jnp.concatenate(Ck_0),
+    # Gauss's law for the density perturbation: E = -(eps / k) sin(kx).
+    "Fk_0": jnp.zeros((6, 1, Nx // 2 + 1, 1), dtype=jnp.complex128).at[0, 0, 1, 0].set(0.5j * eps / k),
+}
+output = simulation(parameters, Nx=Nx, Nn=Nn, Ns=2, timesteps=501, dt=0.01)
+
+t, E1 = output["time"], 2 * jnp.abs(output["Fk"][:, 0, 0, 1, 0])
+x, v = np.linspace(0, 2 * np.pi / k, Nx, endpoint=False), np.linspace(-8, 10, 241)
+f = 0
+for s, a in enumerate(alpha):  # 1D marginal: f(v_x, v_y = v_z = 0) times the transverse integral a^2 pi
+    xi_y, xi_x, xi_z = np.meshgrid([0.0], v / a, [0.0], indexing="ij")
+    f += np.pi * a**2 * inverse_HF_transform(output["Ck"][-1:, s * Nn:(s + 1) * Nn], Nn, 1, 1, Nx, 1, 1,
+                                             xi_x, xi_y, xi_z)[0, 0, :, 0, 0, :, 0]
+
+fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 3.8))
+ax1.semilogy(t, E1, label="SPECTRAX")
+linear = (t > 8) & (t < 20)
+ax1.semilogy(t[linear], E1[linear][0] * np.exp(0.198 * (t[linear] - 8)), "k--", label=r"$\gamma = 0.198$")
+ax1.set(xlabel=r"$t\,\omega_{pe}$", ylabel=r"$|\hat E_{k}|$", title="bump-on-tail, k = 0.3")
+ax1.legend()
+mesh = ax2.pcolormesh(x, v, np.asarray(f).T, shading="auto", cmap="magma")
+ax2.set(xlabel=r"$x/\lambda_D$", ylabel=r"$v/v_t$", title=f"f(x, v, t = {t_max:g})")
+fig.colorbar(mesh, ax=ax2)
+plt.tight_layout()
+plt.show()

Examples/2D_magnetic_reconnection.py

@@ -0,0 +1,228 @@
+"""Periodic double-Harris-sheet magnetic reconnection.
+
+The short defaults are a smoke test. For a clear nonlinear island movie, run
+
+    python Examples/2D_magnetic_reconnection.py --nx 65 --ny 17 \
+        --t-max 500 --chunk-time 50 --snapshots 51 --perturbation 0.1
+
+Use ``--perturbation 1e-4`` instead for a linear tearing-growth diagnostic.
+The parameters follow Camporeale et al. (2006), section IV.C; the periodic
+double-sheet geometry and half-sheet magnetic diagnostic follow Koshkarov et
+al. (2021).
+"""
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+from diffrax import Dopri8
+from matplotlib.animation import FuncAnimation, PillowWriter
+
+from spectrax import simulation
+
+
+def initial_conditions(nx, ny, hermite, perturbation, nu):
+    a = jnp.sqrt(2.0)
+    lx, ly = 32.0 * a, 4.0 * jnp.pi * a
+    x = jnp.linspace(0.0, lx, nx, endpoint=False)
+    y = jnp.linspace(0.0, ly, ny, endpoint=False)
+    X, Y, Z = jnp.meshgrid(x, y, jnp.zeros(1), indexing="xy")
+    sheets = (
+        jnp.cosh((X - lx / 4.0) / a) ** -2,
+        jnp.cosh((X - 3.0 * lx / 4.0) / a) ** -2,
+    )
+    density = jnp.stack((sheets[0], sheets[1], sheets[0], sheets[1]))
+    alpha = a * jnp.array(([1.0 / 6.0] * 6) + ([1.0 / 24.0] * 6))
+    coefficients = jnp.zeros(
+        (4, hermite, hermite, hermite, ny, nx, 1), dtype=jnp.float64
+    )
+    coefficients = coefficients.at[:, 0, 0, 0].set(
+        density / jnp.prod(alpha.reshape(4, 3), axis=1)[:, None, None, None]
+    )
+
+    kx, ky = 2.0 * jnp.pi / lx, 2.0 * jnp.pi / ly
+    magnetic = jnp.array(
+        [
+            -perturbation * ky * jnp.sin(ky * Y) * jnp.sin(kx * X),
+            jnp.tanh((X - lx / 4.0) / a)
+            - jnp.tanh((X - 3.0 * lx / 4.0) / a)
+            - 1.0
+            - perturbation * kx * jnp.cos(ky * Y) * jnp.cos(kx * X),
+            5.0 * jnp.ones_like(X),
+        ]
+    )
+    transform = lambda value: jnp.fft.rfftn(
+        value, axes=(-1, -3, -2), norm="forward"
+    )
+    return {
+        "Lx": lx,
+        "Ly": ly,
+        "Lz": 1.0,
+        "mi_me": 16.0,
+        "qs": jnp.array((-1.0, -1.0, 1.0, 1.0)),
+        "alpha_s": alpha,
+        "u_s": jnp.array(
+            (0.0, 0.0, -1.0 / 6.0, 0.0, 0.0, 1.0 / 6.0,
+             0.0, 0.0, 1.0 / 6.0, 0.0, 0.0, -1.0 / 6.0)
+        ) / a,
+        "Omega_cs": jnp.array((1.0 / 3.0, 1.0 / 3.0, 1.0 / 48.0, 1.0 / 48.0)),
+        "nu": nu,
+        "D": 0.0,
+        "ode_tolerance": 1e-7,
+        "Ck_0": transform(coefficients),
+        "Fk_0": transform(jnp.concatenate((jnp.zeros_like(magnetic), magnetic))),
+    }
+
+
+def fields_and_flux(fk, nx, lx, ly):
+    ny = fk.shape[-3]
+    fields = np.fft.irfftn(
+        fk, s=(1, ny, nx), axes=(-1, -3, -2), norm="forward"
+    ).real
+    kx = 2.0 * np.pi * np.fft.rfftfreq(nx, d=lx / nx)
+    ky = 2.0 * np.pi * np.fft.fftfreq(ny, d=ly / ny)
+    KX, KY = np.meshgrid(kx, ky)
+    bx, by = fk[:, 3, :, :, 0], fk[:, 4, :, :, 0]
+    k2 = KX**2 + KY**2
+    az_k = np.zeros_like(bx)
+    np.divide(-1j * KY * bx + 1j * KX * by, k2, out=az_k, where=k2 != 0.0)
+    jz_k = 1j * KX * by - 1j * KY * bx
+    az = np.fft.irfftn(az_k, s=(ny, nx), axes=(-2, -1), norm="forward").real
+    jz = np.fft.irfftn(jz_k, s=(ny, nx), axes=(-2, -1), norm="forward").real
+    return fields, az, jz
+
+
+def save_diagnostics(output_dir, t, fk, nx, lx, ly):
+    fields, az, jz = fields_and_flux(fk, nx, lx, ly)
+    half = (nx + 1) // 2
+    magnetic = fields[:, 3:, :, :half, 0]
+    phase = np.exp(2j * np.pi * np.arange(magnetic.shape[2]) / magnetic.shape[2])
+    amplitude = np.linalg.norm(
+        np.sum(magnetic * phase[None, None, :, None], axis=(-2, -1)), axis=1
+    ) / np.sqrt(magnetic.shape[2] * magnetic.shape[3])
+    peaks = np.flatnonzero(
+        (amplitude[1:-1] > amplitude[:-2]) & (amplitude[1:-1] > amplitude[2:])
+    ) + 1
+    peaks = peaks[t[peaks] >= t[-1] / 4.0]
+    fit = None
+    if len(peaks) >= 3:
+        gamma, intercept = np.polyfit(t[peaks], np.log(amplitude[peaks]), 1)
+        fit = np.exp(intercept + gamma * t)
+    else:
+        gamma = np.nan
+
+    sheet = np.argmin(abs(np.linspace(0.0, lx, nx, endpoint=False) - lx / 4.0))
+    flux = np.ptp(az[:, :, sheet], axis=1)
+    fig, axes = plt.subplots(2, 1, figsize=(6.4, 6.0), sharex=True, constrained_layout=True)
+    axes[0].semilogy(t, amplitude, color="0.65", label="half-sheet magnetic mode")
+    if fit is not None:
+        axes[0].semilogy(t[peaks], amplitude[peaks], "o", ms=3)
+        axes[0].semilogy(t, fit, "k--", label=rf"envelope $\gamma={gamma:.5f}$")
+    axes[0].set_ylabel("magnetic-mode amplitude")
+    axes[0].legend(fontsize=8)
+    axes[1].plot(t, flux)
+    axes[1].set(xlabel=r"$\omega_{pe}t$", ylabel=r"reconnected flux $\Delta A_z$")
+    fig.savefig(output_dir / "reconnection-growth.png", dpi=180)
+    plt.close(fig)
+
+    x = np.linspace(0.0, lx, nx, endpoint=False)[:half]
+    y = np.linspace(-ly / 2.0, ly / 2.0, fk.shape[-3], endpoint=False)
+    az = np.roll(az, fk.shape[-3] // 2, axis=1)[:, :, :half]
+    jz = np.roll(jz, fk.shape[-3] // 2, axis=1)[:, :, :half]
+    window = (x >= lx / 8.0) & (x <= 3.0 * lx / 8.0)
+    x, az, jz = x[window], az[:, :, window], jz[:, :, window]
+    jz -= jz.mean(axis=1, keepdims=True)
+    limit = max(np.quantile(abs(jz[-1]), 0.99), 1e-12)
+    sheet = np.argmin(abs(x - lx / 4.0))
+    fig, ax = plt.subplots(figsize=(6.4, 3.8), constrained_layout=True)
+    image = ax.pcolormesh(x, y, jz[0], shading="auto", cmap="RdBu_r", vmin=-limit, vmax=limit)
+    contours = [None]
+    title = ax.set_title("")
+    ax.set(xlabel="x", ylabel="y")
+    fig.colorbar(image, ax=ax, label=r"$J_z-\langle J_z\rangle_y$")
+
+    def update(frame):
+        image.set_array(jz[frame].ravel())
+        if contours[0] is not None:
+            contours[0].remove()
+        center = az[frame, :, sheet].mean()
+        span = max(1.25 * np.ptp(az[frame, :, sheet]), 1e-12)
+        contours[0] = ax.contour(
+            x, y, az[frame], center + np.linspace(-span, span, 19),
+            colors="k", linewidths=0.55,
+        )
+        title.set_text(rf"$\omega_{{pe}}t={t[frame]:.1f}$")
+        return image, title, contours[0]
+
+    frames = range(0, len(t), max(1, len(t) // 100))
+    animation = FuncAnimation(fig, update, frames=frames, interval=80, blit=False)
+    animation.save(output_dir / "reconnection-topology.gif", PillowWriter(fps=12), dpi=110)
+    plt.close(fig)
+    return amplitude, gamma
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--nx", type=int, default=33)
+    parser.add_argument("--ny", type=int, default=9)
+    parser.add_argument("--hermite", type=int, default=4)
+    parser.add_argument("--nu", type=float, default=1.0)
+    parser.add_argument("--perturbation", type=float, default=1e-4)
+    parser.add_argument("--t-max", type=float, default=10.0)
+    parser.add_argument("--chunk-time", type=float, default=10.0)
+    parser.add_argument("--snapshots", type=int, default=21)
+    parser.add_argument("--output-dir", type=Path, default=Path("reconnection_output"))
+    args = parser.parse_args()
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    values = initial_conditions(args.nx, args.ny, args.hermite, args.perturbation, args.nu)
+    times, fields, offset = [], [], 0.0
+    start = time.perf_counter()
+    while offset < args.t_max:
+        duration = min(args.chunk_time, args.t_max - offset)
+        values["t_max"] = duration
+        output = simulation(
+            values, Nx=args.nx, Ny=args.ny, Nz=1,
+            Nn=args.hermite, Nm=args.hermite, Np=args.hermite, Ns=4,
+            timesteps=args.snapshots, dt=1e-3, solver=Dopri8(),
+        )
+        jax.block_until_ready(output["Ck"])
+        first = 0 if not times else 1
+        times.append(np.asarray(output["time"])[first:] + offset)
+        fields.append(np.asarray(output["Fk"])[first:])
+        values["Ck_0"] = jnp.asarray(output["Ck"][-1])
+        values["Fk_0"] = jnp.asarray(output["Fk"][-1])
+        del output
+        offset += duration
+        print(f"t={offset:g}")
+
+    t, fk = np.concatenate(times), np.concatenate(fields)
+    amplitude, gamma = save_diagnostics(
+        args.output_dir, t, fk, args.nx, float(values["Lx"]), float(values["Ly"])
+    )
+    coefficients = np.asarray(values["Ck_0"]).reshape(
+        4, args.hermite, args.hermite, args.hermite, args.ny, args.nx // 2 + 1, 1
+    )
+    power = abs(coefficients) ** 2
+    tails = {
+        name: float(np.take(power, -1, axis=axis).sum() / power.sum())
+        for name, axis in (("p", 1), ("m", 2), ("n", 3))
+    }
+    np.savez(args.output_dir / "reconnection.npz", time=t, amplitude=amplitude, Fk=fk)
+    linear = args.perturbation <= 1e-3
+    print(json.dumps({
+        "runtime_s": time.perf_counter() - start,
+        "regime": "linear diagnostic" if linear else "finite-seed nonlinear",
+        "gamma": None if not np.isfinite(gamma) else float(gamma),
+        "expected_gamma": 0.256 / 48.0 if linear else None,
+        "hermite_tail_fraction": tails,
+    }, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

README.md

@@ -201,6 +201,9 @@ most important keys. Keys absent from the file fall back to sensible defaults sp
 | `dt` | Initial step size provided to the ODE solver. |
 | `solver` | Name of Diffrax solver (e.g., `Tsit5`, `Dopri5`, `Dopri8`, `ImplicitMidpoint`). |
 | `adaptive_time_step` | Timestep adaptability (`true` by default). |
+| `dtmin` | Smallest adaptive step; below it the solve stops with an error instead of stalling. Default `t_max / max_steps`, `0` disables it. A collapsing step means an under-resolved run: more Hermite modes, a larger `nu`, or a wider `alpha_s`. |
+| `max_steps` | Maximum number of solver steps (`1000000` by default). |
+| `throw` | `True` (default) raises on a failed solve; `False` returns and reports it in `output["solver_result"]`. |
 
 Many of these parameters can be arrays; for example `alpha_s` must contain three values per species (one
 for each velocity dimension) and can be used to represent anisotropic plasmas.

pyproject.toml

@@ -29,7 +29,7 @@ keywords = ["Plasma",
      "Simulation",
      "JAX"]
 
-dependencies = [ "jax", "jaxlib", "jax_tqdm", "matplotlib"]
+dependencies = [ "jax", "jaxlib", "jax_tqdm", "matplotlib", "diffrax", "orthax", "solvax>=0.8.0"]
 
 requires-python = ">=3.5"
 

requirements.txt

@@ -3,5 +3,6 @@ jaxlib
 jax_tqdm
 diffrax
 orthax
+solvax>=0.8.0
 matplotlib
 pandas

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