diff --git a/essos/boozer.py b/essos/boozer.py index e226912d..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, @@ -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): diff --git a/tests/test_boozer.py b/tests/test_boozer.py index 734a3fb2..568146f3 100644 --- a/tests/test_boozer.py +++ b/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