Keep Boozer splines continuous and avoid duplicate m=0 storage

essos/boozer.py

@@ -62,7 +62,7 @@ class BoozerField(eqx.Module):
     psi0: float
     nfp: int = eqx.field(static=True)
     sine_coef: jax.Array | None = None
-    table: jax.Array | None = None  # (nint, ntheta + DEGREE, nzeta + DEGREE, 2, 4)
+    table: jax.Array | None = None  # theta row 0: m=0; remaining rows: m>0
 
     @classmethod
     def from_booz(cls, s, bmnc, xm, xn, iota, G, I, psi0, nfp, mode_tolerance=1e-6, *, bmns=None,
@@ -78,7 +78,7 @@ def from_booz(cls, s, bmnc, xm, xn, iota, G, I, psi0, nfp, mode_tolerance=1e-6,
         With ``angle_grid`` (default 6 points per shortest wavelength, or a
         ``(ntheta, nzeta)`` pair) the spectrum is summed once on a uniform grid in
         ``theta`` and one field period of ``zeta``, and evaluated by periodic
-        quintic B-splines in the angles, keeping the radial splines exact. A step
+        quintic B-splines in the angles with continuous radial splines. A step
         then costs the same for any number of modes; a scalar ``angle_grid`` is
         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
@@ -177,8 +177,9 @@ def _angle_table(b_coef, sine_coef, xm, xn, nfp, nt, nz):
     for axis in (1, 2):
         table = spline_filter1d(table, DEGREE, axis=axis, mode="grid-wrap")
     pad = ((DEGREE - 1) // 2, (DEGREE + 1) // 2)
-    # float32 halves the per-device copy; its rounding is far below the interpolation error
-    return jnp.asarray(np.pad(table, ((0, 0), pad, pad, (0, 0), (0, 0)), mode="wrap"), jnp.float32)
+    # Keep radial pieces continuous; float32 rounding introduces knot jumps.
+    table = np.pad(table, ((0, 0), pad, pad, (0, 0), (0, 0)), mode="wrap")
+    return jnp.asarray(np.concatenate([table[:, :1, :, 0], table[:, :, :, 1]], axis=1))
 
 
 def _bspline_matrix():
@@ -202,20 +203,24 @@ def _bspline(t):
 def _table_derivatives(knots, table, nfp, r, theta, zeta):
     k = jnp.clip(jnp.searchsorted(knots, r, side="right", method="compare_all") - 1, 0, knots.size - 2)
     d = r - knots[k]
-    n = jnp.array(table.shape[1:3]) - DEGREE
+    n = jnp.array(table.shape[1:3]) - jnp.array([DEGREE + 1, DEGREE])
     h = jnp.array([2 * np.pi, 2 * np.pi / nfp]) / n
     x = jnp.mod(jnp.stack([theta, zeta]), n * h) / h
     i = jnp.minimum(jnp.floor(x), n - 1).astype(k.dtype)
     (wt, dt), (wz, dz) = _bspline(x[0] - i[0]), _bspline(x[1] - i[1])
-    block = jax.lax.dynamic_slice(table, (k, i[0], i[1], 0 * k, 0 * k),
-                                  (1, DEGREE + 1, DEGREE + 1, 2, 4))[0].astype(d.dtype)
+    m0 = jax.lax.dynamic_slice(table, (k, 0*k, i[1], 0*k),
+                               (1, 1, DEGREE+1, 4))[0, 0].astype(d.dtype)
+    mp = jax.lax.dynamic_slice(table, (k, i[0]+1, i[1], 0*k),
+                               (1, DEGREE+1, DEGREE+1, 4))[0].astype(d.dtype)
     one = jnp.ones_like(d)
-    value = block @ jnp.stack([d**3, d**2, d, one])  # radial sums: (theta, zeta, channel)
-    slope = block @ jnp.stack([3 * d**2, 2 * d, one, 0 * d])
-    z_value, z_slope = jnp.einsum("bcq,c->bq", value, wz), jnp.einsum("bcq,c->bq", slope, wz)
-    v, vr, vt = wt @ z_value, wt @ z_slope, dt @ z_value / h[0]
-    vz = wt @ jnp.einsum("bcq,c->bq", value, dz) / h[1]
-    return v[0] + r * v[1], vr[0] + v[1] + r * vr[1], vt[1], vz[0] + r * vz[1]
+    powers, deriv = jnp.stack([d**3, d**2, d, one]), jnp.stack([3*d**2, 2*d, one, 0*d])
+    a, b = m0 @ powers, mp @ powers
+    v0, v1 = wz @ a, wt @ (b @ wz)
+    vr0, vr1 = wz @ (m0 @ deriv), wt @ ((mp @ deriv) @ wz)
+    vt = dt @ (b @ wz) / h[0]
+    vz0, vz1 = dz @ a / h[1], wt @ (b @ dz) / h[1]
+    return v0 + r*v1, vr0 + v1 + r*vr1, vt, vz0 + r*vz1
+
 
 
 def _chart(y):

tests/test_boozer.py

@@ -44,13 +44,73 @@ def test_angle_table_matches_the_mode_sum():
     table = BoozerField.from_booz(*args, bmns=0.5 * bmnc, angle_grid=(16, 16))
     assert BoozerField.from_booz(*args).table is None  # four modes: summed directly
     exact = BoozerField.from_booz(*args, bmns=0.5 * bmnc, angle_grid=None)
-    assert table.table.shape[1:3] == (21, 21)
+    assert table.table.shape[1:3] == (22, 21)
     points = np.random.default_rng(0).uniform(-7, 7, (3, 20))
     points[0] = np.abs(points[0]) / 7
     np.testing.assert_allclose(jax.vmap(table.modB_derivatives)(*points),
                                jax.vmap(exact.modB_derivatives)(*points), atol=1e-4 * B0)
 
 
+def test_angle_table_preserves_radial_polynomial_across_knots():
+    # Constant angles isolate the cubic radial spline.
+    r = np.array([.07, .135, .24, .37, .49, .61, .74, .88, .96])
+    s = r**2
+    polynomial = lambda x: np.pi + .23*x - .31*x**2 + .071*x**3
+    field = BoozerField.from_booz(
+        s, polynomial(r)[None, :], [0], [0], np.full_like(s, .43),
+        np.ones_like(s), np.zeros_like(s), 1., 2, angle_grid=(16, 16))
+    points = np.stack([r[1:-1]-1e-11, r[1:-1]+1e-11], axis=1)
+
+    def evaluate(x):
+        value, slope = field.modB_derivatives(x, .381, .197)[:2]
+        curvature = jax.grad(lambda t: field.modB_derivatives(t, .381, .197)[1])(x)
+        return jnp.array([value, slope, curvature])
+
+    actual = np.asarray(jax.jit(jax.vmap(jax.vmap(evaluate)))(jnp.asarray(points)))
+    expected = np.stack([polynomial(points), .23-.62*points+.213*points**2,
+                         -.62+.426*points], axis=-1)
+    for derivative, tolerance in enumerate([2e-12, 2e-11, 2e-10]):
+        np.testing.assert_allclose(actual[..., derivative], expected[..., derivative],
+                                   rtol=0, atol=tolerance)
+        # Remove the known smooth variation to measure each knot jump.
+        jump = np.diff(actual[..., derivative]-expected[..., derivative], axis=1)
+        np.testing.assert_allclose(jump, 0., rtol=0, atol=tolerance)
+
+
+def test_angle_table_orbit_conserves_analytic_energy_across_radial_knots():
+    from diffrax import Dopri8, ODETerm, PIDController, SaveAt, diffeqsolve
+
+    r = np.linspace(.08, .95, 49)
+    s, nfp = r**2, 2
+    bmnc = np.stack([1.7+.11*r**2, .06*r*(1+.13*r**2), .017*r, .008*r**2])
+    bmns = np.stack([np.zeros_like(r), .011*r, np.zeros_like(r), .003*r**2])
+    field = BoozerField.from_booz(
+        s, bmnc, [0, 1, 1, 2], [0, 0, nfp, nfp], np.full_like(s, .43),
+        np.ones_like(s), np.zeros_like(s), 1., nfp, bmns=bmns, angle_grid=(64, 64))
+    mu = .23
+    initial = jnp.array([.55*np.cos(.7), .55*np.sin(.7), .13, .9])
+    sol = diffeqsolve(
+        ODETerm(lambda t, y, args: guiding_center_rhs(field, y, mu, 1., 1.)),
+        Dopri8(), t0=0., t1=30., dt0=.01, y0=initial,
+        saveat=SaveAt(ts=jnp.linspace(0., 30., 1001)),
+        stepsize_controller=PIDController(rtol=1e-11, atol=1e-13), max_steps=10000)
+    y = np.asarray(sol.ys)
+    radius = np.hypot(y[:, 0], y[:, 1])
+    theta, zeta = np.arctan2(y[:, 1], y[:, 0]), y[:, 2]
+    assert np.count_nonzero(np.diff(np.searchsorted(r, radius))) >= 10
+    # Closed Fourier polynomial, independent of the table's value routine.
+    B = (1.7+.11*radius**2 + .06*radius*(1+.13*radius**2)*np.cos(theta)
+         + .017*radius*np.cos(theta-nfp*zeta) + .008*radius**2*np.cos(2*theta-nfp*zeta)
+         + .011*radius*np.sin(theta) + .003*radius**2*np.sin(2*theta-nfp*zeta))
+    analytic_energy = .5*y[:, 3]**2+mu*B
+    table_B = np.asarray(jax.vmap(field.modB)(jnp.asarray(radius**2),
+                                           jnp.asarray(theta), jnp.asarray(zeta)))
+    table_energy = .5*y[:, 3]**2+mu*table_B
+    # Check the independent analytic energy and continuity of the table energy.
+    errors = [np.max(np.abs(h/h[0]-1)) for h in (analytic_energy, table_energy)]
+    assert max(errors) < 1e-9, errors
+
+
 def test_vmec_flux_sign_sets_the_analytic_radial_drift():
     """For B=B0(1-eps*r*cos theta), positive ions at theta=pi/2 drift outward."""
     import equinox as eqx