diff --git a/README.md b/README.md index 7a59d253..c78988e7 100644 --- a/README.md +++ b/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 diff --git a/essos/boozer.py b/essos/boozer.py index e226912d..1eb3ee41 100644 --- a/essos/boozer.py +++ b/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) diff --git a/tests/test_boozer.py b/tests/test_boozer.py index 734a3fb2..975adf7d 100644 --- a/tests/test_boozer.py +++ b/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 = bmnc.T@np.cos(sample_phase)+bmns.T@np.sin(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])