diff --git a/src/vmecpp/_qs.py b/src/vmecpp/_qs.py new file mode 100644 index 000000000..574dba2fd --- /dev/null +++ b/src/vmecpp/_qs.py @@ -0,0 +1,180 @@ +"""Quasisymmetry objectives built only from the public geometry API. + +Re-exported from vmecpp.simsopt_compat, next to the SIMSOPT wrapper this +mirrors; not part of the solver. + +The metric is the one SIMSOPT's ``QuasisymmetryRatioResidual`` implements: for +quasisymmetry the ratio ``(B x grad B . grad psi) / (B . grad B)`` is constant +on a flux surface, so + + f = sum_j w_j < [ (1/B^3) ( (N - iota M) B x grad B . grad psi + - (M G + N I) B . grad B ) ]^2 > + +with ``< . >`` the flux-surface average, ``G`` and ``I`` the poloidal and +toroidal current profiles, and ``(M, N)`` the desired helicity. Discretized on +a uniform ``(theta, phi)`` grid over one field period this is a sum of squares, + + R = sqrt( w_j nfp dtheta dphi / V' * sqrt(g) ) / B^3 + * ( (N - iota M) B x grad B . grad psi - (M G + N I) B . grad B ), + +which is what ``quasisymmetry_residuals`` returns. + +Everything is computed from the product-basis geometry contract in JAX. The +flux-surface measure ``sqrt(g)`` matters: dropping it changes the objective, +and its derivative is one of the two errors that only show up end to end. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import jax +import jax.numpy as jnp + +from vmecpp.geometry import Geometry, evaluate + +# VMEC++ fixes the sign of the Jacobian to -1, and stores the field components +# with the opposite sign to the one the raw flux derivatives give here. The +# magnitudes agree; only the convention differs. G and I are averages of the +# same covariant components, so the convention has to be applied consistently +# or the two terms of the residual pick up different signs. +_JACOBIAN_MAGNITUDE = jnp.abs + + +def _frame(geometry: Geometry, coordinates: jax.Array): + """Covariant basis vectors, Jacobian and field components at one point.""" + r, z, lambda_, toroidal_flux, poloidal_flux = evaluate(geometry, coordinates) + zeta = coordinates[2] + cos_zeta = jnp.cos(zeta) + sin_zeta = jnp.sin(zeta) + + e_s = jnp.asarray([r[1] * cos_zeta, r[1] * sin_zeta, z[1]]) + e_theta = jnp.asarray([r[2] * cos_zeta, r[2] * sin_zeta, z[2]]) + e_zeta = jnp.asarray( + [ + r[3] * cos_zeta - r[0] * sin_zeta, + r[3] * sin_zeta + r[0] * cos_zeta, + z[3], + ] + ) + sqrt_g = jnp.dot(e_s, jnp.cross(e_theta, e_zeta)) + scale = 2.0 * jnp.pi * _JACOBIAN_MAGNITUDE(sqrt_g) + b_sup_theta = (poloidal_flux[1] - toroidal_flux[1] * lambda_[3]) / scale + b_sup_zeta = toroidal_flux[1] * (1.0 + lambda_[2]) / scale + magnetic_field = b_sup_theta * e_theta + b_sup_zeta * e_zeta + return { + "sqrt_g": sqrt_g, + "b_sup_theta": b_sup_theta, + "b_sup_zeta": b_sup_zeta, + "b_sub_theta": jnp.dot(magnetic_field, e_theta), + "b_sub_zeta": jnp.dot(magnetic_field, e_zeta), + "mod_b": jnp.linalg.norm(magnetic_field), + "iota": poloidal_flux[1] / toroidal_flux[1], + "toroidal_flux_derivative": toroidal_flux[1], + } + + +def magnetic_field_strength(geometry: Geometry, coordinates: jax.Array) -> jax.Array: + """Return ``|B|`` reconstructed from ``R``, ``Z``, ``lambda``, and fluxes.""" + return _frame(geometry, coordinates)["mod_b"] + + +def _surface_fields(geometry: Geometry, s, theta, zeta): + """Evaluate the point quantities plus the |B| angular derivatives.""" + + def mod_b(angles): + return magnetic_field_strength(geometry, jnp.asarray([s, angles[0], angles[1]])) + + flat_theta = theta.ravel() + flat_zeta = zeta.ravel() + coordinates = jnp.stack( + (jnp.full(flat_theta.shape, s), flat_theta, flat_zeta), axis=1 + ) + fields = jax.vmap(_frame, in_axes=(None, 0))(geometry, coordinates) + gradients = jax.vmap(jax.grad(mod_b))(jnp.stack((flat_theta, flat_zeta), axis=1)) + fields["d_mod_b_d_theta"] = gradients[:, 0] + fields["d_mod_b_d_zeta"] = gradients[:, 1] + return fields + + +def quasisymmetry_residuals( + geometry: Geometry, + surfaces: Sequence[float] = (0.5,), + *, + helicity_m: int = 1, + helicity_n: int = 0, + weights: Sequence[float] | None = None, + ntheta: int = 63, + nphi: int = 64, +) -> jax.Array: + """Return the flat vector of quasisymmetry residuals ``R``. + + ``helicity_n = 0`` is quasi-axisymmetry. The residuals are normalized so + that the sum of their squares is the objective ``f``; this matches + SIMSOPT's ``QuasisymmetryRatioResidual.residuals()`` term by term. + """ + surface_values = jnp.asarray(surfaces, dtype=jnp.float64) + surface_weights = ( + jnp.ones_like(surface_values) + if weights is None + else jnp.asarray(weights, dtype=jnp.float64) + ) + nfp = geometry.nfp + + theta_1d = 2.0 * jnp.pi * jnp.arange(ntheta) / ntheta + phi_1d = 2.0 * jnp.pi * jnp.arange(nphi) / (nfp * nphi) + d_theta = theta_1d[1] - theta_1d[0] + d_phi = phi_1d[1] - phi_1d[0] + theta_grid, phi_grid = jnp.meshgrid(theta_1d, phi_1d, indexing="ij") + + # 2 pi psi is the toroidal flux, and VMEC++'s negative Jacobian carries the + # sign SIMSOPT writes explicitly as -phi_edge / (2 pi). + edge_flux = evaluate(geometry, jnp.asarray([1.0, 0.0, 0.0]))[3][0] + d_psi_d_s = -edge_flux / (2.0 * jnp.pi) + + toroidal_mode = helicity_n * nfp + + def one_surface(s, weight): + fields = _surface_fields(geometry, s, theta_grid, phi_grid) + sqrt_g = fields["sqrt_g"] + mod_b = fields["mod_b"] + d_theta_b = fields["d_mod_b_d_theta"] + d_zeta_b = fields["d_mod_b_d_zeta"] + + b_dot_grad_b = ( + fields["b_sup_theta"] * d_theta_b + fields["b_sup_zeta"] * d_zeta_b + ) + b_cross_grad_b_dot_grad_psi = ( + d_psi_d_s + * (fields["b_sub_theta"] * d_zeta_b - fields["b_sub_zeta"] * d_theta_b) + / sqrt_g + ) + # G and I are the (0, 0) harmonics of the covariant components, i.e. + # their plain angular averages on the uniform grid. + current_g = jnp.mean(fields["b_sub_zeta"]) + current_i = jnp.mean(fields["b_sub_theta"]) + iota = jnp.mean(fields["iota"]) + + volume_derivative = nfp * d_theta * d_phi * jnp.sum(sqrt_g) + measure = jnp.sqrt(weight * nfp * d_theta * d_phi / volume_derivative * sqrt_g) + return ( + measure + * ( + b_cross_grad_b_dot_grad_psi * (toroidal_mode - iota * helicity_m) + - b_dot_grad_b * (helicity_m * current_g + toroidal_mode * current_i) + ) + / mod_b**3 + ) + + return jnp.concatenate( + [ + one_surface(s, w) + for s, w in zip(surface_values, surface_weights, strict=True) + ] + ) + + +def quasisymmetry_total(geometry: Geometry, *args, **kwargs) -> jax.Array: + """Return the scalar quasisymmetry error ``f``, the sum of squared residuals.""" + residuals = quasisymmetry_residuals(geometry, *args, **kwargs) + return jnp.sum(residuals**2) diff --git a/src/vmecpp/simsopt_compat.py b/src/vmecpp/simsopt_compat.py index 9778c8313..f034dd4c5 100644 --- a/src/vmecpp/simsopt_compat.py +++ b/src/vmecpp/simsopt_compat.py @@ -21,6 +21,11 @@ ensure_vmecpp_input, is_vmec2000_input, ) +from vmecpp._qs import ( # noqa: F401 + magnetic_field_strength, + quasisymmetry_residuals, + quasisymmetry_total, +) logger = logging.getLogger(__name__) diff --git a/tests/test_qs.py b/tests/test_qs.py new file mode 100644 index 000000000..ecceaff59 --- /dev/null +++ b/tests/test_qs.py @@ -0,0 +1,140 @@ +import jax +import jax.numpy as jnp +import numpy as np + +import vmecpp +from vmecpp import geometry +from vmecpp import simsopt_compat as qs +from vmecpp.cpp import _vmecpp # type: ignore + +jax.config.update("jax_enable_x64", True) + + +def _run(indata): + return _vmecpp.run(indata._to_cpp_vmecindata(), verbose=_vmecpp.OutputMode.SILENT) + + +def _solovev(ns: int = 51) -> vmecpp.VmecInput: + source = vmecpp.VmecInput.from_file("examples/data/solovev.json") + return source.model_copy( + update={ + "ns_array": np.asarray([ns]), + "ftol_array": np.asarray([1.0e-14]), + "niter_array": np.asarray([20000]), + } + ) + + +def _rippled_solovev(ns: int = 51, ripple: float = 0.01) -> vmecpp.VmecInput: + """Solovev with one toroidal boundary mode, so it is genuinely 3D.""" + source = vmecpp.VmecInput.from_file("examples/data/solovev.json") + mpol = source.mpol + rbc = np.zeros((mpol, 3)) + zbs = np.zeros((mpol, 3)) + rbc[:, 1] = np.asarray(source.rbc)[:, 0] + zbs[:, 1] = np.asarray(source.zbs)[:, 0] + rbc[1, 2] = ripple + zbs[1, 2] = ripple + return source.model_copy( + update={ + "ntor": 1, + "rbc": rbc, + "zbs": zbs, + "raxis_c": np.asarray([4.0, 0.0]), + "zaxis_s": np.asarray([0.0, 0.0]), + "ns_array": np.asarray([ns]), + "ftol_array": np.asarray([1.0e-14]), + "niter_array": np.asarray([20000]), + } + ) + + +def test_reconstructed_field_strength_matches_the_vmec_spectrum() -> None: + """|B| rebuilt from the geometry jets must be VMEC's own |B|. + + A quasisymmetry residual near zero cannot establish this on its own: an + axisymmetric equilibrium gives zero non-quasi-axisymmetric power for any + zeta-independent function of the geometry, correct or not. VMEC's ``bmnc`` + spectrum is an independent oracle for the reconstruction itself. + + ``bmnc`` is a half-grid quantity while the evaluator interpolates on the + full grid, so the agreement is first order in the radial spacing: the + relative difference at one point falls 5.19e-3, 2.59e-3, 1.33e-3, 6.62e-4 + for ns = 31, 61, 121, 241. A wrong reconstruction would be off by tens of + percent. + """ + output = _run(_solovev(ns=51)) + wout = output.wout + jax_geometry = geometry.from_cpp(_vmecpp.make_geometry(output)) + + bmnc = np.asarray(wout.bmnc) + if bmnc.shape[1] == wout.ns: + bmnc = bmnc.T + xm = np.asarray(wout.xm_nyq) + xn = np.asarray(wout.xn_nyq) + + surface = int(0.6 * (wout.ns - 1)) + s = (surface + 0.5) / (wout.ns - 1) + for theta, zeta in ((0.3, 0.2), (1.1, -0.4), (2.4, 0.9)): + expected = float(np.sum(bmnc[surface] * np.cos(xm * theta - xn * zeta))) + actual = float( + qs.magnetic_field_strength(jax_geometry, jnp.asarray([s, theta, zeta])) + ) + assert abs(actual - expected) / abs(expected) < 4.0e-3 + + +def test_matches_simsopt_quasisymmetry_ratio_residual(tmp_path) -> None: + """The objective is SIMSOPT's, so SIMSOPT is the oracle. + + This is the whole point of the rewrite: a pure Python and JAX objective + that computes the same number as ``QuasisymmetryRatioResidual``, including + the ``sqrt(g)`` flux-surface measure and the multi-surface sum. + """ + simsopt_vmec = __import__("simsopt.mhd", fromlist=["Vmec"]) + diagnostics = __import__( + "simsopt.mhd.vmec_diagnostics", fromlist=["QuasisymmetryRatioResidual"] + ) + + indata = _rippled_solovev(ns=51) + output = _run(indata) + jax_geometry = geometry.from_cpp(_vmecpp.make_geometry(output)) + + wout_path = tmp_path / "wout_qs_reference.nc" + vmecpp.run(indata, verbose=False).wout.save(wout_path) + + surfaces = [0.3, 0.6, 0.9] + reference = diagnostics.QuasisymmetryRatioResidual( + simsopt_vmec.Vmec(str(wout_path)), + surfaces, + helicity_m=1, + helicity_n=0, + ntheta=63, + nphi=64, + ).compute() + + actual = np.asarray( + qs.quasisymmetry_residuals( + jax_geometry, surfaces, helicity_m=1, helicity_n=0, ntheta=63, nphi=64 + ) + ) + expected = reference.residuals1d + cosine = float( + expected @ actual / np.linalg.norm(expected) / np.linalg.norm(actual) + ) + assert cosine > 0.9999 + np.testing.assert_allclose(float(np.sum(actual**2)), reference.total, rtol=1.0e-3) + + +def test_qs_objective_is_differentiable_through_geometry() -> None: + output = _run(_rippled_solovev(ns=31)) + base = geometry.from_cpp(_vmecpp.make_geometry(output)) + + def objective(scale): + scaled = jax.tree_util.tree_map(lambda leaf: leaf * scale, base) + return qs.quasisymmetry_total( + scaled, [0.6], helicity_m=1, helicity_n=0, ntheta=16, nphi=16 + ) + + value, derivative = jax.value_and_grad(objective)(jnp.asarray(1.0)) + assert float(value) > 0.0 + assert np.isfinite(float(derivative))