From 4be22bb743d9d83b2d608543d28d413c6d072d88 Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Thu, 15 Oct 2026 16:53:47 -0500 Subject: [PATCH 1/3] Preserve radial continuity in Boozer field tables --- essos/boozer.py | 8 +++--- tests/test_boozer.py | 63 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/essos/boozer.py b/essos/boozer.py index e226912d..4b88439e 100644 --- a/essos/boozer.py +++ b/essos/boozer.py @@ -78,7 +78,8 @@ 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, retaining the radial spline precision + so adjacent polynomial pieces remain continuous to roundoff. 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 +178,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) + # Rounding each radial polynomial independently can break continuity at + # shared knots, imposing an orbit-energy error floor even with tight steps. + return jnp.asarray(np.pad(table, ((0, 0), pad, pad, (0, 0), (0, 0)), mode="wrap")) def _bspline_matrix(): diff --git a/tests/test_boozer.py b/tests/test_boozer.py index d54dcd19..ad3b8ae0 100644 --- a/tests/test_boozer.py +++ b/tests/test_boozer.py @@ -49,6 +49,69 @@ def test_angle_table_matches_the_mode_sum(): jax.vmap(exact.modB_derivatives)(*points), atol=1e-4 * B0) +def test_angle_table_preserves_radial_polynomial_across_knots(): + # A global cubic is reproduced exactly by the radial spline. Constant + # angles remove angular interpolation error from this independent oracle. + 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) + # Subtract the known smooth variation before measuring the C0/C1/C2 + # jump. Independently rounded pieces need not meet at a shared knot. + 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 + # Both Hamiltonians matter: the analytic oracle checks the orbit, while + # the interpolated Hamiltonian also detects discontinuous field values. + 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 From 2a62bb11b9b4603af283457a47bfbca38fcecfb4 Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Fri, 23 Oct 2026 15:05:35 -0500 Subject: [PATCH 2/3] Condense angle-table precision comments --- essos/boozer.py | 6 ++---- tests/test_boozer.py | 9 +++------ 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/essos/boozer.py b/essos/boozer.py index 4b88439e..efb3695f 100644 --- a/essos/boozer.py +++ b/essos/boozer.py @@ -78,8 +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, retaining the radial spline precision - so adjacent polynomial pieces remain continuous to roundoff. 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 @@ -178,8 +177,7 @@ 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) - # Rounding each radial polynomial independently can break continuity at - # shared knots, imposing an orbit-energy error floor even with tight steps. + # Keep radial pieces continuous; float32 rounding introduces knot jumps. return jnp.asarray(np.pad(table, ((0, 0), pad, pad, (0, 0), (0, 0)), mode="wrap")) diff --git a/tests/test_boozer.py b/tests/test_boozer.py index ddcf4b7a..183ed962 100644 --- a/tests/test_boozer.py +++ b/tests/test_boozer.py @@ -52,8 +52,7 @@ def test_angle_table_matches_the_mode_sum(): def test_angle_table_preserves_radial_polynomial_across_knots(): - # A global cubic is reproduced exactly by the radial spline. Constant - # angles remove angular interpolation error from this independent oracle. + # 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 @@ -73,8 +72,7 @@ def evaluate(x): for derivative, tolerance in enumerate([2e-12, 2e-11, 2e-10]): np.testing.assert_allclose(actual[..., derivative], expected[..., derivative], rtol=0, atol=tolerance) - # Subtract the known smooth variation before measuring the C0/C1/C2 - # jump. Independently rounded pieces need not meet at a shared knot. + # 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) @@ -108,8 +106,7 @@ def test_angle_table_orbit_conserves_analytic_energy_across_radial_knots(): 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 - # Both Hamiltonians matter: the analytic oracle checks the orbit, while - # the interpolated Hamiltonian also detects discontinuous field values. + # 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 From 3a7a471b8f471d3863bf92e05193acba536e1e99 Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Sat, 24 Oct 2026 06:23:42 -0500 Subject: [PATCH 3/3] Store theta-independent Boozer coefficients once --- essos/boozer.py | 27 ++++++++++++++++----------- tests/test_boozer.py | 2 +- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/essos/boozer.py b/essos/boozer.py index efb3695f..d72e9173 100644 --- a/essos/boozer.py +++ b/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, @@ -178,7 +178,8 @@ def _angle_table(b_coef, sine_coef, xm, xn, nfp, nt, nz): table = spline_filter1d(table, DEGREE, axis=axis, mode="grid-wrap") pad = ((DEGREE - 1) // 2, (DEGREE + 1) // 2) # Keep radial pieces continuous; float32 rounding introduces knot jumps. - return jnp.asarray(np.pad(table, ((0, 0), pad, pad, (0, 0), (0, 0)), mode="wrap")) + 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): diff --git a/tests/test_boozer.py b/tests/test_boozer.py index 183ed962..568146f3 100644 --- a/tests/test_boozer.py +++ b/tests/test_boozer.py @@ -44,7 +44,7 @@ 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),