From 17750c8873a298ec394e2bd6bdaa6a34de7e9d52 Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Mon, 5 Oct 2026 15:37:31 -0500 Subject: [PATCH 1/7] Keep exterior terminal step out of Boozer tracer diagnostics The step that crosses the LCFS ends outside the |B| and profile tables, where the spline extrapolation can return a meaningless |B|. Its end point was kept as the held state and its energy error entered energy_error. Saved states now hold the last interior state; the exterior end point and its energy error are reported separately in terminal_states and terminal_energy_error (NaN unless lost). Loss and failure classification is unchanged. Fixes #115 --- essos/boozer.py | 52 +++++++++++++++++++++++++++++++------------- tests/test_boozer.py | 46 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 15 deletions(-) diff --git a/essos/boozer.py b/essos/boozer.py index b1802f91..412c92cd 100644 --- a/essos/boozer.py +++ b/essos/boozer.py @@ -193,10 +193,17 @@ class BoozerTrace: """Result of :func:`trace_boozer` (NumPy arrays). ``states`` is ``(particles, n_save, 5)``: ``s, theta, zeta, v_par`` and the - speed ``v``, held at the loss or thermalisation point afterwards. + speed ``v``. Every saved state lies inside the plasma: after a loss it is + held at the last interior state, and after thermalisation at that point. ``loss_times`` is ``-1`` for particles that were not lost; ``thermalized_times`` and ``failed_times`` likewise. A failed orbit is frozen at its last finite state and is not counted as confined. + + ``energy_error`` covers interior steps only. The step that crosses the + LCFS ends outside the field tables, where ``|B|`` is an extrapolation, so + its end point and energy error are kept apart in ``terminal_states`` + (``(particles, 5)``, as ``states``) and ``terminal_energy_error``; both are + NaN for particles that were not lost and neither is a physical diagnostic. """ times: np.ndarray @@ -205,6 +212,8 @@ class BoozerTrace: thermalized_times: np.ndarray energy_error: np.ndarray # max |E/E0 - 1| per particle (collisionless drift only) failed_times: np.ndarray | None = None + terminal_states: np.ndarray | None = None + terminal_energy_error: np.ndarray | None = None @property def lost(self): @@ -246,7 +255,8 @@ def one(y, mu): s, r, theta, *_ = _chart(y) e0 = 0.5 * y[3] ** 2 + mu * field.modB_derivatives(r, theta, y[2])[0] first = jnp.array([s, theta, y[2], y[3], jnp.sqrt(2 * e0)]) - return (y, mu, 0.0, jnp.asarray(True), -1.0, -1.0, -1.0, 0.0, e0), first + nan = jnp.full(5, jnp.nan) + return (y, mu, 0.0, jnp.asarray(True), -1.0, -1.0, -1.0, 0.0, e0, nan, jnp.nan), first return jax.vmap(one)(y0, mu0) @@ -255,13 +265,14 @@ def _advance(field, dt, n_sub, mass, charge, species, thermal_cutoff, carry, keys, first_interval, count, method="rk4"): def one(carry, key): def step(carry, k): - y, mu, t, alive, t_loss, t_therm, t_fail, err, e0 = carry + y, mu, t, alive, t_loss, t_therm, t_fail, err, e0, terminal, terminal_err = carry rhs = lambda state: guiding_center_rhs(field, state, mu, mass, charge) y1 = _rk_step(rhs, y, dt, method) s1, r1, th1, *_ = _chart(y1) B1 = field.modB_derivatives(r1, th1, y1[2])[0] e_orbit = 0.5 * y1[3] ** 2 + mu * B1 - err = jnp.where(alive, jnp.maximum(err, jnp.abs(e_orbit / e0 - 1)), err) + step_err = jnp.abs(e_orbit / e0 - 1) + exterior = jnp.array([s1, th1, y1[2], y1[3], jnp.sqrt(2 * e_orbit)]) thermal = jnp.asarray(False) if species is not None: v = jnp.sqrt(2 * e_orbit) @@ -281,15 +292,19 @@ def step(carry, k): lost = alive & finite & (s1 >= 1.0) failed = alive & ~lost & ~finite t_fail = jnp.where(failed, t + dt, t_fail) + err = jnp.where(alive & finite & ~lost, jnp.maximum(err, step_err), err) err = jnp.where(alive & ~finite, jnp.inf, err) + terminal = jnp.where(lost, exterior, terminal) + terminal_err = jnp.where(lost, step_err, terminal_err) therm = alive & ~lost & finite & thermal t_loss = jnp.where(lost, t + dt, t_loss) t_therm = jnp.where(therm, t + dt, t_therm) keep = alive & ~therm & finite - y = jnp.where(keep | lost, y1, y) - mu = jnp.where(keep, mu1, mu) - e0 = jnp.where(keep, e0_new, e0) - return (y, mu, t + dt, keep & ~lost, t_loss, t_therm, t_fail, err, e0), None + y = jnp.where(keep & ~lost, y1, y) + mu = jnp.where(keep & ~lost, mu1, mu) + e0 = jnp.where(keep & ~lost, e0_new, e0) + return (y, mu, t + dt, keep & ~lost, t_loss, t_therm, t_fail, err, e0, + terminal, terminal_err), None def interval(carry, i): carry = jax.lax.fori_loop(0, n_sub, lambda k, state: step(state, i * n_sub + k)[0], carry) @@ -389,13 +404,11 @@ def trace_boozer(field, s, theta, zeta, pitch, *, speed, mass, charge, tmax, tim if progress is not None: progress(stop, n_int) first_interval = stop - status = [np.asarray(carry[i])[:n].copy() for i in (4, 5, 6, 7)] + status = [np.asarray(carry[i])[:n].copy() for i in _STATUS] if selected.size: - for result, i in zip(status, (4, 5, 6, 7)): + for result, i in zip(status, _STATUS): result[selected] = np.asarray(running[i])[:selected.size] - t_loss, t_therm, t_fail, err = status - return BoozerTrace(np.linspace(0.0, float(tmax), n_int + 1), states, - t_loss, t_therm, err, t_fail) + return _result(tmax, n_int, states, *status) carry = initial saved = [] for first_interval in range(0, n_int, chunk): @@ -407,5 +420,14 @@ def trace_boozer(field, s, theta, zeta, pitch, *, speed, mass, charge, tmax, tim jax.block_until_ready(part) progress(first_interval + count, n_int) states = np.concatenate([np.asarray(first)[:, None]] + [np.asarray(p) for p in saved], axis=1)[:n] - t_loss, t_therm, t_fail, err = (np.asarray(x)[:n] for x in (carry[4], carry[5], carry[6], carry[7])) - return BoozerTrace(np.linspace(0.0, float(tmax), n_int + 1), states, t_loss, t_therm, err, t_fail) + return _result(tmax, n_int, states, *(np.asarray(carry[i])[:n] for i in _STATUS)) + + +# Carry slots: loss, thermalisation and failure times, interior energy error, +# terminal (exterior) state and its energy error. +_STATUS = (4, 5, 6, 7, 9, 10) + + +def _result(tmax, n_int, states, t_loss, t_therm, t_fail, err, terminal, terminal_err): + return BoozerTrace(np.linspace(0.0, float(tmax), n_int + 1), states, t_loss, t_therm, err, + t_fail, terminal, terminal_err) diff --git a/tests/test_boozer.py b/tests/test_boozer.py index 7ea050f5..e1020c97 100644 --- a/tests/test_boozer.py +++ b/tests/test_boozer.py @@ -406,3 +406,49 @@ def test_substep_counts_reuse_compilation_and_preserve_static_outputs(collisions np.testing.assert_array_equal(x, y) assert np.count_nonzero(np.asarray(b[0][4]) >= 0) == 2 assert static._cache_size() == 2 and dynamic._cache_size() == 1 + + +def _bundled_wout_field(): + """Boozer-like field from the bundled n3are WOUT's half-mesh |B| spectrum. + + VMEC angles are not Boozer angles, but the tables have the shape and + scale of a real device, which is all the loss bookkeeping needs. + """ + from pathlib import Path + from netCDF4 import Dataset + wout = Path(__file__).parents[1] / "examples" / "input_files" / "wout_n3are_R7.75B5.7.nc" + with Dataset(wout) as d: + v = {k: np.asarray(d.variables[k][:]) for k in + ("bmnc", "xm_nyq", "xn_nyq", "iotas", "bvco", "buco", "phi", "nfp")} + ns = v["iotas"].size + s = (np.arange(1, ns) - 0.5) / (ns - 1) + keep = (v["xm_nyq"] <= 6) & (np.abs(v["xn_nyq"]) <= 6 * v["nfp"]) + return BoozerField.from_booz(s, v["bmnc"][1:, keep].T, v["xm_nyq"][keep], v["xn_nyq"][keep], + v["iotas"][1:], v["bvco"][1:], v["buco"][1:], + psi0_from_vmec(v["phi"][-1]), int(v["nfp"])) + + +@pytest.mark.parametrize("compact", [True, False]) +def test_lost_orbits_keep_exterior_terminal_diagnostics_apart(compact): + """Issue #115: coarse steps that leave the tables must not pollute interior diagnostics.""" + field = _bundled_wout_field() + n = 16 + rng = np.random.default_rng(0) + out = trace_boozer(field, np.full(n, .5), rng.uniform(0, 2 * np.pi, n), rng.uniform(0, 2 * np.pi, n), + rng.uniform(-1, 1, n), speed=V0, mass=300 * M, charge=Q, tmax=1e-6, + timestep=1e-7, n_save=11, compact=compact) + lost = out.lost + assert lost.any() and lost.sum() + out.failed.sum() + (~lost & ~out.failed).sum() == n + assert not (lost & out.failed).any() + # Saved states stay inside the plasma and finite, including after the loss. + assert np.all(out.states[~out.failed, :, 0] < 1) + assert np.isfinite(out.states).all() + # The exterior extrapolation is reported separately, only for lost orbits. + assert out.terminal_states.shape == (n, 5) + assert np.all(out.terminal_states[lost, 0] >= 1) + assert np.isnan(out.terminal_states[~lost]).all() and np.isnan(out.terminal_energy_error[~lost]).all() + assert np.nanmax(out.terminal_energy_error[lost]) > 1 # the extrapolated step is garbage... + assert np.all(out.energy_error[lost] < 0.5) # ...and stays out of the interior error + # The held state is the start of the crossing step: its energy matches the birth energy. + v_end = out.states[lost, -1, 4] + np.testing.assert_allclose(v_end / V0, 1, atol=0.3) From 4be22bb743d9d83b2d608543d28d413c6d072d88 Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Thu, 15 Oct 2026 16:53:47 -0500 Subject: [PATCH 2/7] 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 1282ae6db413923f9c31af07f3c2f4eb97deffbf Mon Sep 17 00:00:00 2001 From: rogeriojorge Date: Tue, 20 Oct 2026 09:25:23 -0500 Subject: [PATCH 3/7] Fix native RHS preparation on JAX 0.10 --- essos/rk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/essos/rk.py b/essos/rk.py index 55e7bb37..eed576ba 100644 --- a/essos/rk.py +++ b/essos/rk.py @@ -12,7 +12,7 @@ import jax import jax.numpy as jnp from jax import lax, tree_util -from jax.interpreters import partial_eval +from jax._src.interpreters import partial_eval import numpy as np from .rk_dense import DOPRI8_OUTPUT From 8884287054e4069f479642d684898e35f56ec41b Mon Sep 17 00:00:00 2001 From: rogeriojorge Date: Tue, 20 Oct 2026 09:46:04 -0500 Subject: [PATCH 4/7] Test step-limit forwarding for both tracing backends --- tests/test_dynamics.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/test_dynamics.py b/tests/test_dynamics.py index 5ddd43d3..2411a703 100644 --- a/tests/test_dynamics.py +++ b/tests/test_dynamics.py @@ -577,19 +577,43 @@ def test_tracing_trace_collisions_adaptative(field, particles,electric_field): pytest.main() -def test_tracing_max_steps_is_configurable_and_bounded(): +@pytest.mark.parametrize("integrator", ["diffrax", "essos_rk45", "essos_dopri8"]) +def test_tracing_max_steps_is_configurable_and_bounded(monkeypatch, integrator): """The Diffrax ceiling used to be 1e10, so a trace that could not finish ran until the process was killed rather than returning.""" import inspect + from types import SimpleNamespace + import essos.dynamics as dynamics from essos.dynamics import Tracing default = inspect.signature(Tracing).parameters["max_steps"].default assert default == 1_000_000 - source = inspect.getsource(Tracing) - assert "max_steps=10000000000" not in source - assert source.count("max_steps=self.max_steps") == 1 # every solve goes through Tracing._solve + class BackendCalled(Exception): + pass + + forwarded = [] + + def backend(*args, **kwargs): + forwarded.append(kwargs["max_steps"]) + raise BackendCalled + + # Stop before either backend integrates; exercise each dispatch branch. + monkeypatch.setattr(dynamics, "diffeqsolve", backend) + monkeypatch.setattr(dynamics, "solve_rk", backend) + tracing = object.__new__(Tracing) + tracing.__dict__.update( + integrator=integrator, max_steps=3, maxtime=1., timestep=.1, + times=jnp.array([0., 1.]), rtol=1e-7, atol=1e-7, + native_error_control="standard", native_controller="i", + native_pcoeff=.4, native_icoeff=.3, native_dense_output="standard", + progress_meter=None) + event = SimpleNamespace(cond_fn=lambda *args: False, root_finder=None) + with pytest.raises(BackendCalled): + tracing._solve(terms=None, solver=None, controller=False, t0=0., + y0=jnp.zeros(1), args=None, event=event) + assert forwarded == [3] From 624091139496afd13850e354d19194e1bd646698 Mon Sep 17 00:00:00 2001 From: rogeriojorge Date: Fri, 23 Oct 2026 15:03:09 -0500 Subject: [PATCH 5/7] Simplify tracing step-limit regression --- tests/test_dynamics.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/test_dynamics.py b/tests/test_dynamics.py index 2411a703..e2d6f7c1 100644 --- a/tests/test_dynamics.py +++ b/tests/test_dynamics.py @@ -579,13 +579,11 @@ def test_tracing_trace_collisions_adaptative(field, particles,electric_field): @pytest.mark.parametrize("integrator", ["diffrax", "essos_rk45", "essos_dopri8"]) def test_tracing_max_steps_is_configurable_and_bounded(monkeypatch, integrator): - """The Diffrax ceiling used to be 1e10, so a trace that could not finish ran - until the process was killed rather than returning.""" + """Keep a bounded default and forward explicit limits to each backend.""" import inspect from types import SimpleNamespace import essos.dynamics as dynamics - from essos.dynamics import Tracing default = inspect.signature(Tracing).parameters["max_steps"].default assert default == 1_000_000 @@ -593,10 +591,8 @@ def test_tracing_max_steps_is_configurable_and_bounded(monkeypatch, integrator): class BackendCalled(Exception): pass - forwarded = [] - def backend(*args, **kwargs): - forwarded.append(kwargs["max_steps"]) + assert kwargs["max_steps"] == 3 raise BackendCalled # Stop before either backend integrates; exercise each dispatch branch. @@ -613,7 +609,6 @@ def backend(*args, **kwargs): with pytest.raises(BackendCalled): tracing._solve(terms=None, solver=None, controller=False, t0=0., y0=jnp.zeros(1), args=None, event=event) - assert forwarded == [3] From 2a62bb11b9b4603af283457a47bfbca38fcecfb4 Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Fri, 23 Oct 2026 15:05:35 -0500 Subject: [PATCH 6/7] 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 7/7] 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),