Add opt-in regular harmonics for Boozer interpolation

README.md

@@ -203,6 +203,21 @@ is lost at `s = 1`; `result.loss_times` and `result.states` hold when and where.
 Boozer tracing evaluates `|B| = sum [bc_mn(s) cos(m theta - n zeta) + bs_mn(s) sin(m theta - n zeta)]`
 and its derivatives, together with the flux functions `G`, `I` and `iota`.
 
+For smooth near-axis spectra, `axis_regular=True` interpolates each harmonic
+as `s**(m/2)` times a cubic spline in `s`, preserving the input samples:
+
+```python
+field = BoozerField.from_booz_xform(
+    booz, psi0=-phi_edge / (2 * np.pi),
+    axis_regular=True, axis_regular_tolerance=0.01)
+```
+
+Choose the positive consistency tolerance before tracing (`0.01` is an example).
+The option requires at least five positive increasing surfaces and rejects inconsistent
+scaled inner profiles; scaling can amplify noise. Passing this local check does
+not qualify the equilibrium or bound extrapolation error. It uses direct Fourier
+sums regardless of `angle_grid`; the default interpolation is unchanged.
+
 ## Comparison with other codes
 
 256 collisionless 3.5 MeV alphas per equilibrium, started at `s = 0.5` with

essos/boozer.py

@@ -1,7 +1,8 @@
 """Fixed-step K=0 Boozer guiding-centre tracing (White, 2014).
 
-Interpolate the |B| spectrum in sqrt(s), with m>=1 coefficients divided by
-sqrt(s) for axis regularity; interpolate iota, G and I in s. Explicit RK runs in
+By default interpolate the |B| spectrum in sqrt(s), with m>=1 coefficients
+divided by sqrt(s). An explicit axis-regular option uses s**(m/2) times
+splines in s; interpolate iota, G and I in s. Explicit RK runs in
 sqrt(s)*(cos(theta), sin(theta)). LCFS crossings are lost; nonfinite steps
 are failed. Optional Euler-Maruyama collisions apply pitch scattering,
 slowing down and energy diffusion (Boozer & Kuo-Petravic, 1981).
@@ -53,7 +54,7 @@ def _evaluate(knots, coef, x):
 class BoozerField(eqx.Module):
     """``|B|`` spectrum and ``iota, G, I`` profiles in Boozer coordinates."""
 
-    r_knots: jax.Array
+    r_knots: jax.Array  # sqrt(s), or s with axis_regular=True
     b_coef: jax.Array  # (nint, modes, 4)
     s_knots: jax.Array
     profile_coef: jax.Array  # (nint, 3, 4): iota, G, I
@@ -63,10 +64,11 @@ class BoozerField(eqx.Module):
     nfp: int = eqx.field(static=True)
     sine_coef: jax.Array | None = None
     table: jax.Array | None = None  # (nint, ntheta + DEGREE, nzeta + DEGREE, 2, 4)
+    axis_regular: bool = eqx.field(static=True, default=False)
 
     @classmethod
     def from_booz(cls, s, bmnc, xm, xn, iota, G, I, psi0, nfp, mode_tolerance=1e-6, *, bmns=None,
-                  angle_grid=6):
+                  angle_grid=6, axis_regular=False, axis_regular_tolerance=None):
         """Build from half-mesh Boozer tables.
 
         ``bmnc`` and optional sine coefficients ``bmns`` are ``(modes, len(s))``;
@@ -83,6 +85,14 @@ def from_booz(cls, s, bmnc, xm, xn, iota, G, I, psi0, nfp, mode_tolerance=1e-6,
         used only above ``TABLE_MODES`` modes, below which the sum is cheaper. ``angle_grid=None``, or a
         toroidal mode number that is not a multiple of ``nfp``, sums the modes at
         every evaluation.
+
+        ``axis_regular=True`` uses ``s**(m/2)`` times splines in ``s`` and
+        direct Fourier sums. It requires at least five positive increasing
+        surfaces and a positive ``axis_regular_tolerance``: for each retained m>0 component,
+        quadratic extrapolation from scaled surfaces 2--5 must match surface 1
+        within this relative local-amplitude tolerance. Scaling can amplify
+        noise; this diagnostic neither qualifies the input equilibrium nor
+        bounds extrapolation error. Samples and default behavior are unchanged.
         """
         s = np.asarray(s, float)
         bmnc = np.asarray(bmnc, float)
@@ -94,14 +104,41 @@ def from_booz(cls, s, bmnc, xm, xn, iota, G, I, psi0, nfp, mode_tolerance=1e-6,
         keep = amplitude > mode_tolerance * amplitude.max()
         bmnc, xm, xn = bmnc[keep], xm[keep], np.asarray(xn, int)[keep]
         r = np.sqrt(s)
-        scaled = np.where(xm[:, None] > 0, bmnc / r, bmnc)
-        r_knots, b_coef = _spline(r, scaled.T)
+        if axis_regular:
+            if (s.ndim != 1 or len(s) < 5 or not np.isfinite(s).all()
+                    or np.any(s <= 0) or np.any(np.diff(s) <= 0) or np.any(xm < 0)):
+                raise ValueError("axis_regular requires at least five positive increasing surfaces and m >= 0")
+            if (axis_regular_tolerance is None or not np.isfinite(axis_regular_tolerance)
+                    or axis_regular_tolerance <= 0):
+                raise ValueError("axis_regular requires a positive finite axis_regular_tolerance")
+            if np.min(np.log(r)) * np.max(xm) < np.log(np.finfo(float).tiny):
+                raise ValueError("axis_regular mode scaling underflows at the first surface")
+        divisor = r[None, :] ** (xm[:, None] if axis_regular else (xm[:, None] > 0))
+        with np.errstate(over="ignore", divide="ignore", invalid="ignore"):
+            scaled = bmnc / divisor
+            scaled_sine = None if sine is None else sine[keep] / divisor
+        if axis_regular:
+            for component, values in (("cosine", scaled), ("sine", scaled_sine)):
+                if values is None:
+                    continue
+                if not np.isfinite(values).all():
+                    raise ValueError("axis_regular scaling is not finite")
+                for mode in np.flatnonzero(xm > 0):
+                    row = values[mode, :5]
+                    estimate = np.polyval(np.polyfit(s[1:5], row[1:5], 2), s[0])
+                    residual = abs(row[0] - estimate) / max(np.max(np.abs(row)), np.finfo(float).tiny)
+                    if residual > axis_regular_tolerance:
+                        raise ValueError(
+                            f"axis_regular inner profile is unresolved for m={xm[mode]}, "
+                            f"n={xn[mode]} ({component} relative discrepancy {residual:.6g} "
+                            f"> {axis_regular_tolerance:.6g})")
+        coordinates = s if axis_regular else r
+        knots, b_coef = _spline(coordinates, scaled.T)
         sine_coef = None
         if sine is not None and np.any(sine[keep]):
-            scaled_sine = np.where(xm[:, None] > 0, sine[keep] / r, sine[keep])
-            _, sine_coef = _spline(r, scaled_sine.T)
+            _, sine_coef = _spline(coordinates, scaled_sine.T)
         table = None
-        if (angle_grid is not None and np.all(xn % int(nfp) == 0)
+        if (not axis_regular and angle_grid is not None and np.all(xn % int(nfp) == 0)
                 and (np.ndim(angle_grid) or xm.size > TABLE_MODES)):
             nt, nz = ((max(16, angle_grid * int(np.abs(xm).max())),
                        max(16, angle_grid * int(np.abs(xn).max()) // int(nfp)))
@@ -115,18 +152,20 @@ def from_booz(cls, s, bmnc, xm, xn, iota, G, I, psi0, nfp, mode_tolerance=1e-6,
         s_prof = np.concatenate([[0.0], s])
         profiles = np.vstack([axis, profiles])
         s_knots, profile_coef = _spline(s_prof, profiles)
-        return cls(r_knots, b_coef, s_knots, profile_coef, jnp.asarray(xm), jnp.asarray(xn),
-                   float(psi0), int(nfp), sine_coef, table)
+        return cls(knots, b_coef, s_knots, profile_coef, jnp.asarray(xm), jnp.asarray(xn),
+                   float(psi0), int(nfp), sine_coef, table, bool(axis_regular))
 
     @classmethod
-    def from_booz_xform(cls, booz, psi0, mode_tolerance=1e-6, angle_grid=6):
+    def from_booz_xform(cls, booz, psi0, mode_tolerance=1e-6, angle_grid=6, *,
+                        axis_regular=False, axis_regular_tolerance=None):
         """Build from a run ``Booz_xform`` (every surface computed)."""
         sine = getattr(booz, "bmns_b", None)
         if bool(getattr(booz, "asym", False)) and sine is None:
             raise ValueError("Asymmetric Boozer transform requires bmns_b")
         return cls.from_booz(booz.s_b, booz.bmnc_b, booz.xm_b, booz.xn_b, booz.iota,
                              booz.Boozer_G, booz.Boozer_I, psi0, int(booz.nfp), mode_tolerance,
-                             bmns=sine, angle_grid=angle_grid)
+                             bmns=sine, angle_grid=angle_grid, axis_regular=axis_regular,
+                             axis_regular_tolerance=axis_regular_tolerance)
 
     def profiles(self, s):
         """``(iota, G, I)`` and their ``s`` derivatives."""
@@ -136,6 +175,22 @@ def modB_derivatives(self, r, theta, zeta):
         """``|B|``, ``d|B|/dr``, ``(d|B|/dtheta)/r`` and ``d|B|/dzeta``."""
         if self.table is not None:
             return _table_derivatives(self.r_knots, self.table, self.nfp, r, theta, zeta)
+        if self.axis_regular:
+            a, da = _evaluate(self.r_knots, self.b_coef, r*r)
+            phase = self.xm * theta - self.xn * zeta
+            c, s = jnp.cos(phase), jnp.sin(phase)
+            power = r ** self.xm
+            angular_power = r ** jnp.maximum(self.xm - 1, 0)
+            df = self.xm * angular_power * a + 2*r*power*da
+            values = (jnp.sum(power*a*c), jnp.sum(df*c),
+                      -jnp.sum(self.xm*angular_power*a*s), jnp.sum(self.xn*power*a*s))
+            if self.sine_coef is not None:
+                a, da = _evaluate(self.r_knots, self.sine_coef, r*r)
+                df = self.xm * angular_power * a + 2*r*power*da
+                values = tuple(x+y for x, y in zip(values, (
+                    jnp.sum(power*a*s), jnp.sum(df*s),
+                    jnp.sum(self.xm*angular_power*a*c), -jnp.sum(self.xn*power*a*c))))
+            return values
         a, da = _evaluate(self.r_knots, self.b_coef, r)
         phase = self.xm * theta - self.xn * zeta
         c, s = jnp.cos(phase), jnp.sin(phase)

tests/test_boozer.py

@@ -1,6 +1,7 @@
 """Boozer-coordinate guiding-centre tracing and Monte Carlo collisions."""
 
 import dataclasses
+from types import SimpleNamespace
 
 import jax
 import jax.numpy as jnp
@@ -36,6 +37,98 @@ def test_field_reproduces_the_spectrum_and_is_regular_on_the_axis():
     assert (iota, G_s, I_s) == pytest.approx((IOTA, G, 0.0))
 
 
+def test_explicit_axis_regular_modes_match_analytic_field_and_derivatives():
+    s = np.linspace(.0025, .9975, 200)
+    r = np.sqrt(s)
+    bmnc = np.stack([2 + .1*s, .2*r*(1+.3*s), .08*r**5*np.exp(.2*s)])
+    bmns = np.stack([np.zeros_like(s), np.zeros_like(s), .04*r**5*np.exp(.2*s)])
+    booz = SimpleNamespace(s_b=s, bmnc_b=bmnc, bmns_b=bmns, xm_b=np.array([0, 1, 5]),
+                           xn_b=np.array([0, 0, 2]), iota=np.full_like(s, .4),
+                           Boozer_G=np.ones_like(s), Boozer_I=np.zeros_like(s), nfp=2)
+    field = BoozerField.from_booz_xform(booz, 1., angle_grid=(16, 16),
+                                      axis_regular=True, axis_regular_tolerance=1e-8)
+    assert field.axis_regular and field.table is None
+    # Preserve all supplied samples, including sine modes, without a fit.
+    sample_theta, sample_zeta = .7, .2
+    sample_phase = np.array([0, 1, 5])*sample_theta-np.array([0, 0, 2])*sample_zeta
+    expected_samples = [email protected](sample_phase)[email protected](sample_phase)
+    np.testing.assert_allclose(jax.vmap(field.modB)(s, jnp.full_like(s, sample_theta),
+                                                  jnp.full_like(s, sample_zeta)),
+                               expected_samples, rtol=2e-15, atol=2e-15)
+    theta, zeta = .7, .2
+    phase = 5*theta-2*zeta
+    harmonic = .08*np.cos(phase)+.04*np.sin(phase)
+    angular = -.4*np.sin(phase)+.2*np.cos(phase)
+    toroidal = .16*np.sin(phase)-.08*np.cos(phase)
+    evaluate = jax.jit(lambda r, t, z: field.modB_derivatives(r, t, z))
+    radial_grad = jax.jit(jax.grad(lambda x: field.modB_derivatives(x, theta, zeta)[0]))
+    for radius in (0., .001, .025, .05, .18, .8):
+        actual = np.asarray(field.modB_derivatives(radius, theta, zeta))
+        expected = [2+.1*radius**2+.2*radius*(1+.3*radius**2)*np.cos(theta)
+                    + radius**5*np.exp(.2*radius**2)*harmonic,
+                    .2*radius+.2*(1+.9*radius**2)*np.cos(theta)
+                    + (5*radius**4+.4*radius**6)*np.exp(.2*radius**2)*harmonic,
+                    -.2*(1+.3*radius**2)*np.sin(theta)
+                    + radius**4*np.exp(.2*radius**2)*angular,
+                    radius**5*np.exp(.2*radius**2)*toroidal]
+        np.testing.assert_allclose(actual, expected, rtol=0, atol=2e-12)
+        np.testing.assert_allclose(evaluate(radius, theta, zeta), expected, rtol=0, atol=2e-12)
+        assert float(radial_grad(radius)) == pytest.approx(expected[1], abs=2e-12)
+
+
+@pytest.mark.parametrize("sine", [False, True])
+def test_axis_regular_option_rejects_unresolved_inner_mode(sine):
+    s = np.linspace(.0025, .9975, 200)
+    r = np.sqrt(s)
+    mode = r**5
+    mode[0] *= 2
+    bmnc = np.stack([2+0*s, 0*s if sine else mode])
+    bmns = np.stack([0*s, mode]) if sine else None
+    with pytest.raises(ValueError, match="inner profile is unresolved for m=5"):
+        BoozerField.from_booz(s, bmnc, [0, 5], [0, 0],
+                              .4+0*s, 1+0*s, 0*s, 1., 1, axis_regular=True,
+                              axis_regular_tolerance=.01, bmns=bmns)
+
+
+@pytest.mark.parametrize("tolerance", [None, 0., -1., np.inf, np.nan])
+def test_axis_regular_option_requires_explicit_finite_positive_tolerance(tolerance):
+    s = np.linspace(.0025, .9975, 200)
+    with pytest.raises(ValueError, match="positive finite axis_regular_tolerance"):
+        BoozerField.from_booz(s, np.stack([8+0*s, s**2.5]), [0, 5], [0, 0],
+                              .4+0*s, 1+0*s, 0*s, 1., 1, axis_regular=True,
+                              axis_regular_tolerance=tolerance)
+
+
+def test_axis_regular_option_removes_manufactured_forbidden_linear_harmonic():
+    # Smooth Cartesian polynomial Re((x+iy)**5), exact noiseless half-mesh data.
+    s = np.linspace(.0025, .9975, 200)
+    args = (s, np.stack([8+0*s, s**2.5]), [0, 5], [0, 0],
+            .4+0*s, 1+0*s, 0*s, 1., 1)
+    default = BoozerField.from_booz(*args, angle_grid=None)
+    regular = BoozerField.from_booz(*args, axis_regular=True, axis_regular_tolerance=1e-12)
+    assert abs(float(default.modB_derivatives(0., 0., 0.)[1])) > 5e-5
+    for radius in (0., 1e-8, 1e-6, .001, .025, .9):
+        np.testing.assert_allclose(regular.modB_derivatives(radius, 0., 0.),
+                                   [8+radius**5, 5*radius**4, 0., 0.], atol=1e-14)
+
+
+def test_axis_regular_manufactured_inner_extrapolation_converges_under_refinement():
+    # Divide out known powers only in the oracle error to expose the small m=5
+    # harmonic rather than hiding it in an O(1) axisymmetric field tolerance.
+    radii = np.array([.02, .06, .15])
+    errors = []
+    for count in (16, 32, 64, 128):
+        s = (np.arange(count)+.5)/count
+        field = BoozerField.from_booz(s, (s**2.5*np.exp(3*s))[None, :], [5], [0],
+                                      .4+0*s, 1+0*s, 0*s, 1., 1, axis_regular=True,
+                                      axis_regular_tolerance=.01)
+        actual = np.asarray(jax.vmap(lambda r: jnp.array(field.modB_derivatives(r, 0., 0.)))(radii))
+        expected = np.exp(3*radii**2)
+        errors.append(max(np.max(np.abs(actual[:, 0]/radii**5-expected)),
+                          np.max(np.abs(actual[:, 1]/radii**4-(5+6*radii**2)*expected))))
+    assert np.all(np.array(errors[:-1])/errors[1:] > 8), errors
+
+
 def test_angle_table_matches_the_mode_sum():
     s = np.linspace(0.005, 0.995, 50)
     bmnc = np.stack([np.full_like(s, B0), -B0 * EPS * np.sqrt(s), 0.1 * s, 0.05 * s])