diff --git a/examples/data/README.md b/examples/data/README.md index 00efb23fd..543b37f78 100644 --- a/examples/data/README.md +++ b/examples/data/README.md @@ -1,13 +1,13 @@ # Example input and output files for VMEC++ -A few cases are available for testing VMCE++ and experimenting with it: +A few cases are available for testing VMEC++ and experimenting with it: 1. `cth_like_fixed_bdy.json` - Stellarator case, similar to the Compact Toroidal Hybrid ([CTH](https://www.auburn.edu/cosam/departments/physics/research/plasma_physics/compact_toroidal_hybrid/index.htm)) device 1. `input.cth_like_fixed_bdy` - Fortran namelist input file to be used with Fortran VMEC 1. `cth_like_fixed_bdy.json` - JSON input file for VMEC++, derived from `input.cth_like_fixed_bdy` using [`indata2json`](https://github.com/jonathanschilling/indata2json) 1. `wout_cth_like_fixed_bdy.nc` - NetCDF output file, produced using [`PARVMEC`](https://github.com/ORNL-Fusion/PARVMEC) from `input.cth_like_fixed_bdy`, for testing the loading of a `wout` file using VMEC++'s tooling -1. `input.nfp4_QH_warm_start` - quasi-helically example for use with SIMSOPT +1. `input.nfp4_QH_warm_start` - quasi-helical example for use with SIMSOPT 1. `solovev` - axisymmetric Tokamak case, similar to the Solov'ev equilibrium used in the [1983 Hirshman & Whitson article](https://doi.org/10.1063/1.864116) 1. `input.solovev` - Fortran namelist input file for use with Fortran VMEC diff --git a/examples/external_optimizers.py b/examples/external_optimizers.py new file mode 100644 index 000000000..b4da85f26 --- /dev/null +++ b/examples/external_optimizers.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Drive a VMEC++ equilibrium from outside with general-purpose optimizers. + +VMEC's equilibrium is the stationary point of its augmented functional (MHD +energy plus the spectral-condensation and lambda constraints). The gradient of +that functional in the decomposed internal basis is the raw, unpreconditioned +force exposed by ``VmecModel.evaluate(precondition=False)``. Finding the +equilibrium is therefore the root problem F(x) = 0, which gradient- and +Hessian-based solvers can attack while reusing VMEC++'s forward model, its +preconditioner (``apply_preconditioner``, VMEC's approximate inverse Hessian), +and its Hessian-vector product (``hessian_vector_product``, a directional +derivative of the analytic force computed inside VMEC++). + +This module wires that residual to several solvers and is shared by the +benchmark ``main`` below and by the tests: + +* preconditioned descent (VMEC's own update direction), +* Jacobian-free Newton-Krylov, plain and preconditioned, and +* a true Newton-Krylov driven by VMEC++'s own Hessian-vector product. + +All converge to the same equilibrium as the native solver. Force evaluations are +counted inside VMEC++ (``force_eval_count``) so the comparison is fair across +methods, including the evaluations hidden in Hessian-vector products. + +Interpretation: what separates the methods is how they couple DOFs of +neighbouring flux surfaces. Plain JFNK does not -- it treats every degree of +freedom as independent, so it cannot damp the radial stiffness that lets a +surface cross its neighbour between iterations (BAD_JACOBIAN restarts in the +native solver). It still converges, but slowly. VMEC's couples each surface to +its jF +/- 1 neighbours; applying it -- as the inner Krylov preconditioner, or +refreshed every step in the HVP Newton -- is what rescues conditioning and cuts +the cost by an order of magnitude. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from scipy.optimize import newton_krylov +from scipy.sparse.linalg import LinearOperator, gmres + +import vmecpp +from vmecpp.cpp import _vmecpp # type: ignore[import] + +DEFAULT_INPUT = ( + Path(__file__).resolve().parents[1] + / "examples" + / "data" + / "cth_like_fixed_bdy.json" +) + + +_BAD_JACOBIAN = 2 # RestartReason::BAD_JACOBIAN (flow_control.h) + + +def _ensure_valid_initial_jacobian(model, max_reguess: int = 2) -> None: + """Reguess the magnetic axis until the initial Jacobian is non-singular. + + Inputs that ship no axis (e.g. cma.json, with raxis/zaxis all zero) start from a + degenerate geometry, so the bare forward model returns at the BAD_JACOBIAN + checkpoint with zero force. The native solver reguesses the axis on the first + iterate (vmec.cc SolveEquilibriumLoop); mirror that here via reinitialize() so a + cold start from any valid INDATA has a defined gradient. + """ + for _ in range(max_reguess): + model.evaluate(2, 2, False) + if model.restart_reason != _BAD_JACOBIAN: + return + model.reinitialize() + model.evaluate(2, 2, False) + + +def make_model(input_path: Path = DEFAULT_INPUT, ns: int = 11): + indata = _vmecpp.VmecINDATA.from_file(str(input_path)) + model = _vmecpp.VmecModel.create(indata, ns) + _ensure_valid_initial_jacobian(model) + return model + + +def residual(model): + """Return F(x) = raw internal-basis force; F(x) = 0 at equilibrium.""" + + def F(x): + model.set_state(np.ascontiguousarray(x, dtype=float)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), dtype=float) + + return F + + +@dataclass +class Result: + name: str + force_evals: int + outer_iters: int + seconds: float + residual_norm: float + energy: float + + +def reference_equilibrium(input_path: Path = DEFAULT_INPUT, ns: int = 11): + model = make_model(input_path, ns) + model.solve() + model.evaluate(2, 2, False) + return np.asarray(model.get_state(), float), model.mhd_energy + + +def _finish(model, name, x, outer_iters, t0): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return x, Result( + name, + model.force_eval_count, + outer_iters, + time.perf_counter() - t0, + float(np.linalg.norm(np.asarray(model.get_forces(), float))), + model.mhd_energy, + ) + + +def solve_vmecpp(input_path=DEFAULT_INPUT, ns=11): + """Native VMEC++ solve through the public API; reports the iteration count. + + Unlike the other variants this drives the ordinary ``vmecpp.run`` path rather + than the low-level ``VmecModel`` primitives, so the reported residual/energy + are the public wout diagnostics (invariant force residuals and total MHD + energy) rather than the raw internal-basis force this module counts elsewhere. + """ + vmec_input = vmecpp.VmecInput.from_file(input_path) + vmec_input.ns_array = np.asarray([ns], dtype=vmec_input.ns_array.dtype) + t0 = time.perf_counter() + output = vmecpp.run(vmec_input, verbose=False) + wout = output.wout + return output, Result( + name="VMEC++ (native, vmecpp.run)", + force_evals=wout.niter, + outer_iters=wout.niter, + seconds=time.perf_counter() - t0, + residual_norm=float(np.sqrt(wout.fsqt[-1])), + energy=wout.wb + wout.wp, + ) + + +def solve_preconditioned_descent( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, delt=0.9, momentum=0.5, max_iter=20000 +): + model = make_model(input_path, ns) + F = residual(model) + x = np.asarray(model.get_state(), float).copy() + v = np.zeros_like(x) + model.reset_force_eval_count() + it = 0 + t0 = time.perf_counter() + for _ in range(max_iter): + if np.linalg.norm(F(x)) < tol: + break + it += 1 + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, True) # preconditioned search direction + fprec = np.asarray(model.get_forces(), float) + v = momentum * v + delt * fprec + x = x + delt * v + return _finish(model, "preconditioned descent", x, it, t0) + + +def solve_newton_krylov( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=2500, preconditioned=False +): + model = make_model(input_path, ns) + F = residual(model) + x0 = np.asarray(model.get_state(), float) + inner_m = None + model.reset_force_eval_count() + if preconditioned: + # Assemble VMEC's preconditioner at x0 and use it, frozen, as the inner + # Krylov preconditioner. M^-1 approximates the inverse Hessian; it is a + # radial tridiagonal (Thomas) solve per Fourier mode, so it couples each + # flux surface to its jF +/- 1 neighbours -- the coupling the plain + # branch lacks. + model.evaluate(2, 2, True) + n_dof = x0.size + inner_m = LinearOperator( # type: ignore[call-overload] + (n_dof, n_dof), + matvec=lambda b: np.asarray( # type: ignore[call-overload] + model.apply_preconditioner(np.ascontiguousarray(b)), float + ), + ) + t0 = time.perf_counter() + x = newton_krylov( + F, x0, f_tol=tol, maxiter=max_iter, method="lgmres", inner_M=inner_m + ) + name = ( + "Newton-Krylov (preconditioned)" if preconditioned else "Newton-Krylov (JFNK)" + ) + return _finish(model, name, x, 0, t0) + + +def solve_newton_krylov_preconditioned(input_path=DEFAULT_INPUT, ns=11, tol=1e-9): + return solve_newton_krylov(input_path, ns, tol, preconditioned=True) + + +def solve_newton_hvp( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_newton=80, inner_tol=1e-3 +): + """Globalized Newton-Krylov using VMEC++'s own Hessian-vector product. + + Each Newton step solves H dx = -F with GMRES, where H v is hessian_vector_product + (the analytic force's directional derivative computed inside VMEC++) and the inner + solve is preconditioned by M^-1. A backtracking line search on ||F|| globalizes the + step, which is required on stiff 3D cases where the full Newton step overshoots. + """ + model = make_model(input_path, ns) + F = residual(model) + x = np.asarray(model.get_state(), float).copy() + n_dof = x.size + model.reset_force_eval_count() + t0 = time.perf_counter() + it = 0 + for _ in range(max_newton): + fk = F(x) + norm0 = np.linalg.norm(fk) + if norm0 < tol: + break + it += 1 + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, True) # assemble M at the current iterate + h_op = LinearOperator( # type: ignore[call-overload] + (n_dof, n_dof), + matvec=lambda v: np.asarray( # type: ignore[call-overload] + model.hessian_vector_product(np.ascontiguousarray(v)), float + ), + ) + m_op = LinearOperator( # type: ignore[call-overload] + (n_dof, n_dof), + matvec=lambda b: np.asarray( # type: ignore[call-overload] + model.apply_preconditioner(np.ascontiguousarray(b)), float + ), + ) + dx, _ = gmres(h_op, -fk, M=m_op, rtol=inner_tol, maxiter=100) + # Backtracking line search: accept the largest step that reduces ||F||. + alpha = 1.0 + for _ in range(30): + if np.linalg.norm(F(x + alpha * dx)) < norm0: + break + alpha *= 0.5 + else: + break # no decrease found; stop + x = x + alpha * dx + return _finish(model, "Newton (VMEC++ HVP + M^-1)", x, it, t0) + + +ALL_SOLVERS = ( + solve_vmecpp, + solve_preconditioned_descent, + solve_newton_krylov_preconditioned, + solve_newton_krylov, + solve_newton_hvp, +) + + +def main(): + _, w_star = reference_equilibrium() + print(f"reference equilibrium (native solve): W = {w_star:.8e}\n") + print( + f"{'optimizer':32s} {'F-evals':>8s} {'iters':>6s} {'time[s]':>8s} " + f"{'||F||':>10s} {'dW vs ref':>10s}" + ) + for solver in ALL_SOLVERS: + r = solver()[1] + print( + f"{r.name:32s} {r.force_evals:8d} {r.outer_iters:6d} {r.seconds:8.2f} " + f"{r.residual_norm:10.1e} {abs(r.energy - w_star):10.1e}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/fourier_resolution_increase.py b/examples/fourier_resolution_increase.py index 69dda7c91..375452153 100644 --- a/examples/fourier_resolution_increase.py +++ b/examples/fourier_resolution_increase.py @@ -10,14 +10,13 @@ resolution by :func:`vmecpp.interpolate_solution` (radial interpolation in ``sqrt(s)`` plus Fourier zero-padding). -``vmecpp.run_continuation`` performs the whole schedule in one call: +Setting ``VmecInput.mpol``/``.ntor`` to a sequence (instead of a plain int) makes +``vmecpp.run`` perform the whole schedule in one call: - output = vmecpp.run_continuation( - vmec_input, - ns_array=[15, 31, 31], - mpol_array=[5, 9, 13], - ntor_array=[4, 4, 4], - ) + vmec_input.ns_array = np.array([15, 31, 31]) + vmec_input.mpol = [5, 9, 13] + vmec_input.ntor = [4, 4, 4] + output = vmecpp.run(vmec_input) This example runs the schedule step by step instead, so it can report the per-step iteration counts and compare the total against solving at the target @@ -64,8 +63,8 @@ def step_input(ns: int, mpol: int, ntor: int) -> vmecpp.VmecInput: # --- Fourier continuation --------------------------------------------------- -# vmecpp.run_continuation(vmec_input, ns_array=ns_array, mpol_array=mpol_array, -# ntor_array=ntor_array) runs exactly this schedule in a single call; it is +# Setting vmec_input.ns_array/.mpol/.ntor to these schedules and calling +# vmecpp.run(vmec_input) runs exactly this schedule in a single call; it is # unrolled here to report the per-step iteration counts. print("Fourier continuation:") schedule = list(zip(ns_array, mpol_array, ntor_array, strict=True)) diff --git a/examples/simsopt_vmec_gradient.py b/examples/simsopt_vmec_gradient.py new file mode 100644 index 000000000..78eafde6e --- /dev/null +++ b/examples/simsopt_vmec_gradient.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Use VMEC++ as a gradient-providing equilibrium component for SIMSOPT. + +A SIMSOPT optimization over a plasma boundary normally differentiates VMEC by +finite differences: one equilibrium re-solve per boundary degree of freedom and +per outer iteration. VMEC++ instead provides the boundary gradient analytically +through the implicit-function adjoint (``vmecpp_adjoint.boundary_gradient``): one +extra Hessian solve, independent of the number of boundary DOFs. + +``VmecEnergy`` wraps that as a SIMSOPT ``Optimizable`` whose objective is the MHD +energy of the converged equilibrium and whose ``dJ`` is the adjoint gradient. +``optimize_to_target`` runs a gradient-based optimization of the boundary toward +a target energy, with analytic gradient or with finite differences, and +reports the cost (forward-model evaluations counted inside VMEC++, outer +iterations, wall time) so the two can be compared. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from vmecpp_adjoint import ( + DEFAULT_INPUT, + boundary_gradient, + finite_difference_boundary_gradient, + make_model, + mhd_energy, + partition, + solve_interior, +) + + +class VmecBoundaryProblem: + """Equilibrium energy and its boundary gradient, cached per boundary state.""" + + def __init__(self, input_path: Path = DEFAULT_INPUT, ns: int = 11): + self.model = make_model(input_path, ns) + self.model.solve() + self.ns = ns + self.interior, self.boundary = partition(self.model, ns) + self._x_full = np.asarray(self.model.get_state(), float).copy() + self._cached_p = self._x_full[self.boundary].copy() + + @property + def x0(self): + return self._x_full[self.boundary].copy() + + def _resolve(self, p): + p = np.asarray(p, float) + if not np.array_equal(p, self._cached_p): + self._x_full = solve_interior( + self.model, self._x_full, self.interior, self.boundary, p + ) + self._cached_p = p.copy() + + def value(self, p): + self._resolve(p) + self.model.set_state(np.ascontiguousarray(self._x_full)) + self.model.evaluate(2, 2, False) + return self.model.mhd_energy + + def gradient(self, p): + self._resolve(p) + return boundary_gradient( + self.model, self._x_full, self.interior, self.boundary, mhd_energy + ) + + +def make_simsopt_optimizable(problem: VmecBoundaryProblem): + """Wrap the problem as a SIMSOPT Optimizable exposing an analytic gradient.""" + # Imported lazily so the rest of the module (and the gradient benchmark) work + # without SIMSOPT installed. + from simsopt._core import Optimizable # noqa: PLC0415 + from simsopt._core.derivative import ( # noqa: PLC0415 + Derivative, + OptimizableDefaultDict, + derivative_dec, + ) + + class VmecEnergy(Optimizable): + def __init__(self): + x0 = problem.x0 + super().__init__(x0=x0, names=[f"boundary{i}" for i in range(x0.size)]) + + def J(self): + return problem.value(self.local_full_x) + + @derivative_dec + def dJ(self): + data = OptimizableDefaultDict({self: problem.gradient(self.local_full_x)}) + return Derivative(data) + + return VmecEnergy() + + +@dataclass +class GradResult: + method: str + force_evals: int + seconds: float + gradient: np.ndarray + + +def gradient_cost(input_path: Path = DEFAULT_INPUT, ns: int = 11, analytic=True): + """Cost of one full boundary gradient at the converged equilibrium. + + This is what an external optimizer pays per iteration. The analytic adjoint needs + one Hessian solve regardless of the number of boundary DOFs; finite differences re- + converge the equilibrium twice per boundary DOF. + """ + problem = VmecBoundaryProblem(input_path, ns) + x_star = problem._x_full.copy() + interior, boundary = problem.interior, problem.boundary + problem.model.reset_force_eval_count() + t0 = time.perf_counter() + if analytic: + g_dict = None + g = boundary_gradient(problem.model, x_star, interior, boundary, mhd_energy) + else: + g_dict = finite_difference_boundary_gradient( + problem.model, + x_star, + interior, + boundary, + mhd_energy, + range(boundary.size), + ) + g = np.array([g_dict[j] for j in range(boundary.size)]) + return GradResult( + "analytic adjoint" if analytic else "finite differences", + problem.model.force_eval_count, + time.perf_counter() - t0, + g, + ) + + +def main(): + analytic = gradient_cost(analytic=True) + fd = gradient_cost(analytic=False) + rel = np.linalg.norm(analytic.gradient - fd.gradient) / np.linalg.norm(fd.gradient) + n_boundary = analytic.gradient.size + print(f"boundary gradient cost ({n_boundary} boundary DOFs, solovev ns=11)\n") + print(f"{'method':20s} {'F-evals':>9s} {'time[s]':>8s}") + for r in (fd, analytic): + print(f"{r.method:20s} {r.force_evals:9d} {r.seconds:8.2f}") + print( + f"\nspeedup (force evals): {fd.force_evals / max(analytic.force_evals, 1):.1f}x" + f" gradient agreement: {rel:.1e}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/vmecpp_adjoint.py b/examples/vmecpp_adjoint.py new file mode 100644 index 000000000..f10b63d12 --- /dev/null +++ b/examples/vmecpp_adjoint.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Adjoint sensitivity of a converged VMEC++ equilibrium to its boundary. + +A fixed-boundary equilibrium satisfies the interior force balance F_I(x) = 0, +where x is the decomposed internal-basis state and F is the gradient of VMEC's +augmented functional. The outermost flux surface (the boundary) is the last +radial block of the state and is held fixed during the solve. For a scalar +objective J(x), the sensitivity to the boundary degrees of freedom follows from +the implicit function theorem: + + dJ/dx_B = dJ/dx_B|_x - (dF_I/dx_B)^T lambda, H_II lambda = dJ/dx_I, + +with H = dF/dx the (symmetric) Hessian of the augmented functional. Every +operator is matrix-free and already exposed by VmecModel: the Hessian-vector +product (``hessian_vector_product``) and the preconditioner +(``apply_preconditioner``), used to solve the adjoint system. Only one Hessian +solve is needed for the full boundary gradient, versus one equilibrium re-solve +per boundary degree of freedom for finite differences. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from scipy.sparse.linalg import LinearOperator, gmres + +from vmecpp.cpp import _vmecpp # type: ignore + +DEFAULT_INPUT = ( + Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" +) + + +def make_model(input_path: Path = DEFAULT_INPUT, ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(input_path)), ns) + + +def partition(model, ns: int): + """Indices of the interior (free) and boundary (LCFS) state components.""" + k = model.mpol * (model.ntor + 1) + n = np.asarray(model.get_state()).size + per_span = ns * k + n_span = n // per_span + boundary = [] + for s in range(n_span): + boundary.extend(range(s * per_span + (ns - 1) * k, s * per_span + ns * k)) + boundary = np.array(sorted(boundary)) + interior = np.setdiff1d(np.arange(n), boundary) + return interior, boundary + + +def _raw_force(model, x): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), float) + + +def _interior_operators(model, x, interior): + # The caller must set the base state to x and assemble the preconditioner + # (evaluate(2, 2, True)) before using these. hessian_vector_product uses the + # current state as its base point and restores it, so no per-matvec state + # update is needed: that keeps each Hessian matvec at two force evaluations. + n = x.size + ni = interior.size + + def hii(vi): + v = np.zeros(n) + v[interior] = vi + return np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float)[ + interior + ] + + def mii(bi): + v = np.zeros(n) + v[interior] = bi + return np.asarray(model.apply_preconditioner(np.ascontiguousarray(v)), float)[ + interior + ] + + return ( + LinearOperator((ni, ni), matvec=hii), # type: ignore[call-overload] + LinearOperator((ni, ni), matvec=mii), # type: ignore[call-overload] + ) + + +def solve_interior(model, x0, interior, boundary, x_boundary, tol=1e-10, max_newton=80): + """Converge the interior to force balance with the boundary held fixed. + + Preconditioned Newton-Krylov on the interior residual with a backtracking line + search; the line search is required for stiff 3D equilibria, where the full Newton + step overshoots. + """ + x = np.asarray(x0, float).copy() + x[boundary] = x_boundary + for _ in range(max_newton): + f = _raw_force(model, x) + norm0 = np.linalg.norm(f[interior]) + if norm0 < tol: + break + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, True) # assemble preconditioner + set base state + h_op, m_op = _interior_operators(model, x, interior) + dxi, _ = gmres(h_op, -f[interior], M=m_op, rtol=1e-4, maxiter=300) + alpha = 1.0 + for _ in range(30): + xt = x.copy() + xt[interior] += alpha * dxi + if np.linalg.norm(_raw_force(model, xt)[interior]) < norm0: + break + alpha *= 0.5 + else: + break # no decrease found; stop + x[interior] += alpha * dxi + return x + + +def objective_state_gradient(model, x, objective, h=1e-6): + """Partial derivative dJ/dx at fixed state, by central finite differences.""" + n = x.size + g = np.zeros(n) + for i in range(n): + xp = x.copy() + xp[i] += h + model.set_state(np.ascontiguousarray(xp)) + model.evaluate(2, 2, False) + jp = objective(model) + xm = x.copy() + xm[i] -= h + model.set_state(np.ascontiguousarray(xm)) + model.evaluate(2, 2, False) + jm = objective(model) + g[i] = (jp - jm) / (2 * h) + return g + + +def boundary_gradient(model, x_star, interior, boundary, objective, h=1e-6): + """Adjoint gradient dJ/dx_B at the converged equilibrium x_star.""" + n = x_star.size + dj = objective_state_gradient(model, x_star, objective, h) + model.set_state(np.ascontiguousarray(x_star)) + model.evaluate(2, 2, True) # assemble preconditioner + set base state to x_star + h_op, m_op = _interior_operators(model, x_star, interior) + lam, _ = gmres(h_op, dj[interior], M=m_op, rtol=1e-6, restart=100, maxiter=30) + embedded = np.zeros(n) + embedded[interior] = lam + model.set_state(np.ascontiguousarray(x_star)) + model.evaluate(2, 2, False) + coupling = np.asarray( + model.hessian_vector_product(np.ascontiguousarray(embedded)), float + )[boundary] + return dj[boundary] - coupling + + +def finite_difference_boundary_gradient( + model, x_star, interior, boundary, objective, dofs, h=1e-5 +): + """Reference gradient: re-solve the interior for each perturbed boundary DOF.""" + g = {} + for j in dofs: + xbp = x_star[boundary].copy() + xbp[j] += h + xp = solve_interior(model, x_star, interior, boundary, xbp) + model.set_state(np.ascontiguousarray(xp)) + model.evaluate(2, 2, False) + jp = objective(model) + xbm = x_star[boundary].copy() + xbm[j] -= h + xm = solve_interior(model, x_star, interior, boundary, xbm) + model.set_state(np.ascontiguousarray(xm)) + model.evaluate(2, 2, False) + jm = objective(model) + g[j] = (jp - jm) / (2 * h) + return g + + +def mhd_energy(model): + return model.mhd_energy diff --git a/src/vmecpp/__init__.py b/src/vmecpp/__init__.py index 9883157d5..492b87cc9 100644 --- a/src/vmecpp/__init__.py +++ b/src/vmecpp/__init__.py @@ -20,7 +20,7 @@ import pydantic from vmecpp import _util -from vmecpp._continuation import interpolate_solution, run_continuation +from vmecpp._continuation import _run_fourier_continuation, interpolate_solution from vmecpp._free_boundary import ( MagneticFieldResponseTable, MakegridParameters, @@ -77,6 +77,32 @@ def _wrap_int_as_float( pydantic.BeforeValidator(lambda x: np.array(x).astype(np.int64)), ] + +def _coerce_mpol_ntor(value: typing.Any) -> int | np.ndarray: + """Normalizes an ``mpol``/``ntor`` field value. + + A length-1 sequence is equivalent to a scalar and is collapsed to one; anything + longer is kept as an int array representing a per-``ns_array``-step Fourier + resolution continuation schedule. + """ + if isinstance(value, (int, np.integer)): + return int(value) + array = np.atleast_1d(np.asarray(value, dtype=np.int64)) + return int(array[0]) if array.size == 1 else array + + +MpolNtorField: typing.TypeAlias = typing.Annotated[ + int | jt.Int[np.ndarray, "num_fourier_steps"], + pydantic.BeforeValidator(_coerce_mpol_ntor), +] + + +def _final_resolution(value: int | np.ndarray) -> int: + """The target Fourier resolution: itself if scalar, else the schedule's last + (finest) entry.""" + return value if isinstance(value, int) else int(value[-1]) + + AuxFType = typing.Annotated[ _ArrayType, pydantic.BeforeValidator(lambda x: _util.right_pad(x, ndfmax, 0.0)), @@ -188,12 +214,23 @@ class VmecInput(BaseModelWithNumpy): nfp: int = 1 """Number of toroidal field periods (=1 for Tokamak)""" - mpol: int = 6 - """Number of poloidal Fourier harmonics; m = 0, 1, ..., (mpol-1)""" + mpol: MpolNtorField = 6 + """Number of poloidal Fourier harmonics; m = 0, 1, ..., (mpol-1). - ntor: int = 0 + May also be a sequence of ints, with one entry per ``ns_array`` step (a scalar + broadcasts to every step), to request continuation in Fourier resolution: + ``vmecpp.run()`` then solves each step in turn, hot-restarting from the + previous step's solution interpolated to the new resolution (see + :func:`interpolate_solution`). The boundary coefficients (``rbc``, ``zbs``, ...) + are always defined at the final (largest-index) entry's resolution. + """ + + ntor: MpolNtorField = 0 """Number of toroidal Fourier harmonics; n = -ntor, -ntor+1, ..., -1, 0, 1, ..., - ntor-1, ntor.""" + ntor-1, ntor. + + May be a sequence of ints, analogous to :attr:`mpol`; see its docstring. + """ mpol_geometry: int = -1 """Optional reduced poloidal resolution for the geometry (R, Z). @@ -465,7 +502,9 @@ def _validate_fourier_coefficients_shapes(self) -> VmecInput: if self.lasym: mpol_two_ntor_plus_one_fields.extend(["rbs", "zbc"]) - expected_shape = (self.mpol, 2 * self.ntor + 1) + mpol_final = _final_resolution(self.mpol) + ntor_final = _final_resolution(self.ntor) + expected_shape = (mpol_final, 2 * ntor_final + 1) for field in mpol_two_ntor_plus_one_fields: current_value = getattr(self, field) @@ -480,8 +519,8 @@ def _validate_fourier_coefficients_shapes(self) -> VmecInput: field, VmecInput.resize_2d_coeff( current_value, - mpol_new=self.mpol, - ntor_new=self.ntor, + mpol_new=mpol_final, + ntor_new=ntor_final, ), ) return self @@ -634,7 +673,9 @@ def _to_cpp_vmecindata(self) -> _vmecpp.VmecINDATA: ) # this also resizes the readonly_attrs - cpp_indata._set_mpol_ntor(self.mpol, self.ntor) + cpp_indata._set_mpol_ntor( + _final_resolution(self.mpol), _final_resolution(self.ntor) + ) for attr in readonly_attrs - {"mpol", "ntor"}: # now we can set the elements of the readonly_attrs value = getattr(self, attr) @@ -2234,6 +2275,14 @@ def run( restart_from: if present, VMEC++ is initialized using the converged equilibrium from the provided VmecOutput. This can dramatically decrease the number of iterations to convergence when running VMEC++ on a configuration that is very similar to the `restart_from` equilibrium. + If `input.mpol`/`input.ntor` is a sequence (see below), this is used to hot-restart + only the first continuation step; later steps always hot-restart from the previous one. + + If `input.mpol` and/or `input.ntor` is a sequence rather than a plain int, `run` performs + continuation in Fourier resolution: each entry pairs with the corresponding `input.ns_array` + entry (a scalar mpol/ntor broadcasts to every step), and each step is solved in turn, + hot-restarting from the previous step's solution interpolated to the new resolution (see + `interpolate_solution`). Example: >>> import vmecpp @@ -2244,6 +2293,16 @@ def run( 0.2033313711 """ input = VmecInput.model_validate(input) + + if not isinstance(input.mpol, int) or not isinstance(input.ntor, int): + return _run_fourier_continuation( + input, + magnetic_field, + max_threads=max_threads, + verbose=verbose, + restart_from=restart_from, + ) + cpp_indata = input._to_cpp_vmecindata() if restart_from is None: @@ -2490,7 +2549,6 @@ def set_profile( # items in the generated documentation. __all__ = [ # noqa: RUF022 "run", - "run_continuation", "interpolate_solution", "VmecInput", "VmecOutput", diff --git a/src/vmecpp/_continuation.py b/src/vmecpp/_continuation.py index b80cb5539..a4fe57d1f 100644 --- a/src/vmecpp/_continuation.py +++ b/src/vmecpp/_continuation.py @@ -1,10 +1,16 @@ """Generalized resolution interpolation and the Python-side continuation driver. VMEC++ converges much more reliably when a hard equilibrium is approached through -a sequence of increasing resolutions (the classic ``ns_array`` multi-grid, and now -also ``mpol_array`` / ``ntor_array`` Fourier continuation). Each step solves a single -resolution and hot-restarts from the previous step's solution, interpolated to the -new resolution by :func:`interpolate_solution`. +a sequence of increasing resolutions (the classic ``ns_array`` multi-grid, and also +Fourier continuation via a sequence-valued ``VmecInput.mpol`` / ``.ntor``). Each step +solves a single resolution and hot-restarts from the previous step's solution, +interpolated to the new resolution by :func:`interpolate_solution`. + +:func:`vmecpp.run` dispatches to :func:`_run_fourier_continuation` (this module) +whenever ``input.mpol`` and/or ``input.ntor`` is a sequence rather than a plain int; +:func:`interpolate_solution` itself is public API and can also be used directly, e.g. +to hand-roll a custom continuation schedule (see +``examples/fourier_resolution_increase.py``). The interpolation is purely a Python operation on a converged :class:`VmecOutput`: the flux-surface geometry is interpolated radially along the normalized toroidal @@ -22,7 +28,8 @@ import numpy as np if typing.TYPE_CHECKING: - from vmecpp import VmecInput, VmecOutput + from vmecpp import OutputMode, VmecInput, VmecOutput + from vmecpp._free_boundary import MagneticFieldResponseTable # State-vector geometry arrays, shape [mn_mode, n_surfaces]. These are the only # quantities VMEC++ reads back when hot-restarting, so they must be interpolated. @@ -285,63 +292,59 @@ def _step_input( return step -def run_continuation( +def _run_fourier_continuation( input: VmecInput, + magnetic_field: MagneticFieldResponseTable | None, *, - ns_array: typing.Sequence[int] | None = None, - mpol_array: typing.Sequence[int] | None = None, - ntor_array: typing.Sequence[int] | None = None, - ftol_array: typing.Sequence[float] | None = None, - niter_array: typing.Sequence[int] | None = None, - **run_kwargs: typing.Any, + max_threads: int | None, + verbose: bool | int | OutputMode, + restart_from: VmecOutput | None, ) -> VmecOutput: - """Solve an equilibrium by continuation in radial and Fourier resolution. + """Solves an equilibrium by continuation in Fourier resolution. + Called by :func:`vmecpp.run` whenever ``input.mpol`` and/or ``input.ntor`` is a + sequence rather than a plain int. Each entry pairs with the corresponding + ``input.ns_array`` entry (a scalar ``mpol``/``ntor`` broadcasts to every step). Each step solves a single ``(ns, mpol, ntor)`` resolution and hot-restarts from the previous step's solution interpolated to the new resolution (see - :func:`interpolate_solution`). This drives the classic ``ns_array`` multi-grid and, - by also increasing ``mpol`` / ``ntor`` along the schedule, Fourier continuation, - entirely from Python. + :func:`interpolate_solution`); if ``restart_from`` is given, it seeds the first + step instead of a cold start. Args: - input: the target configuration. Its boundary is the final-resolution boundary; - each step truncates it. Schedule arrays default to the corresponding fields - of ``input``; ``mpol_array`` / ``ntor_array`` default to constant - ``input.mpol`` / ``input.ntor`` (i.e. the classic fixed-Fourier multi-grid). - ns_array, mpol_array, ntor_array, ftol_array, niter_array: per-step schedules. - All provided arrays must share one length (a length-1 array is broadcast). - **run_kwargs: forwarded to :func:`vmecpp.run` for every step (e.g. ``verbose``, - ``max_threads``). + input: the target configuration. Its boundary is the final-resolution + boundary; each step truncates or zero-pads it to that step's resolution. + magnetic_field, max_threads, verbose, restart_from: forwarded to + :func:`vmecpp.run` for every step (``restart_from`` only seeds the first). Returns: - The converged :class:`VmecOutput` at the final resolution. + The converged :class:`VmecOutput` at the final resolution, with ``input`` set + to the original (full-schedule) ``input`` argument. """ import vmecpp # noqa: PLC0415 (lazy import avoids a circular import) - ns_schedule = [int(x) for x in (input.ns_array if ns_array is None else ns_array)] + ns_schedule = [int(x) for x in input.ns_array] n_steps = len(ns_schedule) - if n_steps == 0: - msg = "ns_array must have at least one entry" - raise ValueError(msg) - - def _resolve(values: typing.Sequence[float] | None, default: list) -> list: - resolved = list(default) if values is None else list(values) - if len(resolved) == 1: - resolved = resolved * n_steps + + def _resolve(value: int | np.ndarray, name: str) -> list[int]: + if isinstance(value, int): + return [value] * n_steps + resolved = [int(x) for x in value] if len(resolved) != n_steps: msg = ( - f"continuation schedule length {len(resolved)} does not match " - f"ns_array length {n_steps}" + f"'{name}' has {len(resolved)} entries, but 'ns_array' has " + f"{n_steps}; a Fourier-resolution continuation schedule must have " + "one entry per ns_array step (or be a scalar, broadcast to every " + "step)." ) raise ValueError(msg) return resolved - mpol_schedule = _resolve(mpol_array, [int(input.mpol)] * n_steps) - ntor_schedule = _resolve(ntor_array, [int(input.ntor)] * n_steps) - ftol_schedule = _resolve(ftol_array, list(np.asarray(input.ftol_array))) - niter_schedule = _resolve(niter_array, list(np.asarray(input.niter_array))) + mpol_schedule = _resolve(input.mpol, "mpol") + ntor_schedule = _resolve(input.ntor, "ntor") + ftol_schedule = [float(x) for x in input.ftol_array] + niter_schedule = [int(x) for x in input.niter_array] - output: VmecOutput | None = None + output = restart_from for i in range(n_steps): step_input = _step_input( input, @@ -351,10 +354,14 @@ def _resolve(values: typing.Sequence[float] | None, default: list) -> list: ftol_schedule[i], niter_schedule[i], ) - if output is None: - output = vmecpp.run(step_input, **run_kwargs) - else: - guess = interpolate_solution(output, step_input) - output = vmecpp.run(step_input, restart_from=guess, **run_kwargs) + guess = None if output is None else interpolate_solution(output, step_input) + output = vmecpp.run( + step_input, + magnetic_field, + max_threads=max_threads, + verbose=verbose, + restart_from=guess, + ) + assert output is not None # n_steps >= 1, so the loop always assigns output - return output + return output.model_copy(update={"input": input}) diff --git a/src/vmecpp/_iteration.py b/src/vmecpp/_iteration.py index b019a79fa..e54a8cfc5 100644 --- a/src/vmecpp/_iteration.py +++ b/src/vmecpp/_iteration.py @@ -259,7 +259,7 @@ def solve_equilibrium( iter2 = force_iteration - bad_resets # Evolve: forward model, then convergence / damping / time step. - model.evaluate(iter1, iter2) + model.evaluate(iter1, iter2, precondition=True, always_fix_m1_gauge=False) fsqr, fsqz, fsql = model.fsqr, model.fsqz, model.fsql rr = model.restart_reason finite = math.isfinite(fsqr) and math.isfinite(fsqz) and math.isfinite(fsql) diff --git a/src/vmecpp/cpp/docker/tsan/README.md b/src/vmecpp/cpp/docker/tsan/README.md index 33bd6bbe2..c60eea9a8 100644 --- a/src/vmecpp/cpp/docker/tsan/README.md +++ b/src/vmecpp/cpp/docker/tsan/README.md @@ -1,6 +1,6 @@ # OpenMP-aware ThreadSanitizer -This is a tutorial how to build `clang` and `libomp` with ThreadSanitizer (TSan) +This is a tutorial on how to build `clang` and `libomp` with ThreadSanitizer (TSan) support for being able to check OpenMP-parallelized C/C++ code. ## Perform ThreadSanitizer runs on code in the Proxima Fusion monorepo diff --git a/src/vmecpp/cpp/docker/tsan/background_info.md b/src/vmecpp/cpp/docker/tsan/background_info.md index 60035b0dc..0a3b9a58d 100644 --- a/src/vmecpp/cpp/docker/tsan/background_info.md +++ b/src/vmecpp/cpp/docker/tsan/background_info.md @@ -28,7 +28,7 @@ found here: - https://www.vi-hps.org/cms/upload/material/tw30/Archer.pdf (slide 8 and onwards) Turns out, this is **deprecated** as well and Archer was integrated into the -LLVM project in the mean time. +LLVM project in the meantime. This seems to be the **state of things as of Oct 2023**. Indeed, Archer can be found in the `llvm` source tree: @@ -46,7 +46,7 @@ https://packages.ubuntu.com/jammy/clang ## Setting up a first test case TL;DR: This is an example commonly presented as a case for data races. As will -be seen when running this example, the errornous output due to a data race is +be seen when running this example, the erroneous output due to a data race is not reliably reproduced. Thus, skip this section if you are looking for an example that breaks reliably. @@ -223,7 +223,7 @@ int main(int argc, char** argv) { } ``` -Compile it (for now without TSan to not clobber the console when things to +Compile it (for now without TSan to not clobber the console when things go sideways): ```bash @@ -313,7 +313,7 @@ Now run it (two threads should be enough to trigger the data race checks): OMP_NUM_THREADS=2 ./pi_example ``` -However, contrary to the expection, TSan report data races, even though the +However, contrary to the expectation, TSan reports data races, even though the result is always computed correctly: ``` @@ -427,7 +427,7 @@ Now, adjust the command line as requested: OMP_NUM_THREADS=1 ARCHER_OPTIONS='verbose=1' TSAN_OPTIONS='ignore_noninstrumented_modules=1' ./pi_example ``` -and we gete: +and we get: ``` Archer detected OpenMP application with TSan, supplying OpenMP synchronization semantics @@ -505,7 +505,7 @@ This is what we work on next. ## Build the stage1 demo using a custom toolchain -from: https://bazel.build/tutorials/ccp-toolchain-config +from: https://bazel.build/tutorials/cpp-toolchain-config This uses a system-provided `clang` installation (instead of the default compiler Bazel uses). @@ -529,7 +529,7 @@ bazel build --config=clang_config //abseil-hello:hello_main bazel-bin/abseil-hello/hello_main "from Abseil" ``` -Fixed by added `-lm` to standard linker flags. This was inspired by: +Fixed by adding `-lm` to standard linker flags. This was inspired by: https://github.com/bazelbuild/bazel/issues/934#issuecomment-193474914 Even the unit test works if a recent version of `googletest` is used: @@ -633,7 +633,7 @@ Delete the Bazel cache: bazel clean --expunge ``` -The PDF files of the slides mentioned in this articles are mirrored locally +The PDF files of the slides mentioned in this article are mirrored locally here: https://drive.google.com/drive/folders/1NBNTr4jDQy951CoG-AYqpDfKKpNKNSzh?usp=sharing diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc index be9acf83c..d2e92218e 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dft_toroidal.cc @@ -51,95 +51,75 @@ void ForcesToFourier3DSymmFastPoloidal( const auto& fzcon = m_even ? d.fzcon_e : d.fzcon_o; for (int k = 0; k < s.nZeta; ++k) { + double rmkcc = 0.0; + double rmkcc_n = 0.0; + double rmkss = 0.0; + double rmkss_n = 0.0; + double zmksc = 0.0; + double zmksc_n = 0.0; + double zmkcs = 0.0; + double zmkcs_n = 0.0; + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; const int idx_ml_base = m * s.nThetaReduced; - // Vectorized poloidal loop using Eigen operations - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); - - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - auto crmn_seg = Eigen::Map( - crmn.data() + idx_kl_base, s.nThetaReduced); - auto czmn_seg = Eigen::Map( - czmn.data() + idx_kl_base, s.nThetaReduced); - auto armn_seg = Eigen::Map( - armn.data() + idx_kl_base, s.nThetaReduced); - auto azmn_seg = Eigen::Map( - azmn.data() + idx_kl_base, s.nThetaReduced); - auto brmn_seg = Eigen::Map( - brmn.data() + idx_kl_base, s.nThetaReduced); - auto bzmn_seg = Eigen::Map( - bzmn.data() + idx_kl_base, s.nThetaReduced); - auto frcon_seg = Eigen::Map( - frcon.data() + idx_kl_base, s.nThetaReduced); - auto fzcon_seg = Eigen::Map( - fzcon.data() + idx_kl_base, s.nThetaReduced); - - double lmksc = blmn_seg.dot(cosmumi_seg); - double lmkcs = blmn_seg.dot(sinmumi_seg); - double lmkcs_n = -clmn_seg.dot(cosmui_seg); - double lmksc_n = -clmn_seg.dot(sinmui_seg); - - double rmkcc_n = -crmn_seg.dot(cosmui_seg); - double zmkcs_n = -czmn_seg.dot(cosmui_seg); - - double rmkss_n = -crmn_seg.dot(sinmui_seg); - double zmksc_n = -czmn_seg.dot(sinmui_seg); - - // Assemble effective R and Z forces from MHD and spectral condensation - // contributions. Materialize to avoid re-evaluation in each dot - // product. - // Per-thread scratch reused across iterations instead of a heap - // temporary in this innermost loop; still materialized once and then - // used in the two dot products below. - thread_local Eigen::VectorXd tempR_seg; - thread_local Eigen::VectorXd tempZ_seg; - tempR_seg = armn_seg + xmpq[m] * frcon_seg; - tempZ_seg = azmn_seg + xmpq[m] * fzcon_seg; - - double rmkcc = tempR_seg.dot(cosmui_seg) + brmn_seg.dot(sinmumi_seg); - double rmkss = tempR_seg.dot(sinmui_seg) + brmn_seg.dot(cosmumi_seg); - double zmksc = tempZ_seg.dot(sinmui_seg) + bzmn_seg.dot(cosmumi_seg); - double zmkcs = tempZ_seg.dot(cosmui_seg) + bzmn_seg.dot(sinmumi_seg); - - // Vectorized toroidal scatter: segment ops replace scalar n-loop - const int ntorp1 = s.ntor + 1; - const int idx_mn_base = ((jF - rp.nsMinF) * s.mpol + m) * ntorp1; - const int idx_kn_base = k * (s.nnyq2 + 1); - - auto cosnv_seg = fb.cosnv.segment(idx_kn_base, ntorp1); - auto sinnv_seg = fb.sinnv.segment(idx_kn_base, ntorp1); - auto cosnvn_seg = fb.cosnvn.segment(idx_kn_base, ntorp1); - auto sinnvn_seg = fb.sinnvn.segment(idx_kn_base, ntorp1); - - Eigen::Map frcc_seg( - m_physical_forces.frcc.data() + idx_mn_base, ntorp1); - Eigen::Map frss_seg( - m_physical_forces.frss.data() + idx_mn_base, ntorp1); - Eigen::Map fzsc_seg( - m_physical_forces.fzsc.data() + idx_mn_base, ntorp1); - Eigen::Map fzcs_seg( - m_physical_forces.fzcs.data() + idx_mn_base, ntorp1); - - frcc_seg += rmkcc * cosnv_seg + rmkcc_n * sinnvn_seg; - frss_seg += rmkss * sinnv_seg + rmkss_n * cosnvn_seg; - fzsc_seg += zmksc * cosnv_seg + zmksc_n * sinnvn_seg; - fzcs_seg += zmkcs * sinnv_seg + zmkcs_n * cosnvn_seg; - - if (jMinL <= jF) { - Eigen::Map flsc_seg( - m_physical_forces.flsc.data() + idx_mn_base, ntorp1); - Eigen::Map flcs_seg( - m_physical_forces.flcs.data() + idx_mn_base, ntorp1); - flsc_seg += lmksc * cosnv_seg + lmksc_n * sinnvn_seg; - flcs_seg += lmkcs * sinnv_seg + lmkcs_n * cosnvn_seg; - } + // NOTE: nThetaReduced is usually pretty small, 9 for cma.json + // and 16 for w7x_ref_167_12_12.json, so in our benchmark forcing + // the compiler to auto-vectorize this loop was a pessimization. + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; // --> flsc (no A) + lmkcs += blmn[idx_kl] * sinmumi; // --> flcs + lmkcs_n -= clmn[idx_kl] * cosmui; // --> flcs + lmksc_n -= clmn[idx_kl] * sinmui; // --> flsc + + rmkcc_n -= crmn[idx_kl] * cosmui; // --> frcc + zmkcs_n -= czmn[idx_kl] * cosmui; // --> fzcs + + rmkss_n -= crmn[idx_kl] * sinmui; // --> frss + zmksc_n -= czmn[idx_kl] * sinmui; // --> fzsc + + // assemble effective R and Z forces from MHD and spectral + // condensation contributions + const double tempR = armn[idx_kl] + xmpq[m] * frcon[idx_kl]; + const double tempZ = azmn[idx_kl] + xmpq[m] * fzcon[idx_kl]; + + rmkcc += tempR * cosmui + brmn[idx_kl] * sinmumi; // --> frcc + rmkss += tempR * sinmui + brmn[idx_kl] * cosmumi; // --> frss + zmksc += tempZ * sinmui + bzmn[idx_kl] * cosmumi; // --> fzsc + zmkcs += tempZ * cosmui + bzmn[idx_kl] * sinmumi; // --> fzcs + } // l + + for (int n = 0; n < s.ntor + 1; ++n) { + const int idx_mn = ((jF - rp.nsMinF) * s.mpol + m) * (s.ntor + 1) + n; + const int idx_kn = k * (s.nnyq2 + 1) + n; + + const double cosnv = fb.cosnv[idx_kn]; + const double sinnv = fb.sinnv[idx_kn]; + const double cosnvn = fb.cosnvn[idx_kn]; + const double sinnvn = fb.sinnvn[idx_kn]; + + m_physical_forces.frcc[idx_mn] += rmkcc * cosnv + rmkcc_n * sinnvn; + m_physical_forces.frss[idx_mn] += rmkss * sinnv + rmkss_n * cosnvn; + m_physical_forces.fzsc[idx_mn] += zmksc * cosnv + zmksc_n * sinnvn; + m_physical_forces.fzcs[idx_mn] += zmkcs * sinnv + zmkcs_n * cosnvn; + + if (jMinL <= jF) { + m_physical_forces.flsc[idx_mn] += lmksc * cosnv + lmksc_n * sinnvn; + m_physical_forces.flcs[idx_mn] += lmkcs * sinnv + lmkcs_n * cosnvn; + } + } // n } // k } // m } // jF @@ -154,42 +134,41 @@ void ForcesToFourier3DSymmFastPoloidal( const auto& clmn = m_even ? d.clmn_e : d.clmn_o; for (int k = 0; k < s.nZeta; ++k) { + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; const int idx_ml_base = m * s.nThetaReduced; - // Vectorized poloidal loop using Eigen operations - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); - - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - - double lmksc = blmn_seg.dot(cosmumi_seg); - double lmkcs = blmn_seg.dot(sinmumi_seg); - double lmkcs_n = -clmn_seg.dot(cosmui_seg); - double lmksc_n = -clmn_seg.dot(sinmui_seg); - - // Vectorized toroidal scatter for lambda-only section - const int ntorp1 = s.ntor + 1; - const int idx_mn_base = ((jF - rp.nsMinF) * s.mpol + m) * ntorp1; - const int idx_kn_base = k * (s.nnyq2 + 1); - - auto cosnv_seg = fb.cosnv.segment(idx_kn_base, ntorp1); - auto sinnv_seg = fb.sinnv.segment(idx_kn_base, ntorp1); - auto cosnvn_seg = fb.cosnvn.segment(idx_kn_base, ntorp1); - auto sinnvn_seg = fb.sinnvn.segment(idx_kn_base, ntorp1); - - Eigen::Map flsc_seg( - m_physical_forces.flsc.data() + idx_mn_base, ntorp1); - Eigen::Map flcs_seg( - m_physical_forces.flcs.data() + idx_mn_base, ntorp1); - - flsc_seg += lmksc * cosnv_seg + lmksc_n * sinnvn_seg; - flcs_seg += lmkcs * sinnv_seg + lmkcs_n * cosnvn_seg; + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; // --> flsc (no A) + lmkcs += blmn[idx_kl] * sinmumi; // --> flcs + lmkcs_n -= clmn[idx_kl] * cosmui; // --> flcs + lmksc_n -= clmn[idx_kl] * sinmui; // --> flsc + } // l + + for (int n = 0; n < s.ntor + 1; ++n) { + const int idx_mn = ((jF - rp.nsMinF) * s.mpol + m) * (s.ntor + 1) + n; + const int idx_kn = k * (s.nnyq2 + 1) + n; + + const double cosnv = fb.cosnv[idx_kn]; + const double sinnv = fb.sinnv[idx_kn]; + const double cosnvn = fb.cosnvn[idx_kn]; + const double sinnvn = fb.sinnvn[idx_kn]; + + m_physical_forces.flsc[idx_mn] += lmksc * cosnv + lmksc_n * sinnvn; + m_physical_forces.flcs[idx_mn] += lmkcs * sinnv + lmkcs_n * cosnvn; + } // n } // k } // m } // jF @@ -257,97 +236,100 @@ void FourierToReal3DSymmFastPoloidal(const FourierGeometry& physical_x, } for (int k = 0; k < s.nZeta; ++k) { - // INVERSE TRANSFORM IN N-ZETA, FOR FIXED M - // Vectorized toroidal accumulation loop - const int idx_kn_base = k * (s.nnyq2 + 1); - const int idx_mn_base = ((jF - nsMinF1) * s.mpol + m) * (s.ntor + 1); - - auto cosnv_seg = fb.cosnv.segment(idx_kn_base, s.ntor + 1); - auto sinnv_seg = fb.sinnv.segment(idx_kn_base, s.ntor + 1); - auto sinnvn_seg = fb.sinnvn.segment(idx_kn_base, s.ntor + 1); - auto cosnvn_seg = fb.cosnvn.segment(idx_kn_base, s.ntor + 1); - - auto rmncc_seg = Eigen::Map( - physical_x.rmncc.data() + idx_mn_base, s.ntor + 1); - auto rmnss_seg = Eigen::Map( - physical_x.rmnss.data() + idx_mn_base, s.ntor + 1); - auto zmnsc_seg = Eigen::Map( - physical_x.zmnsc.data() + idx_mn_base, s.ntor + 1); - auto zmncs_seg = Eigen::Map( - physical_x.zmncs.data() + idx_mn_base, s.ntor + 1); - auto lmnsc_seg = Eigen::Map( - physical_x.lmnsc.data() + idx_mn_base, s.ntor + 1); - auto lmncs_seg = Eigen::Map( - physical_x.lmncs.data() + idx_mn_base, s.ntor + 1); - - double rmkcc = rmncc_seg.dot(cosnv_seg); - double rmkcc_n = rmncc_seg.dot(sinnvn_seg); - double rmkss = rmnss_seg.dot(sinnv_seg); - double rmkss_n = rmnss_seg.dot(cosnvn_seg); - double zmksc = zmnsc_seg.dot(cosnv_seg); - double zmksc_n = zmnsc_seg.dot(sinnvn_seg); - double zmkcs = zmncs_seg.dot(sinnv_seg); - double zmkcs_n = zmncs_seg.dot(cosnvn_seg); - double lmksc = lmnsc_seg.dot(cosnv_seg); - double lmksc_n = lmnsc_seg.dot(sinnvn_seg); - double lmkcs = lmncs_seg.dot(sinnv_seg); - double lmkcs_n = lmncs_seg.dot(cosnvn_seg); + double rmkcc = 0.0; + double rmkcc_n = 0.0; + double rmkss = 0.0; + double rmkss_n = 0.0; + double zmksc = 0.0; + double zmksc_n = 0.0; + double zmkcs = 0.0; + double zmkcs_n = 0.0; + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + + for (int n = 0; n < s.ntor + 1; ++n) { + // INVERSE TRANSFORM IN N-ZETA, FOR FIXED M + + const int idx_kn = k * (s.nnyq2 + 1) + n; + + double cosnv = fb.cosnv[idx_kn]; + double sinnv = fb.sinnv[idx_kn]; + double sinnvn = fb.sinnvn[idx_kn]; + double cosnvn = fb.cosnvn[idx_kn]; + + int idx_mn = ((jF - nsMinF1) * s.mpol + m) * (s.ntor + 1) + n; + + rmkcc += physical_x.rmncc[idx_mn] * cosnv; + rmkcc_n += physical_x.rmncc[idx_mn] * sinnvn; + rmkss += physical_x.rmnss[idx_mn] * sinnv; + rmkss_n += physical_x.rmnss[idx_mn] * cosnvn; + zmksc += physical_x.zmnsc[idx_mn] * cosnv; + zmksc_n += physical_x.zmnsc[idx_mn] * sinnvn; + zmkcs += physical_x.zmncs[idx_mn] * sinnv; + zmkcs_n += physical_x.zmncs[idx_mn] * cosnvn; + lmksc += physical_x.lmnsc[idx_mn] * cosnv; + lmksc_n += physical_x.lmnsc[idx_mn] * sinnvn; + lmkcs += physical_x.lmncs[idx_mn] * sinnv; + lmkcs_n += physical_x.lmncs[idx_mn] * cosnvn; + } // n // INVERSE TRANSFORM IN M-THETA, FOR ALL RADIAL, ZETA VALUES const int idx_kl_base = ((jF - nsMinF1) * s.nZeta + k) * s.nThetaEff; - // Vectorized poloidal loops using Eigen operations - auto sinmum_seg = fb.sinmum.segment(idx_ml_base, s.nThetaReduced); - auto cosmum_seg = fb.cosmum.segment(idx_ml_base, s.nThetaReduced); - - auto ru_seg = Eigen::Map(ru.data() + idx_kl_base, - s.nThetaReduced); - auto zu_seg = Eigen::Map(zu.data() + idx_kl_base, - s.nThetaReduced); - auto lu_seg = Eigen::Map(lu.data() + idx_kl_base, - s.nThetaReduced); - - // NOTE: element-wise multiplication - ru_seg += rmkcc * sinmum_seg + rmkss * cosmum_seg; - zu_seg += zmksc * cosmum_seg + zmkcs * sinmum_seg; - lu_seg += lmksc * cosmum_seg + lmkcs * sinmum_seg; - - auto cosmu_seg = fb.cosmu.segment(idx_ml_base, s.nThetaReduced); - auto sinmu_seg = fb.sinmu.segment(idx_ml_base, s.nThetaReduced); - - auto rv_seg = Eigen::Map(rv.data() + idx_kl_base, - s.nThetaReduced); - auto zv_seg = Eigen::Map(zv.data() + idx_kl_base, - s.nThetaReduced); - auto lv_seg = Eigen::Map(lv.data() + idx_kl_base, - s.nThetaReduced); - - // NOTE: element-wise multiplication - rv_seg += rmkcc_n * cosmu_seg + rmkss_n * sinmu_seg; - zv_seg += zmksc_n * sinmu_seg + zmkcs_n * cosmu_seg; - // it is here that lv gets a negative sign! - lv_seg -= lmksc_n * sinmu_seg + lmkcs_n * cosmu_seg; - - auto r1_seg = Eigen::Map(r1.data() + idx_kl_base, - s.nThetaReduced); - auto z1_seg = Eigen::Map(z1.data() + idx_kl_base, - s.nThetaReduced); - - r1_seg += rmkcc * cosmu_seg + rmkss * sinmu_seg; - z1_seg += zmksc * sinmu_seg + zmkcs * cosmu_seg; + // the loop over l is split to help compiler auto-vectorization + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_ml = idx_ml_base + l; - if (nsMinF <= jF && jF < r.nsMaxFIncludingLcfs) { - // spectral condensation is local per flux surface - const int idx_con_base = ((jF - nsMinF) * s.nZeta + k) * s.nThetaEff; + const double sinmum = fb.sinmum[idx_ml]; + const double cosmum = fb.cosmum[idx_ml]; + + const int idx_kl = idx_kl_base + l; + ru[idx_kl] += rmkcc * sinmum + rmkss * cosmum; + zu[idx_kl] += zmksc * cosmum + zmkcs * sinmum; + lu[idx_kl] += lmksc * cosmum + lmkcs * sinmum; + } // l + + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; - auto rCon_seg = Eigen::Map( - m_geometry.rCon.data() + idx_con_base, s.nThetaReduced); - auto zCon_seg = Eigen::Map( - m_geometry.zCon.data() + idx_con_base, s.nThetaReduced); + const double cosmu = fb.cosmu[idx_ml]; + const double sinmu = fb.sinmu[idx_ml]; + rv[idx_kl] += rmkcc_n * cosmu + rmkss_n * sinmu; + zv[idx_kl] += zmksc_n * sinmu + zmkcs_n * cosmu; + // it is here that lv gets a negative sign! + lv[idx_kl] -= lmksc_n * sinmu + lmkcs_n * cosmu; + } // l - rCon_seg += (rmkcc * cosmu_seg + rmkss * sinmu_seg) * con_factor; - zCon_seg += (zmksc * sinmu_seg + zmkcs * cosmu_seg) * con_factor; - } + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_ml = idx_ml_base + l; + + const double cosmu = fb.cosmu[idx_ml]; + const double sinmu = fb.sinmu[idx_ml]; + + const int idx_kl = idx_kl_base + l; + + r1[idx_kl] += rmkcc * cosmu + rmkss * sinmu; + z1[idx_kl] += zmksc * sinmu + zmkcs * cosmu; + } // l + + if (nsMinF <= jF && jF < r.nsMaxFIncludingLcfs) { + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_ml = idx_ml_base + l; + const double cosmu = fb.cosmu[idx_ml]; + const double sinmu = fb.sinmu[idx_ml]; + + // spectral condensation is local per flux surface + // --> no need for numFull1 + const int idx_con = ((jF - nsMinF) * s.nZeta + k) * s.nThetaEff + l; + m_geometry.rCon[idx_con] += + (rmkcc * cosmu + rmkss * sinmu) * con_factor; + m_geometry.zCon[idx_con] += + (zmksc * sinmu + zmkcs * cosmu) * con_factor; + } + } // l } // k } // m } // j diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc index 8b845c752..4323a1c30 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal.cc @@ -459,10 +459,6 @@ void ForcesToFourier3DSymmFastPoloidalFft( const auto& fzcon = m_even ? d.fzcon_e : d.fzcon_o; const int idx_ml_base = m * s.nThetaReduced; - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); double* rmkcc_buf = in_slot(m, kRmkcc); double* rmkss_buf = in_slot(m, kRmkss); @@ -477,47 +473,65 @@ void ForcesToFourier3DSymmFastPoloidalFft( double* lmksc_n_buf = in_slot(m, kLmkscN); double* lmkcs_n_buf = in_slot(m, kLmkcsN); + const double xmpq_m = xmpq[m]; for (int k = 0; k < s.nZeta; ++k) { const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - auto crmn_seg = Eigen::Map( - crmn.data() + idx_kl_base, s.nThetaReduced); - auto czmn_seg = Eigen::Map( - czmn.data() + idx_kl_base, s.nThetaReduced); - auto armn_seg = Eigen::Map( - armn.data() + idx_kl_base, s.nThetaReduced); - auto azmn_seg = Eigen::Map( - azmn.data() + idx_kl_base, s.nThetaReduced); - auto brmn_seg = Eigen::Map( - brmn.data() + idx_kl_base, s.nThetaReduced); - auto bzmn_seg = Eigen::Map( - bzmn.data() + idx_kl_base, s.nThetaReduced); - auto frcon_seg = Eigen::Map( - frcon.data() + idx_kl_base, s.nThetaReduced); - auto fzcon_seg = Eigen::Map( - fzcon.data() + idx_kl_base, s.nThetaReduced); - - lmksc_buf[k] = blmn_seg.dot(cosmumi_seg); - lmkcs_buf[k] = blmn_seg.dot(sinmumi_seg); - lmkcs_n_buf[k] = -clmn_seg.dot(cosmui_seg); - lmksc_n_buf[k] = -clmn_seg.dot(sinmui_seg); - - rmkcc_n_buf[k] = -crmn_seg.dot(cosmui_seg); - zmkcs_n_buf[k] = -czmn_seg.dot(cosmui_seg); - rmkss_n_buf[k] = -crmn_seg.dot(sinmui_seg); - zmksc_n_buf[k] = -czmn_seg.dot(sinmui_seg); - - const Eigen::VectorXd tempR = (armn_seg + xmpq[m] * frcon_seg).eval(); - const Eigen::VectorXd tempZ = (azmn_seg + xmpq[m] * fzcon_seg).eval(); - - rmkcc_buf[k] = tempR.dot(cosmui_seg) + brmn_seg.dot(sinmumi_seg); - rmkss_buf[k] = tempR.dot(sinmui_seg) + brmn_seg.dot(cosmumi_seg); - zmksc_buf[k] = tempZ.dot(sinmui_seg) + bzmn_seg.dot(cosmumi_seg); - zmkcs_buf[k] = tempZ.dot(cosmui_seg) + bzmn_seg.dot(sinmumi_seg); + double rmkcc = 0.0; + double rmkcc_n = 0.0; + double rmkss = 0.0; + double rmkss_n = 0.0; + double zmksc = 0.0; + double zmksc_n = 0.0; + double zmkcs = 0.0; + double zmkcs_n = 0.0; + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + + // Fused poloidal reduction, matching ForcesToFourier3DSymmFastPoloidal: + // one pass over the short theta axis, no per-iteration heap temporary. + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; + lmkcs += blmn[idx_kl] * sinmumi; + lmkcs_n -= clmn[idx_kl] * cosmui; + lmksc_n -= clmn[idx_kl] * sinmui; + + rmkcc_n -= crmn[idx_kl] * cosmui; + zmkcs_n -= czmn[idx_kl] * cosmui; + rmkss_n -= crmn[idx_kl] * sinmui; + zmksc_n -= czmn[idx_kl] * sinmui; + + const double tempR = armn[idx_kl] + xmpq_m * frcon[idx_kl]; + const double tempZ = azmn[idx_kl] + xmpq_m * fzcon[idx_kl]; + + rmkcc += tempR * cosmui + brmn[idx_kl] * sinmumi; + rmkss += tempR * sinmui + brmn[idx_kl] * cosmumi; + zmksc += tempZ * sinmui + bzmn[idx_kl] * cosmumi; + zmkcs += tempZ * cosmui + bzmn[idx_kl] * sinmumi; + } // l + + rmkcc_buf[k] = rmkcc; + rmkss_buf[k] = rmkss; + rmkcc_n_buf[k] = rmkcc_n; + rmkss_n_buf[k] = rmkss_n; + zmksc_buf[k] = zmksc; + zmkcs_buf[k] = zmkcs; + zmksc_n_buf[k] = zmksc_n; + zmkcs_n_buf[k] = zmkcs_n; + lmksc_buf[k] = lmksc; + lmkcs_buf[k] = lmkcs; + lmksc_n_buf[k] = lmksc_n; + lmkcs_n_buf[k] = lmkcs_n; } // k } // m (fill) @@ -590,10 +604,6 @@ void ForcesToFourier3DSymmFastPoloidalFft( const auto& clmn = m_even ? d.clmn_e : d.clmn_o; const int idx_ml_base = m * s.nThetaReduced; - auto cosmui_seg = fb.cosmui.segment(idx_ml_base, s.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(idx_ml_base, s.nThetaReduced); - auto cosmumi_seg = fb.cosmumi.segment(idx_ml_base, s.nThetaReduced); - auto sinmumi_seg = fb.sinmumi.segment(idx_ml_base, s.nThetaReduced); double* lmksc_buf = in_slot(m, kLmksc); double* lmkcs_buf = in_slot(m, kLmkcs); @@ -602,15 +612,32 @@ void ForcesToFourier3DSymmFastPoloidalFft( for (int k = 0; k < s.nZeta; ++k) { const int idx_kl_base = ((jF - rp.nsMinF) * s.nZeta + k) * s.nThetaEff; - auto blmn_seg = Eigen::Map( - blmn.data() + idx_kl_base, s.nThetaReduced); - auto clmn_seg = Eigen::Map( - clmn.data() + idx_kl_base, s.nThetaReduced); - - lmksc_buf[k] = blmn_seg.dot(cosmumi_seg); - lmkcs_buf[k] = blmn_seg.dot(sinmumi_seg); - lmkcs_n_buf[k] = -clmn_seg.dot(cosmui_seg); - lmksc_n_buf[k] = -clmn_seg.dot(sinmui_seg); + + double lmksc = 0.0; + double lmksc_n = 0.0; + double lmkcs = 0.0; + double lmkcs_n = 0.0; + + // Fused poloidal reduction (see main-loop note). + for (int l = 0; l < s.nThetaReduced; ++l) { + const int idx_kl = idx_kl_base + l; + const int idx_ml = idx_ml_base + l; + + const double cosmui = fb.cosmui[idx_ml]; + const double sinmui = fb.sinmui[idx_ml]; + const double cosmumi = fb.cosmumi[idx_ml]; + const double sinmumi = fb.sinmumi[idx_ml]; + + lmksc += blmn[idx_kl] * cosmumi; + lmkcs += blmn[idx_kl] * sinmumi; + lmkcs_n -= clmn[idx_kl] * cosmui; + lmksc_n -= clmn[idx_kl] * sinmui; + } // l + + lmksc_buf[k] = lmksc; + lmkcs_buf[k] = lmkcs; + lmksc_n_buf[k] = lmksc_n; + lmkcs_n_buf[k] = lmkcs_n; } // k } // m (fill) diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc index 1ec40ce51..66b746e1c 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc @@ -421,7 +421,8 @@ absl::StatusOr IdealMhdModel::update( bool& m_need_restart, int& m_last_preconditioner_update, int& m_last_full_update_nestor, FlowControl& m_fc, const int iter1, const int iter2, const VmecCheckpoint& checkpoint, - const int iterations_before_checkpointing, bool verbose) { + const int iterations_before_checkpointing, bool verbose, + bool always_fix_m1_gauge) { // An axis re-guess after a bad Jacobian can repopulate high geometry modes // directly, bypassing the force mask; clear them on the state each iteration. if (s_.mpolGeometry < s_.mpol || s_.ntorGeometry < s_.ntor) { @@ -813,7 +814,9 @@ absl::StatusOr IdealMhdModel::update( m_decomposed_f.m1Constraint(1.0 / std::numbers::sqrt2); // v8.50: ADD iter2<2 so reset= works - if (m_fc.fsqz < 1.0e-6 || iter2 < 2) { + const bool fix_m1_gauge = + always_fix_m1_gauge || m_fc.fsqz < 1.0e-6 || iter2 < 2; + if (fix_m1_gauge) { // ensure that the m=1 constraint is satisfied exactly // --> the corresponding m=1 coeffs of R,Z are constrained to be zero // and thus must not be "forced" (by the time evol using gc) away from diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h index b753e53f5..da3cf19fc 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h @@ -70,7 +70,8 @@ class IdealMhdModel { bool& m_need_restart, int& m_last_preconditioner_update, int& m_last_full_update_nestor, FlowControl& m_fc, const int iter1, const int iter2, const VmecCheckpoint& checkpoint = VmecCheckpoint::NONE, - const int iterations_before_checkpointing = INT_MAX, bool verbose = true); + const int iterations_before_checkpointing = INT_MAX, bool verbose = true, + bool always_fix_m1_gauge = false); // Coordinates which inverse-DFT routine to call for computing // the flux surface geometry and lambda on it from the provided Fourier diff --git a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc index 063ac9ca0..5e2510d7d 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc @@ -226,7 +226,12 @@ class VmecModel { // lambda-constraint components. That raw gradient is what gradient-based // optimizers minimizing the MHD energy functional need; mhd_energy is already // set earlier in update(), so it is valid at the checkpoint too. - void Evaluate(int iter1, int iter2, bool precondition = true) { + // The native iteration leaves the m=1 gauge free until the previous Z + // residual crosses its threshold. External evaluations fix it immediately so + // F(x) does not depend on the previously evaluated state. + void Evaluate(int iter1, int iter2, bool precondition = true, + bool always_fix_m1_gauge = true) { + ++force_eval_count_; bool need_restart = false; std::string error_message; const vmecpp::VmecCheckpoint checkpoint = @@ -254,7 +259,7 @@ class VmecModel { *vmec_->decomposed_f_[0], *vmec_->physical_f_[0], need_restart, last_preconditioner_update_, last_full_update_nestor_, vmec_->fc_, iter1, iter2, checkpoint, checkpoint_after, - /*verbose=*/false); + /*verbose=*/false, always_fix_m1_gauge); if (!s.ok()) { error_message = std::string(s.status().message()); } @@ -266,6 +271,12 @@ class VmecModel { } bool need_restart() const { return last_need_restart_; } + // Total forward-model (force) evaluations since construction or the last + // reset. Counts every Evaluate, including those inside hessian_vector_product + // and preconditioner assembly, for a fair cross-optimizer cost comparison. + long force_eval_count() const { return force_eval_count_; } + void reset_force_eval_count() { force_eval_count_ = 0; } + // The Garabedian-style time step (PerformTimeStep): for each Fourier // coefficient, v = velocity_scale*(conjugation*v + dt*force); x += dt*v. void PerformTimeStep(double velocity_scale, double conjugation_parameter, @@ -394,6 +405,55 @@ class VmecModel { return FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_); } + // Hessian-vector product of VMEC's augmented functional, computed inside + // VMEC++ by a central directional derivative of the analytic force (which is + // the gradient): H v = (F(x + eps v) - F(x - eps v)) / (2 eps), in the + // decomposed internal basis. This is the matrix-free Hessian information an + // internal or external Newton-Krylov solver needs; F itself is exact, so only + // the directional step is finite-differenced. The current state is restored. + Eigen::VectorXd HessianVectorProduct(const Eigen::VectorXd &v, + double eps_rel = 1e-7) { + const Eigen::VectorXd x = + FlattenActive(*vmec_->decomposed_x_[0], vmec_->s_); + const double vnorm = v.norm(); + if (vnorm == 0.0) { + return Eigen::VectorXd::Zero(x.size()); + } + const double eps = eps_rel * (1.0 + x.norm()) / vnorm; + UnflattenActive(*vmec_->decomposed_x_[0], vmec_->s_, x + eps * v); + Evaluate(2, 2, /*precondition=*/false); + const Eigen::VectorXd fp = + FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_); + UnflattenActive(*vmec_->decomposed_x_[0], vmec_->s_, x - eps * v); + Evaluate(2, 2, /*precondition=*/false); + const Eigen::VectorXd fm = + FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_); + UnflattenActive(*vmec_->decomposed_x_[0], vmec_->s_, x); + return (fp - fm) / (2.0 * eps); + } + + // Apply VMEC's preconditioner M^-1 to a vector in the decomposed internal + // basis, mirroring the native apply sequence (m=1, radial, lambda). This is + // VMEC's hand-built approximate inverse Hessian; gradient-based solvers use + // it as the metric (preconditioned Krylov / quasi-Newton, and as the + // preconditioner for the Hessian solve in adjoint sensitivities). + // + // Requires a prior evaluate(precondition=true) at the current state: the + // radial preconditioner is assembled inside that forward-model call. + Eigen::VectorXd ApplyPreconditioner(const Eigen::VectorXd &v) const { + vmecpp::FourierForces tmp(&vmec_->s_, vmec_->r_[0].get(), vmec_->fc_.ns); + tmp.setZero(); + UnflattenActive(tmp, vmec_->s_, v); + vmecpp::IdealMhdModel &model = *vmec_->m_[0]; + model.applyM1Preconditioner(tmp); + const absl::Status status = model.applyRZPreconditioner(tmp); + if (!status.ok()) { + throw std::runtime_error(std::string(status.message())); + } + model.applyLambdaPreconditioner(tmp); + return FlattenActive(tmp, vmec_->s_); + } + // Residuals (set by Evaluate()): invariant {fsqr,fsqz,fsql} and // preconditioned {fsqr1,fsqz1,fsql1}. double fsqr() const { return vmec_->fc_.fsqr; } @@ -466,6 +526,7 @@ class VmecModel { int last_preconditioner_update_ = 0; int last_full_update_nestor_ = 0; bool last_need_restart_ = false; + long force_eval_count_ = 0; }; } // anonymous namespace @@ -1203,7 +1264,8 @@ PYBIND11_MODULE(_vmecpp, m) { .def_static("create", &VmecModel::Create, py::arg("indata"), py::arg("ns"), py::arg("initial_state") = std::nullopt) .def("evaluate", &VmecModel::Evaluate, py::arg("iter1"), py::arg("iter2"), - py::arg("precondition") = true) + py::arg("precondition") = true, + py::arg("always_fix_m1_gauge") = true) .def_property_readonly("need_restart", &VmecModel::need_restart) .def("perform_time_step", &VmecModel::PerformTimeStep, py::arg("velocity_scale"), py::arg("conjugation_parameter"), @@ -1219,6 +1281,12 @@ PYBIND11_MODULE(_vmecpp, m) { .def("get_state", &VmecModel::GetState) .def("set_state", &VmecModel::SetState, py::arg("state")) .def("get_forces", &VmecModel::GetForces) + .def("apply_preconditioner", &VmecModel::ApplyPreconditioner, + py::arg("v")) + .def("hessian_vector_product", &VmecModel::HessianVectorProduct, + py::arg("v"), py::arg("eps_rel") = 1e-7) + .def_property_readonly("force_eval_count", &VmecModel::force_eval_count) + .def("reset_force_eval_count", &VmecModel::reset_force_eval_count) .def_property_readonly("fsqr", &VmecModel::fsqr) .def_property_readonly("fsqz", &VmecModel::fsqz) .def_property_readonly("fsql", &VmecModel::fsql) diff --git a/src/vmecpp/simsopt_compat.py b/src/vmecpp/simsopt_compat.py index cddc26ab4..9778c8313 100644 --- a/src/vmecpp/simsopt_compat.py +++ b/src/vmecpp/simsopt_compat.py @@ -146,8 +146,9 @@ def __init__( # object, but the mpol/ntor values of either the vmec object # or the boundary surface object can be changed independently # by the user. + mpol, ntor = self._last_mpol_ntor(self.indata) mpol_for_surfacerzfourier, ntor_for_surfacerzfourier = ( - self._surface_rzfourier_resolution(self.indata.mpol, self.indata.ntor) + self._surface_rzfourier_resolution(mpol, ntor) ) self._boundary = SurfaceRZFourier.from_nphi_ntheta( nfp=self.indata.nfp, @@ -162,8 +163,8 @@ def __init__( # Transfer boundary shape data from indata to _boundary: vi = self.indata - for m in range(vi.mpol): - for n in range(2 * vi.ntor + 1): + for m in range(mpol): + for n in range(2 * ntor + 1): self._boundary.rc[m, n] = vi.rbc[m, n] self._boundary.zs[m, n] = vi.zbs[m, n] if vi.lasym: @@ -435,6 +436,18 @@ def boundary(self, boundary: SurfaceRZFourier) -> None: self.append_parent(boundary) self.need_to_run_code = True + @staticmethod + def _last_mpol_ntor(indata: vmecpp.VmecInput) -> tuple[int, int]: + """(indata.mpol, indata.ntor) as plain ints. + + VmecInput.mpol/.ntor may also be a Fourier-resolution continuation schedule (a + sequence); SIMSOPT's SurfaceRZFourier only has a single resolution, so the last + (finest, target) entry of the schedule is used. + """ + mpol = indata.mpol if isinstance(indata.mpol, int) else int(indata.mpol[-1]) + ntor = indata.ntor if isinstance(indata.ntor, int) else int(indata.ntor[-1]) + return mpol, ntor + @staticmethod def _surface_rzfourier_resolution(mpol: int, ntor: int) -> tuple[int, int]: # SurfaceRZFourier uses m up to mpol inclusive, unlike VMEC++. @@ -462,9 +475,8 @@ def set_indata(self) -> None: raise RuntimeError(msg) assert self.indata is not None vi = self.indata # Shorthand - target_mpol, target_ntor = self._surface_rzfourier_resolution( - self.indata.mpol, self.indata.ntor - ) + mpol, ntor = self._last_mpol_ntor(self.indata) + target_mpol, target_ntor = self._surface_rzfourier_resolution(mpol, ntor) boundary_RZFourier = self._resize_surface_rzfourier( self.boundary.to_RZFourier().copy(), target_mpol, @@ -483,8 +495,7 @@ def set_indata(self) -> None: zbc.fill(0.0) # Transfer boundary shape data from the surface object to VMEC: - ntor = self.indata.ntor - for m in range(self.indata.mpol): + for m in range(mpol): for n in range(2 * ntor + 1): vi.rbc[m, n] = boundary_RZFourier.get_rc(m, n - ntor) vi.zbs[m, n] = boundary_RZFourier.get_zs(m, n - ntor) diff --git a/tests/test_adjoint.py b/tests/test_adjoint.py new file mode 100644 index 000000000..54a45c1a1 --- /dev/null +++ b/tests/test_adjoint.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""The adjoint boundary gradient matches brute-force finite differences. + +dJ/d(boundary) from one Hessian solve (implicit-function adjoint) agrees with the +reference gradient obtained by re-converging the interior equilibrium for each perturbed +boundary degree of freedom. J here is the MHD energy of the converged equilibrium. +""" + +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from vmecpp_adjoint import ( # type: ignore + boundary_gradient, + finite_difference_boundary_gradient, + make_model, + mhd_energy, + partition, +) + + +def test_adjoint_matches_finite_difference(): + ns = 11 + model = make_model(ns=ns) + model.solve() + x_star = np.asarray(model.get_state(), float).copy() + interior, boundary = partition(model, ns) + + g_adjoint = boundary_gradient(model, x_star, interior, boundary, mhd_energy) + + # Reference on a representative subset (each requires interior re-solves). + dofs = [0, 2, 9] + g_fd = finite_difference_boundary_gradient( + model, x_star, interior, boundary, mhd_energy, dofs + ) + + scale = max(np.linalg.norm(g_adjoint), 1e-30) + for j in dofs: + assert abs(g_adjoint[j] - g_fd[j]) < 1e-3 * scale + + +if __name__ == "__main__": + import pytest + + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_continuation.py b/tests/test_continuation.py index 41dbf20a0..2dbc4a3c3 100644 --- a/tests/test_continuation.py +++ b/tests/test_continuation.py @@ -125,42 +125,55 @@ def test_interpolate_fourier_pad_and_truncate(solovev_output: vmecpp.VmecOutput) # --- continuation driver ----------------------------------------------------- -def test_single_step_continuation_matches_direct_run(cma_input: vmecpp.VmecInput): - """A single-resolution continuation invokes exactly one solve and must be bit-for- - bit identical to calling ``run`` on the same single-grid input.""" - ns_final = int(np.asarray(cma_input.ns_array)[-1]) - - single_grid = cma_input.model_copy(deep=True) - single_grid.ns_array = np.asarray([ns_final], dtype=np.int64) - single_grid.ftol_array = np.asarray([1e-12], dtype=float) - single_grid.niter_array = np.asarray([60000], dtype=np.int64) - direct = vmecpp.run(single_grid, verbose=False, max_threads=1) - - continued = vmecpp.run_continuation( - cma_input, - ns_array=[ns_final], - ftol_array=[1e-12], - niter_array=[60000], - verbose=False, - max_threads=1, +def test_mpol_ntor_length_one_sequence_collapses_to_scalar( + cma_input: vmecpp.VmecInput, +): + """A length-1 mpol/ntor sequence is equivalent to a scalar and is stored as a plain + int, so it takes the direct (non-continuation) code path. + + This coercion runs during validation (construction / model_validate), like all of + VmecInput's other validators; it goes through vmecpp.VmecInput.model_validate here + rather than a bare attribute assignment for that reason. + """ + data = cma_input.model_dump(mode="json") + data["mpol"] = [int(cma_input.mpol)] + data["ntor"] = [int(cma_input.ntor)] + single_step = vmecpp.VmecInput.model_validate(data) + assert isinstance(single_step.mpol, int) + assert isinstance(single_step.ntor, int) + assert single_step.mpol == cma_input.mpol + assert single_step.ntor == cma_input.ntor + + +def test_mpol_length_mismatch_raises(cma_input: vmecpp.VmecInput): + """A mpol/ntor schedule that doesn't match ns_array's length is rejected with a + clear error, rather than silently misaligning steps.""" + mismatched = cma_input.model_copy(deep=True) + mismatched.ns_array = np.asarray([15, 31], dtype=np.int64) + mismatched.ftol_array = np.asarray([1e-8, 1e-10], dtype=float) + mismatched.niter_array = np.asarray([100, 100], dtype=np.int64) + mismatched.mpol = np.asarray([3, 4, 5], dtype=np.int64) # 3 entries, ns_array has 2 + + with pytest.raises(ValueError, match="mpol"): + vmecpp.run(mismatched, verbose=False, max_threads=1) + + +def test_ns_only_continuation_reproduces_direct_multigrid( + cma_input: vmecpp.VmecInput, cma_direct: vmecpp.VmecOutput +): + """A continuation schedule with a constant (array-valued) mpol/ntor reaches the same + equilibrium as the C++ multi-grid: identical resolution, a converged force balance, + a bit-identical plasma volume, and geometry agreeing at the convergence level.""" + n_steps = len(np.asarray(cma_input.ns_array)) + continuation_input = cma_input.model_copy(deep=True) + continuation_input.mpol = np.asarray( + [int(cma_input.mpol)] * n_steps, dtype=np.int64 ) - for field in ("rmnc", "zmns", "lmns_full"): - np.testing.assert_array_equal( - np.asarray(getattr(continued.wout, field)), - np.asarray(getattr(direct.wout, field)), - ) - np.testing.assert_array_equal( - np.asarray(continued.wout.iotaf), np.asarray(direct.wout.iotaf) + continuation_input.ntor = np.asarray( + [int(cma_input.ntor)] * n_steps, dtype=np.int64 ) - -def test_run_continuation_reproduces_direct( - cma_input: vmecpp.VmecInput, cma_direct: vmecpp.VmecOutput -): - """The default Python continuation reaches the same equilibrium as the C++ - multi-grid: identical resolution, a converged force balance, a bit-identical - plasma volume, and geometry agreeing at the convergence level.""" - continued = vmecpp.run_continuation(cma_input, verbose=False, max_threads=1) + continued = vmecpp.run(continuation_input, verbose=False, max_threads=1) assert int(continued.wout.ns) == int(cma_direct.wout.ns) # Converged to a force residual comparable to the direct solve. @@ -184,14 +197,23 @@ def test_continuation_agreement_tightens_with_ftol(cma_input: vmecpp.VmecInput): """As the force-balance tolerance tightens, the continuation and the direct multi- grid converge toward the same geometry -- the signature of a shared equilibrium rather than a defect in the continuation.""" + n_steps = len(np.asarray(cma_input.ns_array)) def max_geometry_diff(ftol_array: list[float]) -> float: reference = cma_input.model_copy(deep=True) reference.ftol_array = np.asarray(ftol_array, dtype=float) direct = vmecpp.run(reference, verbose=False, max_threads=1) - continued = vmecpp.run_continuation( - cma_input, ftol_array=ftol_array, verbose=False, max_threads=1 + + continuation_input = cma_input.model_copy(deep=True) + continuation_input.mpol = np.asarray( + [int(cma_input.mpol)] * n_steps, dtype=np.int64 ) + continuation_input.ntor = np.asarray( + [int(cma_input.ntor)] * n_steps, dtype=np.int64 + ) + continuation_input.ftol_array = np.asarray(ftol_array, dtype=float) + continued = vmecpp.run(continuation_input, verbose=False, max_threads=1) + return float( np.max( np.abs(np.asarray(direct.wout.rmnc) - np.asarray(continued.wout.rmnc)) @@ -214,16 +236,18 @@ def test_fourier_continuation_converges( ftol_final = float(np.asarray(cma_input.ftol_array)[-1]) niter_final = int(np.asarray(cma_input.niter_array)[-1]) - continued = vmecpp.run_continuation( - cma_input, - ns_array=[ns_final, ns_final], - mpol_array=[max(2, mpol_final - 2), mpol_final], - ntor_array=[ntor, ntor], - ftol_array=[ftol_final, ftol_final], - niter_array=[niter_final, niter_final], - verbose=False, - max_threads=1, + continuation_input = cma_input.model_copy(deep=True) + continuation_input.ns_array = np.asarray([ns_final, ns_final], dtype=np.int64) + continuation_input.mpol = np.asarray( + [max(2, mpol_final - 2), mpol_final], dtype=np.int64 ) + continuation_input.ntor = ntor # scalar broadcasts to every step + continuation_input.ftol_array = np.asarray([ftol_final, ftol_final], dtype=float) + continuation_input.niter_array = np.asarray( + [niter_final, niter_final], dtype=np.int64 + ) + + continued = vmecpp.run(continuation_input, verbose=False, max_threads=1) assert int(continued.wout.mpol) == mpol_final assert _final_force_residual(continued) < 1e-5 @@ -236,3 +260,36 @@ def test_fourier_continuation_converges( atol=4e-3, rtol=1e-2, ) + + +def test_fourier_continuation_hot_restarts_first_step_from_restart_from( + cma_input: vmecpp.VmecInput, +): + """restart_from seeds the first continuation step (interpolated to its resolution) + instead of a cold start, when input.mpol/.ntor is a sequence.""" + ns_final = int(np.asarray(cma_input.ns_array)[-1]) + mpol_final = int(cma_input.mpol) + ntor = int(cma_input.ntor) + + warm_start_input = cma_input.model_copy(deep=True) + warm_start_input.ns_array = np.asarray([ns_final], dtype=np.int64) + warm_start = vmecpp.run(warm_start_input, verbose=False, max_threads=1) + + continuation_input = cma_input.model_copy(deep=True) + continuation_input.ns_array = np.asarray([ns_final, ns_final], dtype=np.int64) + continuation_input.mpol = np.asarray( + [max(2, mpol_final - 2), mpol_final], dtype=np.int64 + ) + continuation_input.ntor = ntor + continuation_input.ftol_array = np.asarray([1e-8, 1e-10], dtype=float) + continuation_input.niter_array = np.asarray([2000, 2000], dtype=np.int64) + + continued = vmecpp.run( + continuation_input, + verbose=False, + max_threads=1, + restart_from=warm_start, + ) + + assert int(continued.wout.mpol) == mpol_final + assert _final_force_residual(continued) < 1e-5 diff --git a/tests/test_external_optimizers.py b/tests/test_external_optimizers.py new file mode 100644 index 000000000..5e20b7fa0 --- /dev/null +++ b/tests/test_external_optimizers.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""External optimizers reach force balance on axisymmetric and 3D cases. + +The Solov'ev equilibrium has a reproducible internal state, so the 2D test compares the +full state vector with the native solver. In 3D, the poloidal parameterization is not +unique and spectral condensation only regularizes that coordinate freedom. The 3D test +therefore compares the residual and energy instead of coordinate-dependent Fourier +coefficients. +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from external_optimizers import ( # type: ignore + make_model, + reference_equilibrium, + residual, + solve_newton_hvp, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_preconditioned_descent, +) + +ROOT = Path(__file__).resolve().parents[1] +SOLOVEV = ROOT / "examples" / "data" / "solovev.json" +CTH_LIKE = ROOT / "examples" / "data" / "cth_like_fixed_bdy.json" +CMA = ROOT / "src" / "vmecpp" / "cpp" / "vmecpp" / "test_data" / "cma.json" + +SOLVERS = ( + solve_preconditioned_descent, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_newton_hvp, +) + + +def _solve_all(input_path): + return {solver.__name__: solver(input_path, ns=11) for solver in SOLVERS} + + +@pytest.fixture(scope="module") +def reference_2d(): + return reference_equilibrium(SOLOVEV, ns=11) + + +@pytest.fixture(scope="module") +def reference_3d(): + return reference_equilibrium(CTH_LIKE, ns=11) + + +@pytest.fixture(scope="module") +def solutions_2d(): + return _solve_all(SOLOVEV) + + +@pytest.fixture(scope="module") +def solutions_3d(): + return _solve_all(CTH_LIKE) + + +def test_optimizers_reach_2d_equilibrium(reference_2d, solutions_2d): + x_star, w_star = reference_2d + for name, (x, result) in solutions_2d.items(): + assert result.residual_norm < 1e-7, name + assert abs(result.energy - w_star) < 1e-8, name + assert np.linalg.norm(x - x_star) < 1e-5, name + + +def test_optimizers_reach_3d_force_balance(reference_3d, solutions_3d): + _, w_star = reference_3d + for name, (_, result) in solutions_3d.items(): + assert result.residual_norm < 1e-7, name + assert np.isclose(result.energy, w_star, rtol=1e-4, atol=0.0), name + + +def test_preconditioner_reduces_newton_krylov_force_evaluations(solutions_3d): + plain = solutions_3d[solve_newton_krylov.__name__][1] + preconditioned = solutions_3d[solve_newton_krylov_preconditioned.__name__][1] + assert preconditioned.force_evals < plain.force_evals + + +def test_hvp_newton_uses_fewer_outer_iterations_than_descent(solutions_3d): + newton = solutions_3d[solve_newton_hvp.__name__][1] + descent = solutions_3d[solve_preconditioned_descent.__name__][1] + assert newton.outer_iters < 20 + assert newton.outer_iters < descent.outer_iters + + +def test_cma_cold_start_exercises_non_axisymmetric_paths(): + model = make_model(CMA, ns=25) + x0 = np.asarray(model.get_state(), float) + f0 = residual(model)(x0) + assert np.all(np.isfinite(f0)) + assert np.linalg.norm(f0) > 0.0 + + rng = np.random.default_rng(0) + v = rng.standard_normal(x0.size) + hv = np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float) + assert np.all(np.isfinite(hv)) + assert np.linalg.norm(hv) > 0.0 diff --git a/tests/test_hessian.py b/tests/test_hessian.py new file mode 100644 index 000000000..8034f5ace --- /dev/null +++ b/tests/test_hessian.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VmecModel.hessian_vector_product gives the augmented functional's curvature. + +The Hessian-vector product is a central directional derivative of the analytic +force (the gradient of VMEC's augmented functional), computed inside VMEC++: +H v = (F(x + eps v) - F(x - eps v)) / (2 eps). It must be linear in v and agree +with an independent finite difference of the force, and it restores the state. +""" + +from pathlib import Path + +import numpy as np + +from vmecpp.cpp import _vmecpp # type: ignore + +SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" + + +def _model(ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(SOLOVEV)), ns) + + +def _raw_force(model, x): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), float) + + +def test_hvp_matches_finite_difference(): + m = _model() + m.evaluate(2, 2, False) + x = np.asarray(m.get_state(), float) + rng = np.random.default_rng(0) + v = rng.standard_normal(x.size) + v /= np.linalg.norm(v) + + hv = np.asarray(m.hessian_vector_product(np.ascontiguousarray(v)), float) + + eps = 1e-6 + fd = (_raw_force(m, x + eps * v) - _raw_force(m, x - eps * v)) / (2 * eps) + assert np.linalg.norm(hv - fd) < 1e-5 * np.linalg.norm(fd) + + +def test_hvp_is_linear(): + m = _model() + m.evaluate(2, 2, False) + rng = np.random.default_rng(1) + v = rng.standard_normal(np.asarray(m.get_state()).size) + hv = np.asarray(m.hessian_vector_product(np.ascontiguousarray(v)), float) + hv2 = np.asarray(m.hessian_vector_product(np.ascontiguousarray(2.0 * v)), float) + assert np.linalg.norm(hv2 - 2.0 * hv) < 1e-9 * np.linalg.norm(hv) + + +def test_hvp_restores_state(): + m = _model() + m.evaluate(2, 2, False) + x0 = np.asarray(m.get_state(), float).copy() + rng = np.random.default_rng(2) + v = rng.standard_normal(x0.size) + m.hessian_vector_product(np.ascontiguousarray(v)) + assert np.allclose(np.asarray(m.get_state(), float), x0) diff --git a/tests/test_internal_gradient.py b/tests/test_internal_gradient.py index 17d36f65f..16d26825e 100644 --- a/tests/test_internal_gradient.py +++ b/tests/test_internal_gradient.py @@ -17,10 +17,17 @@ from pathlib import Path import numpy as np +import pytest from vmecpp.cpp import _vmecpp # type: ignore SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" +CTH_LIKE = ( + Path(__file__).resolve().parents[1] + / "examples" + / "data" + / "cth_like_fixed_bdy.json" +) def _model(ns: int = 11): @@ -63,3 +70,38 @@ def test_cold_start_is_excluded(): m = _model() m.evaluate(2, 2, False) assert np.all(np.isfinite(np.asarray(m.get_forces(), float))) + + +@pytest.fixture(scope="module") +def force_histories(): + indata = _vmecpp.VmecINDATA.from_file(str(CTH_LIKE)) + reference = _vmecpp.VmecModel.create(indata, 11) + reference.solve() + low_residual_state = np.asarray(reference.get_state(), float).copy() + + model = _vmecpp.VmecModel.create(indata, 11) + high_residual_state = np.asarray(model.get_state(), float).copy() + model.evaluate(2, 2, False) + high_after_high = np.asarray(model.get_forces(), float).copy() + + model.set_state(np.ascontiguousarray(low_residual_state)) + model.evaluate(2, 2, False) + low_after_high = np.asarray(model.get_forces(), float).copy() + model.evaluate(2, 2, False) + low_after_low = np.asarray(model.get_forces(), float).copy() + + model.set_state(np.ascontiguousarray(high_residual_state)) + model.evaluate(2, 2, False) + high_after_low = np.asarray(model.get_forces(), float).copy() + + return high_after_high, high_after_low, low_after_high, low_after_low + + +def test_raw_force_is_independent_of_high_residual_history(force_histories): + _, _, low_after_high, low_after_low = force_histories + np.testing.assert_array_equal(low_after_high, low_after_low) + + +def test_raw_force_is_independent_of_low_residual_history(force_histories): + high_after_high, high_after_low, _, _ = force_histories + np.testing.assert_array_equal(high_after_high, high_after_low) diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py new file mode 100644 index 000000000..7009dee8a --- /dev/null +++ b/tests/test_preconditioner.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VmecModel.apply_preconditioner exposes VMEC's preconditioner as an operator. + +The preconditioner M^-1 is VMEC's hand-built approximate inverse Hessian. The native +solver applies it to the raw force to get its search direction, so +apply_preconditioner(raw force) must equal the preconditioned force exactly. The +operator is linear and, once assembled (via evaluate(precondition=True)), does not +depend on the current state, so it can be reused as a frozen preconditioner for +Krylov/quasi-Newton solvers. +""" + +from pathlib import Path + +import numpy as np + +try: + from vmecpp.cpp import _vmecpp +except ImportError: + import _vmecpp + +SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" + + +def _model(ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(SOLOVEV)), ns) + + +def test_preconditioner_matches_native_search_direction(): + m = _model() + m.evaluate(2, 2, True) # assemble preconditioner + preconditioned force + f_prec = np.asarray(m.get_forces(), float) + m.evaluate(2, 2, False) # raw force (does not reassemble) + f_raw = np.asarray(m.get_forces(), float) + minv_fraw = np.asarray(m.apply_preconditioner(f_raw), float) + assert np.linalg.norm(minv_fraw - f_prec) <= 1e-12 * np.linalg.norm(f_prec) + + +def test_preconditioner_is_linear_and_finite(): + m = _model() + m.evaluate(2, 2, True) + rng = np.random.default_rng(0) + v = np.ascontiguousarray(rng.standard_normal(np.asarray(m.get_state()).size)) + mv = np.asarray(m.apply_preconditioner(v), float) + m2v = np.asarray(m.apply_preconditioner(np.ascontiguousarray(2.0 * v)), float) + assert np.all(np.isfinite(mv)) + assert np.linalg.norm(m2v - 2.0 * mv) <= 1e-12 * np.linalg.norm(mv) + + +def test_preconditioner_state_invariant_after_assembly(): + m = _model() + m.evaluate(2, 2, True) + rng = np.random.default_rng(1) + x = np.asarray(m.get_state(), float) + v = np.ascontiguousarray(rng.standard_normal(x.size)) + mv0 = np.asarray(m.apply_preconditioner(v), float) + # Move to a different state and raw-evaluate (no reassembly). + m.set_state(np.ascontiguousarray(x + 0.01 * rng.standard_normal(x.size))) + m.evaluate(2, 2, False) + mv1 = np.asarray(m.apply_preconditioner(v), float) + assert np.linalg.norm(mv1 - mv0) <= 1e-12 * np.linalg.norm(mv0) diff --git a/tests/test_simsopt_gradient.py b/tests/test_simsopt_gradient.py new file mode 100644 index 000000000..90f4f97c5 --- /dev/null +++ b/tests/test_simsopt_gradient.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VMEC++ exposes an analytic boundary gradient to SIMSOPT. + +The VmecEnergy Optimizable's analytic dJ (the implicit-function adjoint) matches finite +differences of its objective, and computing it is much cheaper than the conventional +finite-difference boundary gradient (which re-solves the equilibrium per boundary degree +of freedom). +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) + +pytest.importorskip("simsopt") + +from simsopt_vmec_gradient import ( # type: ignore + VmecBoundaryProblem, + gradient_cost, + make_simsopt_optimizable, +) + + +def test_simsopt_optimizable_gradient_matches_fd(): + problem = VmecBoundaryProblem(ns=11) + opt = make_simsopt_optimizable(problem) + g = np.asarray(opt.dJ(), float) + + p0 = np.asarray(opt.local_full_x, float) + h = 1e-5 + scale = max(np.linalg.norm(g), 1e-30) + for j in (0, 2, 9): + pp = p0.copy() + pp[j] += h + opt.local_full_x = pp + jp = opt.J() + pm = p0.copy() + pm[j] -= h + opt.local_full_x = pm + jm = opt.J() + opt.local_full_x = p0 + assert abs(g[j] - (jp - jm) / (2 * h)) < 1e-3 * scale + + +def test_adjoint_gradient_cheaper_than_finite_difference(): + analytic = gradient_cost(analytic=True) + fd = gradient_cost(analytic=False) + # Same gradient, far fewer force evaluations (advantage grows with #DOFs). + rel = np.linalg.norm(analytic.gradient - fd.gradient) / np.linalg.norm(fd.gradient) + assert rel < 1e-2 + assert analytic.force_evals < fd.force_evals