diff --git a/.github/actions/run-benchmarks/action.yml b/.github/actions/run-benchmarks/action.yml index a8ddf671a..2d24521f3 100644 --- a/.github/actions/run-benchmarks/action.yml +++ b/.github/actions/run-benchmarks/action.yml @@ -53,6 +53,14 @@ runs: export OMP_WAIT_POLICY=passive pytest benchmarks/test_benchmarks.py --benchmark-json=benchmark_results.json + # github-action-benchmark's "pytest" tool always plots pytest-benchmark's + # `ops` field as iter/sec, but the C++ suite emits seconds. + # We convert both to the same format, so plotting is the same. + - name: Convert Python benchmark result to seconds + shell: bash + run: | + python3 benchmarks/to_seconds_json.py pytest benchmark_results.json benchmark_results_seconds.json + - name: Store / compare Python benchmark result # Skip the result upload/compare for fork PRs: their GITHUB_TOKEN is # read-only, so comment-on-alert/auto-push hit 'Resource not accessible @@ -61,8 +69,8 @@ runs: if: ${{ inputs.auto-push == 'true' || github.event.pull_request.head.repo.full_name == github.repository }} uses: benchmark-action/github-action-benchmark@v1.21.0 with: - tool: "pytest" - output-file-path: benchmark_results.json + tool: "customSmallerIsBetter" + output-file-path: benchmark_results_seconds.json gh-pages-branch: benchmark-runs benchmark-data-dir-path: benchmarks alert-threshold: "150%" @@ -76,7 +84,7 @@ runs: # Builds the four benchmark cc_binary targets, runs each with # --benchmark_format=json, and merges their output into a single # cpp_benchmark_results.json (schema: {"context": ..., "benchmarks": [...]}) - # for the store step below. Entries with "error_occurred" (e.g. an + # for the conversion step below. Entries with "error_occurred" (e.g. an # FFT-path benchmark skipped because no FFTX codelet exists for that # resolution) are dropped: they carry no valid timing to track. @@ -142,6 +150,14 @@ runs: f"({skipped} skipped/errored).") PY + # github-action-benchmark's "pytest" tool always plots pytest-benchmark's + # `ops` field as iter/sec, but the C++ suite emits seconds. + # We convert both to the same format, so plotting is the same. + - name: Convert C++ benchmark result to seconds + shell: bash + run: | + python3 benchmarks/to_seconds_json.py googlecpp cpp_benchmark_results.json cpp_benchmark_results_seconds.json + - name: Store / compare C++ benchmark result # Same fork-PR skip rationale as the Python store step above. # @@ -160,8 +176,8 @@ runs: uses: benchmark-action/github-action-benchmark@v1.21.0 with: name: "C++ Microbenchmarks" - tool: "googlecpp" - output-file-path: cpp_benchmark_results.json + tool: "customSmallerIsBetter" + output-file-path: cpp_benchmark_results_seconds.json gh-pages-branch: benchmark-runs benchmark-data-dir-path: benchmarks skip-fetch-gh-pages: true diff --git a/.github/workflows/full_validation.yaml b/.github/workflows/full_validation.yaml new file mode 100644 index 000000000..c378dc018 --- /dev/null +++ b/.github/workflows/full_validation.yaml @@ -0,0 +1,70 @@ +name: Full V&V against reference VMEC + +# Runs the full Verification & Validation suite from proximafusion/vmecpp-validation +# (all ~219 input configurations against reference Fortran VMEC2000), using the +# vmecpp version built from this commit rather than the version pinned on PyPI/GitHub. +# This is a lot more exhaustive (and slower) than the "short" validation that +# vmecpp-validation itself runs on its own PRs, so it is not run on every push here. +on: + workflow_dispatch: + push: + branches: + - main + +concurrency: + group: ${{ github.workflow }}-${{ github.ref || github.run_id }} + cancel-in-progress: ${{ github.ref_name != 'main' }} + +jobs: + full-validation: + name: Run full VMEC++ validation + runs-on: ubuntu-22.04 + # The full parameter scan runs ~219 configurations, each computing a reference + # wout with Fortran VMEC2000 in Docker plus one with VMEC++: budget generously. + timeout-minutes: 360 + steps: + - name: Check out VMEC++ + uses: actions/checkout@v4 + with: + path: vmecpp + lfs: true + + - name: Check out vmecpp-validation + uses: actions/checkout@v4 + with: + repository: proximafusion/vmecpp-validation + path: vmecpp-validation + lfs: true + + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install required system packages for Ubuntu + run: | + sudo apt-get update && sudo apt-get install -y \ + build-essential cmake libnetcdf-dev liblapack-dev liblapacke-dev libopenmpi-dev \ + libomp-dev libeigen3-dev nlohmann-json3-dev libhdf5-dev + + - name: Install vmecpp-validation's Python requirements + run: | + cd vmecpp-validation + # Install everything except vmecpp itself: we build and install vmecpp + # from this commit below instead of the version pinned in requirements.txt. + grep -v '^vmecpp@' requirements.txt > requirements.no-vmecpp.txt + python -m pip install -r requirements.no-vmecpp.txt + + - name: Install VMEC++ from this commit + run: python -m pip install ./vmecpp + + - name: Run full validation + working-directory: vmecpp-validation + run: python validate_vmec.py + + - name: Upload V&V results + if: always() + uses: actions/upload-artifact@v4 + with: + name: vnvresults + path: vmecpp-validation/vnvresults_*/ + retention-days: 30 diff --git a/README.md b/README.md index fa0800d5a..18cc7eefe 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,7 @@ [![CI](https://github.com/proximafusion/vmecpp/actions/workflows/tests.yaml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/tests.yaml) [![C++ core tests](https://github.com/proximafusion/vmecpp/actions/workflows/test_bazel.yaml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/test_bazel.yaml) +[![Full V&V against reference VMEC](https://github.com/proximafusion/vmecpp/actions/workflows/full_validation.yaml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/full_validation.yaml) [![Publish wheels to PyPI](https://github.com/proximafusion/vmecpp/actions/workflows/pypi_publish.yml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/pypi_publish.yml) VMEC++ is a Python-friendly, from-scratch reimplementation in C++ of the Variational Moments Equilibrium Code (VMEC), @@ -315,6 +316,10 @@ The single-thread runtimes as well as the contents of the "wout" file produced b The full validation test can be found at https://github.com/proximafusion/vmecpp-validation, including a set of sensible input configurations, parameter scan values and tolerances that make the comparison pass. See that repo for more information. +This full validation (~219 input configurations) is run against every commit to `main` and can also be triggered +on demand from the [Full V&V against reference VMEC](https://github.com/proximafusion/vmecpp/actions/workflows/full_validation.yaml) workflow +(click "Run workflow"). It builds `vmecpp` from the corresponding commit rather than using the version pinned by `vmecpp-validation`. + ## Differences with respect to PARVMEC/VMEC2000 VMEC++: diff --git a/benchmarks/to_seconds_json.py b/benchmarks/to_seconds_json.py new file mode 100644 index 000000000..15a72b0ed --- /dev/null +++ b/benchmarks/to_seconds_json.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# SPDX-License-Identifier: MIT +"""Convert pytest-benchmark or Google Benchmark JSON to seconds. + +github-action-benchmark's built-in "pytest" and "googlecpp" tool parsers plot whatever +value each framework happens to report natively -- pytest-benchmark's "iter/sec" and +Google Benchmark's raw per-iteration time in its own time_unit (commonly nanoseconds). +Charting those side by side, or even a single suite whose time_unit isn't fixed, +requires the reader (or the chart's JS) to know and convert between units. + +This script normalizes both to plain wall-clock seconds and emits the action's generic +"customSmallerIsBetter" schema, so the chart always plots seconds and the benchmarks.js +frontend needs no unit-detection logic. +""" + +import argparse +import json + + +def convert_pytest(data): + results = [] + for bench in data["benchmarks"]: + stats = bench["stats"] + results.append( + { + "name": bench["fullname"], + "unit": "seconds", + "value": stats["mean"], + "range": f"stddev: {stats['stddev']}", + "extra": f"rounds: {stats['rounds']}", + } + ) + return results + + +_GOOGLECPP_TIME_UNIT_TO_SECONDS = { + "s": 1.0, + "ms": 1e-3, + "us": 1e-6, + "ns": 1e-9, +} + + +def convert_googlecpp(data): + results = [] + for bench in data["benchmarks"]: + if bench.get("error_occurred"): + continue + factor = _GOOGLECPP_TIME_UNIT_TO_SECONDS[bench["time_unit"]] + results.append( + { + "name": bench["name"], + "unit": "seconds", + "value": bench["real_time"] * factor, + "extra": ( + f"iterations: {bench['iterations']}\n" + f"cpu: {bench['cpu_time'] * factor} seconds\n" + f"threads: {bench.get('threads', 1)}" + ), + } + ) + return results + + +_CONVERTERS = { + "pytest": convert_pytest, + "googlecpp": convert_googlecpp, +} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("format", choices=sorted(_CONVERTERS)) + parser.add_argument("input", help="Raw benchmark JSON produced by the tool.") + parser.add_argument("output", help="Where to write the customSmallerIsBetter JSON.") + args = parser.parse_args() + + with open(args.input) as f: + data = json.load(f) + + results = _CONVERTERS[args.format](data) + + with open(args.output, "w") as f: + json.dump(results, f, indent=2) + + +if __name__ == "__main__": + main() diff --git a/docs/_static/benchmarks.js b/docs/_static/benchmarks.js index 7ec772b2f..19754aa72 100644 --- a/docs/_static/benchmarks.js +++ b/docs/_static/benchmarks.js @@ -37,20 +37,6 @@ "#1f6feb", ]; - // Parse mean duration from the extra field (e.g. "mean: 339.88 msec\nrounds: 5") - function parseMeanSeconds(bench) { - if (!bench.extra) return null; - var match = bench.extra.match(/mean:\s*([\d.]+)\s*([\w]+)/); - if (!match) return null; - var value = parseFloat(match[1]); - var unit = match[2]; - if (unit === "msec") return value / 1000; - if (unit === "usec") return value / 1000000; - if (unit === "nsec") return value / 1000000000; - if (unit === "sec") return value; - return null; - } - function renderCharts(data) { var container = document.getElementById("benchmark-charts"); if (!container) return; @@ -130,8 +116,7 @@ { label: benchName, data: dataset.map(function (d) { - var secs = parseMeanSeconds(d.bench); - return secs !== null ? secs : 1.0 / d.bench.value; + return d.bench.value; }), borderColor: color, backgroundColor: color + "30", 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 index 7a8706c4f..b4da85f26 100644 --- a/examples/external_optimizers.py +++ b/examples/external_optimizers.py @@ -7,21 +7,32 @@ 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)`` (see -``tests/test_internal_gradient.py``). Finding the equilibrium is therefore the -root problem F(x) = 0, which any gradient/Hessian-based solver can attack while -reusing VMEC++'s forward model. - -This module wires that residual to two solvers and is used both by the benchmark -``main`` below and by ``tests/test_external_optimizers.py``: - -* native-style preconditioned descent (heavy-ball on the preconditioned search - direction, i.e. VMEC's own update), and -* Jacobian-free Newton-Krylov (matrix-free Hessian information). - -Both converge to the same equilibrium as the native solver. Quasi-Newton -root-finders without a preconditioner diverge on this stiff system, which is why -VMEC's preconditioner matters; exposing it as an operator is a follow-up. +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 @@ -32,11 +43,16 @@ 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" / "solovev.json" + Path(__file__).resolve().parents[1] + / "examples" + / "data" + / "cth_like_fixed_bdy.json" ) @@ -82,6 +98,7 @@ def F(x): class Result: name: str force_evals: int + outer_iters: int seconds: float residual_norm: float energy: float @@ -94,6 +111,42 @@ def reference_equilibrium(input_path: Path = DEFAULT_INPUT, ns: int = 11): 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 ): @@ -101,62 +154,127 @@ def solve_preconditioned_descent( F = residual(model) x = np.asarray(model.get_state(), float).copy() v = np.zeros_like(x) - evals = 0 + model.reset_force_eval_count() + it = 0 t0 = time.perf_counter() for _ in range(max_iter): if np.linalg.norm(F(x)) < tol: break - evals += 1 + 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 - model.set_state(np.ascontiguousarray(x)) - model.evaluate(2, 2, False) - return x, Result( - "preconditioned descent", - evals, - time.perf_counter() - t0, - float(np.linalg.norm(np.asarray(model.get_forces(), float))), - model.mhd_energy, - ) + return _finish(model, "preconditioned descent", x, it, t0) -def solve_newton_krylov(input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=200): +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) - n = [0] - - def counted(x): - n[0] += 1 - return F(x) - 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(counted, x0, f_tol=tol, maxiter=max_iter, method="lgmres") - model.set_state(np.ascontiguousarray(x)) - model.evaluate(2, 2, False) - return x, Result( - "Newton-Krylov (JFNK)", - n[0], - time.perf_counter() - t0, - float(np.linalg.norm(np.asarray(model.get_forces(), float))), - model.mhd_energy, + 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") - rows = [solve_preconditioned_descent()[1], solve_newton_krylov()[1]] print( - f"{'optimizer':28s} {'F-evals':>8s} {'time[s]':>8s} " + f"{'optimizer':32s} {'F-evals':>8s} {'iters':>6s} {'time[s]':>8s} " f"{'||F||':>10s} {'dW vs ref':>10s}" ) - for r in rows: + for solver in ALL_SOLVERS: + r = solver()[1] print( - f"{r.name:28s} {r.force_evals:8d} {r.seconds:8.2f} " + 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}" ) 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/src/vmecpp/__init__.py b/src/vmecpp/__init__.py index 54ffb4077..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)), @@ -113,6 +139,16 @@ class FreeBoundaryMethod(str, enum.Enum): """Boundary Integral Equation Solver for Toroidal systems.""" +class IterationStyle(str, enum.Enum): + """Time-step / restart control scheme for the equilibrium iteration.""" + + VMEC_8_52 = "vmec_8_52" + """The Fortran VMEC 8.52 control (the default).""" + + PARVMEC = "parvmec" + """The PARVMEC / VMEC2000 9.0 control.""" + + class OutputMode(enum.Enum): """Controls the output format of iteration logging..""" @@ -138,6 +174,15 @@ def _validate_free_boundary_method( return FreeBoundaryMethod(str(value)) +def _validate_iteration_style( + value: _vmecpp.IterationStyle | str | IterationStyle, +) -> IterationStyle: + """Convert various representations to IterationStyle.""" + if isinstance(value, _vmecpp.IterationStyle): + return IterationStyle(value.name.lower()) # pyright: ignore[reportAttributeAccessIssue] + return IterationStyle(str(value)) + + # This is a pure Python equivalent of VmecINDATAPyWrapper. # In the future VmecINDATAPyWrapper and the C++ VmecINDATA will merge into one type, # and this will become a Python wrapper around the one C++ VmecINDATA type. @@ -169,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). + + 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: int = 0 + 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). @@ -347,6 +403,14 @@ class VmecInput(BaseModelWithNumpy): ] = FreeBoundaryMethod.NESTOR """Method for handling free-boundary conditions.""" + iteration_style: typing.Annotated[ + IterationStyle, + pydantic.BeforeValidator(_validate_iteration_style), + pydantic.Field(), + ] = IterationStyle.VMEC_8_52 + """Time-step / restart control scheme for the equilibrium iteration (``"vmec_8_52"`` + or ``"parvmec"``).""" + nstep: int = 10 """Printout interval at which convergence progress is logged.""" @@ -438,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) @@ -453,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 @@ -591,7 +657,10 @@ def _to_cpp_vmecindata(self) -> _vmecpp.VmecINDATA: } for attr in VmecInput.model_fields: - if attr in readonly_attrs or attr == "free_boundary_method": + if attr in readonly_attrs or attr in ( + "free_boundary_method", + "iteration_style", + ): continue # these must be set separately setattr(cpp_indata, attr, getattr(self, attr)) @@ -599,9 +668,14 @@ def _to_cpp_vmecindata(self) -> _vmecpp.VmecINDATA: cpp_indata.free_boundary_method = getattr( _vmecpp.FreeBoundaryMethod, self.free_boundary_method.upper() ) + cpp_indata.iteration_style = getattr( + _vmecpp.IterationStyle, self.iteration_style.upper() + ) # 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) @@ -2201,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 @@ -2211,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: @@ -2457,7 +2549,6 @@ def set_profile( # items in the generated documentation. __all__ = [ # noqa: RUF022 "run", - "run_continuation", "interpolate_solution", "VmecInput", "VmecOutput", @@ -2468,6 +2559,7 @@ def set_profile( "MakegridParameters", "MagneticFieldResponseTable", "FreeBoundaryMethod", + "IterationStyle", "set_profile", "iterate", "solve_equilibrium", 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/util/netcdf_io/BUILD.bazel b/src/vmecpp/cpp/util/netcdf_io/BUILD.bazel index 2d02e19d4..5452bd619 100644 --- a/src/vmecpp/cpp/util/netcdf_io/BUILD.bazel +++ b/src/vmecpp/cpp/util/netcdf_io/BUILD.bazel @@ -7,8 +7,8 @@ cc_library( hdrs = ["netcdf_io.h"], visibility = ["//visibility:public"], deps = [ - "@abseil-cpp//absl/log:check", - "@abseil-cpp//absl/log:log", + "@abseil-cpp//absl/status", + "@abseil-cpp//absl/status:statusor", "@abseil-cpp//absl/strings", "@abseil-cpp//absl/strings:str_format", "//third_party/netcdf4", @@ -31,6 +31,7 @@ cc_test( ], deps = [ ":netcdf_io", + "@abseil-cpp//absl/status", "@googletest//:gtest_main", ], ) diff --git a/src/vmecpp/cpp/util/netcdf_io/netcdf_io.cc b/src/vmecpp/cpp/util/netcdf_io/netcdf_io.cc index 54c9340b0..eef621527 100644 --- a/src/vmecpp/cpp/util/netcdf_io/netcdf_io.cc +++ b/src/vmecpp/cpp/util/netcdf_io/netcdf_io.cc @@ -7,156 +7,213 @@ #include #include -#include "absl/log/check.h" -#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/strings/ascii.h" #include "absl/strings/str_format.h" #include "netcdf.h" namespace netcdf_io { -bool NetcdfReadBool(int ncid, const std::string& variable_name) { - // VMEC uses `int` to store booleans: 0 means false, otherwise true. - // Also, the actual variable name is `__logical__`. - // AFAIK this is because NetCDF3 did not have a `boolean` data type. +namespace { - // find variable ID for given variable name +absl::StatusOr FindVariableId(int ncid, const std::string& variable_name) { int variable_id = 0; - CHECK_EQ( - nc_inq_varid(ncid, (variable_name + "__logical__").c_str(), &variable_id), - NC_NOERR) - << "variable '" << variable_name << "' not found"; + if (nc_inq_varid(ncid, variable_name.c_str(), &variable_id) != NC_NOERR) { + return absl::NotFoundError( + absl::StrFormat("variable '%s' not found", variable_name)); + } + return variable_id; +} - // figure out rank of data, i.e., how many dimensions does it have +absl::StatusOr GetVariableRank(int ncid, int variable_id, + const std::string& variable_name) { int rank = 0; - CHECK_EQ(nc_inq_varndims(ncid, variable_id, &rank), NC_NOERR); + if (nc_inq_varndims(ncid, variable_id, &rank) != NC_NOERR) { + return absl::InternalError(absl::StrFormat( + "could not determine rank of variable '%s'", variable_name)); + } + return rank; +} - // only accept zero-dimensional array for scalar - CHECK_EQ(rank, 0) << "Not a rank-0 array: " << variable_name; +absl::StatusOr > GetVariableDimensions( + int ncid, int variable_id, int rank, const std::string& variable_name) { + std::vector dimension_ids(rank, 0); + if (nc_inq_vardimid(ncid, variable_id, dimension_ids.data()) != NC_NOERR) { + return absl::InternalError(absl::StrFormat( + "could not determine dimension ids of variable '%s'", variable_name)); + } - // actually read data - int variable_data = 0; - CHECK_EQ(nc_get_var_int(ncid, variable_id, &variable_data), NC_NOERR); + std::vector dimensions(rank, 0); + for (int i = 0; i < rank; ++i) { + size_t dimension = 0; + if (nc_inq_dimlen(ncid, dimension_ids[i], &dimension) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not determine dimension %d of variable '%s'", + i, variable_name)); + } + dimensions[i] = dimension; + } + return dimensions; +} - return (variable_data != 0); -} // NetcdfReadBool +} // namespace -char NetcdfReadChar(int ncid, const std::string& variable_name) { - // find variable ID for given variable name - int variable_id = 0; - CHECK_EQ(nc_inq_varid(ncid, variable_name.c_str(), &variable_id), NC_NOERR) - << "variable '" << variable_name << "' not found"; +absl::StatusOr NetcdfReadBool(int ncid, + const std::string& variable_name) { + // VMEC uses `int` to store booleans: 0 means false, otherwise true. + // Also, the actual variable name is `__logical__`. + // AFAIK this is because NetCDF3 did not have a `boolean` data type. + const std::string logical_variable_name = variable_name + "__logical__"; - // figure out rank of data, i.e., how many dimensions does it have - int rank = 1; - CHECK_EQ(nc_inq_varndims(ncid, variable_id, &rank), NC_NOERR); + absl::StatusOr variable_id = FindVariableId(ncid, logical_variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } - // only accept zero-dimensional array for scalar - CHECK_EQ(rank, 1) << "Not a rank-1 array: " << variable_name; + absl::StatusOr rank = + GetVariableRank(ncid, *variable_id, logical_variable_name); + if (!rank.ok()) { + return rank.status(); + } + if (*rank != 0) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-0 array: %s", logical_variable_name)); + } - // figure out the dimension IDs - std::vector dimension_ids(rank, 0); - CHECK_EQ(nc_inq_vardimid(ncid, variable_id, dimension_ids.data()), NC_NOERR); + int variable_data = 0; + if (nc_get_var_int(ncid, *variable_id, &variable_data) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", logical_variable_name)); + } - // figure out dimension of data, i.e., length of string - std::vector dimensions(rank, 0); - size_t total_element_count = 1; - for (int i = 0; i < rank; ++i) { - size_t dimension = 0; - CHECK_EQ(nc_inq_dimlen(ncid, dimension_ids[i], &dimension), NC_NOERR); - dimensions[i] = dimension; - total_element_count *= dimension; + return variable_data != 0; +} // NetcdfReadBool + +absl::StatusOr NetcdfReadChar(int ncid, + const std::string& variable_name) { + absl::StatusOr variable_id = FindVariableId(ncid, variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } + + absl::StatusOr rank = GetVariableRank(ncid, *variable_id, variable_name); + if (!rank.ok()) { + return rank.status(); + } + if (*rank != 1) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-1 array: %s", variable_name)); + } + + absl::StatusOr > dimensions = + GetVariableDimensions(ncid, *variable_id, *rank, variable_name); + if (!dimensions.ok()) { + return dimensions.status(); } // for a single char, make sure that the array dimension is 1 - CHECK_EQ(dimensions[0], (size_t)1) - << "Not a length-1 array: " << variable_name; + if ((*dimensions)[0] != 1) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a length-1 array: %s", variable_name)); + } // actually read data - std::vector read_start_indices(rank, 0); - std::vector variable_data(total_element_count, 0); - CHECK_EQ(nc_get_vara(ncid, variable_id, read_start_indices.data(), - dimensions.data(), variable_data.data()), - NC_NOERR); + std::vector read_start_indices(*rank, 0); + std::vector variable_data(1, 0); + if (nc_get_vara(ncid, *variable_id, read_start_indices.data(), + dimensions->data(), variable_data.data()) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", variable_name)); + } return variable_data[0]; } // NetcdfReadChar -int NetcdfReadInt(int ncid, const std::string& variable_name) { - // find variable ID for given variable name - int variable_id = 0; - CHECK_EQ(nc_inq_varid(ncid, variable_name.c_str(), &variable_id), NC_NOERR) - << "variable '" << variable_name << "' not found"; - - // figure out rank of data, i.e., how many dimensions does it have - int rank = 0; - CHECK_EQ(nc_inq_varndims(ncid, variable_id, &rank), NC_NOERR); +absl::StatusOr NetcdfReadInt(int ncid, const std::string& variable_name) { + absl::StatusOr variable_id = FindVariableId(ncid, variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } - // only accept zero-dimensional array for scalar - CHECK_EQ(rank, 0) << "Not a rank-0 array: " << variable_name; + absl::StatusOr rank = GetVariableRank(ncid, *variable_id, variable_name); + if (!rank.ok()) { + return rank.status(); + } + if (*rank != 0) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-0 array: %s", variable_name)); + } - // actually read data int variable_data = 0; - CHECK_EQ(nc_get_var_int(ncid, variable_id, &variable_data), NC_NOERR); + if (nc_get_var_int(ncid, *variable_id, &variable_data) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", variable_name)); + } return variable_data; } // NetcdfReadInt -double NetcdfReadDouble(int ncid, const std::string& variable_name) { - // find variable ID for given variable name - int variable_id = 0; - CHECK_EQ(nc_inq_varid(ncid, variable_name.c_str(), &variable_id), NC_NOERR) - << "variable '" << variable_name << "' not found"; - - // figure out rank of data, i.e., how many dimensions does it have - int rank = 0; - CHECK_EQ(nc_inq_varndims(ncid, variable_id, &rank), NC_NOERR); +absl::StatusOr NetcdfReadDouble(int ncid, + const std::string& variable_name) { + absl::StatusOr variable_id = FindVariableId(ncid, variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } - // only accept zero-dimensional array for scalar - CHECK_EQ(rank, 0) << "Not a rank-0 array: " << variable_name; + absl::StatusOr rank = GetVariableRank(ncid, *variable_id, variable_name); + if (!rank.ok()) { + return rank.status(); + } + if (*rank != 0) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-0 array: %s", variable_name)); + } - // actually read data double variable_data = 0; - CHECK_EQ(nc_get_var_double(ncid, variable_id, &variable_data), NC_NOERR); + if (nc_get_var_double(ncid, *variable_id, &variable_data) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", variable_name)); + } return variable_data; } // NetcdfReadDouble -std::string NetcdfReadString(int ncid, const std::string& variable_name) { - // find variable ID for given variable name - int varid = 0; - CHECK_EQ(nc_inq_varid(ncid, variable_name.c_str(), &varid), NC_NOERR) - << "variable '" << variable_name << "' not found"; - - // figure out rank of data, i.e., how many dimensions does it have - int rank = 0; - CHECK_EQ(nc_inq_varndims(ncid, varid, &rank), NC_NOERR); +absl::StatusOr NetcdfReadString(int ncid, + const std::string& variable_name) { + absl::StatusOr variable_id = FindVariableId(ncid, variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } + absl::StatusOr rank = GetVariableRank(ncid, *variable_id, variable_name); + if (!rank.ok()) { + return rank.status(); + } // only accept one-dimensional array of CHAR for strings - CHECK_EQ(rank, 1) << "Not a rank-1 array: " << variable_name; - - // figure out the dimension IDs - std::vector dimension_ids(rank, 0); - CHECK_EQ(nc_inq_vardimid(ncid, varid, dimension_ids.data()), NC_NOERR); + if (*rank != 1) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-1 array: %s", variable_name)); + } - // figure out dimension of data, i.e., length of string - std::vector dimensions(rank, 0); - size_t total_element_count = 1; - for (int i = 0; i < rank; ++i) { - size_t dimension = 0; - CHECK_EQ(nc_inq_dimlen(ncid, dimension_ids[i], &dimension), NC_NOERR); - dimensions[i] = dimension; - total_element_count *= dimension; + absl::StatusOr > dimensions = + GetVariableDimensions(ncid, *variable_id, *rank, variable_name); + if (!dimensions.ok()) { + return dimensions.status(); } + size_t total_element_count = (*dimensions)[0]; + // actually read data - std::vector read_start_indices(rank, 0); + std::vector read_start_indices(*rank, 0); // one extra element that stays at 0 in order to properly zero-terminate the // string std::vector variable_data(total_element_count + 1, 0); - CHECK_EQ(nc_get_vara(ncid, varid, read_start_indices.data(), - dimensions.data(), variable_data.data()), - NC_NOERR); + if (nc_get_vara(ncid, *variable_id, read_start_indices.data(), + dimensions->data(), variable_data.data()) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", variable_name)); + } std::string string_from_char_array = std::string(variable_data.data()); // Strings are usually whitespace-padded when coming from Fortran @@ -164,78 +221,73 @@ std::string NetcdfReadString(int ncid, const std::string& variable_name) { return std::string(absl::StripAsciiWhitespace(string_from_char_array)); } // NetcdfReadString -std::vector NetcdfReadArray1D(int ncid, - const std::string& variable_name) { - // find variable ID for given variable name - int variable_id = 0; - CHECK_EQ(nc_inq_varid(ncid, variable_name.c_str(), &variable_id), NC_NOERR) - << "variable '" << variable_name << "' not found"; - - // figure out rank of data, i.e., how many dimensions does it have - int rank = 1; - CHECK_EQ(nc_inq_varndims(ncid, variable_id, &rank), NC_NOERR); - - // only accept zero-dimensional array for scalar - CHECK_EQ(rank, 1) << "Not a rank-1 array: " << variable_name; +absl::StatusOr > NetcdfReadArray1D( + int ncid, const std::string& variable_name) { + absl::StatusOr variable_id = FindVariableId(ncid, variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } - // figure out the dimension IDs - std::vector dimension_ids(rank, 0); - CHECK_EQ(nc_inq_vardimid(ncid, variable_id, dimension_ids.data()), NC_NOERR); + absl::StatusOr rank = GetVariableRank(ncid, *variable_id, variable_name); + if (!rank.ok()) { + return rank.status(); + } + if (*rank != 1) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-1 array: %s", variable_name)); + } - // figure out dimension of data, i.e., length of string - std::vector dimensions(rank, 0); - size_t total_element_count = 1; - for (int i = 0; i < rank; ++i) { - size_t dimension = 0; - CHECK_EQ(nc_inq_dimlen(ncid, dimension_ids[i], &dimension), NC_NOERR); - dimensions[i] = dimension; - total_element_count *= dimension; + absl::StatusOr > dimensions = + GetVariableDimensions(ncid, *variable_id, *rank, variable_name); + if (!dimensions.ok()) { + return dimensions.status(); } - // actually read data - std::vector read_start_indices(rank, 0); + size_t total_element_count = (*dimensions)[0]; + + std::vector read_start_indices(*rank, 0); std::vector variable_data(total_element_count, 0.0); - CHECK_EQ(nc_get_vara(ncid, variable_id, read_start_indices.data(), - dimensions.data(), variable_data.data()), - NC_NOERR); + if (nc_get_vara(ncid, *variable_id, read_start_indices.data(), + dimensions->data(), variable_data.data()) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", variable_name)); + } return variable_data; } // NetcdfReadArray1D -std::vector > NetcdfReadArray2D( +absl::StatusOr > > NetcdfReadArray2D( int ncid, const std::string& variable_name) { - // find variable ID for given variable name - int variable_id = 0; - CHECK_EQ(nc_inq_varid(ncid, variable_name.c_str(), &variable_id), NC_NOERR) - << "variable '" << variable_name << "' not found"; - - // figure out rank of data, i.e., how many dimensions does it have - int rank = 1; - CHECK_EQ(nc_inq_varndims(ncid, variable_id, &rank), NC_NOERR); - - // only accept zero-dimensional array for scalar - CHECK_EQ(rank, 2) << "Not a rank-2 array: " << variable_name; + absl::StatusOr variable_id = FindVariableId(ncid, variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } - // figure out the dimension IDs - std::vector dimension_ids(rank, 0); - CHECK_EQ(nc_inq_vardimid(ncid, variable_id, dimension_ids.data()), NC_NOERR); + absl::StatusOr rank = GetVariableRank(ncid, *variable_id, variable_name); + if (!rank.ok()) { + return rank.status(); + } + if (*rank != 2) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-2 array: %s", variable_name)); + } - // figure out dimension of data, i.e., length of string - std::vector dimensions(rank, 0); - size_t total_element_count = 1; - for (int i = 0; i < rank; ++i) { - size_t dimension = 0; - CHECK_EQ(nc_inq_dimlen(ncid, dimension_ids[i], &dimension), NC_NOERR); - dimensions[i] = dimension; - total_element_count *= dimension; + absl::StatusOr > dimensions_or = + GetVariableDimensions(ncid, *variable_id, *rank, variable_name); + if (!dimensions_or.ok()) { + return dimensions_or.status(); } + auto dimensions = dimensions_or.value(); - // actually read data - std::vector read_start_indices(rank, 0); + size_t total_element_count = dimensions[0] * dimensions[1]; + + std::vector read_start_indices(*rank, 0); std::vector variable_data(total_element_count, 0.0); - CHECK_EQ(nc_get_vara(ncid, variable_id, read_start_indices.data(), - dimensions.data(), variable_data.data()), - NC_NOERR); + if (nc_get_vara(ncid, *variable_id, read_start_indices.data(), + dimensions.data(), variable_data.data()) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", variable_name)); + } // copy from flattened vector into two-dimensional vector of vectors std::vector > two_dimensional_data(dimensions[0]); @@ -249,40 +301,37 @@ std::vector > NetcdfReadArray2D( return two_dimensional_data; } // NetcdfReadArray2D -std::vector > > NetcdfReadArray3D( - int ncid, const std::string& variable_name) { - // find variable ID for given variable name - int variable_id = 0; - CHECK_EQ(nc_inq_varid(ncid, variable_name.c_str(), &variable_id), NC_NOERR) - << "variable '" << variable_name << "' not found"; - - // figure out rank of data, i.e., how many dimensions does it have - int rank = 1; - CHECK_EQ(nc_inq_varndims(ncid, variable_id, &rank), NC_NOERR); - - // only accept zero-dimensional array for scalar - CHECK_EQ(rank, 3) << "Not a rank-3 array: " << variable_name; +absl::StatusOr > > > +NetcdfReadArray3D(int ncid, const std::string& variable_name) { + absl::StatusOr variable_id = FindVariableId(ncid, variable_name); + if (!variable_id.ok()) { + return variable_id.status(); + } - // figure out the dimension IDs - std::vector dimension_ids(rank, 0); - CHECK_EQ(nc_inq_vardimid(ncid, variable_id, dimension_ids.data()), NC_NOERR); + absl::StatusOr rank = GetVariableRank(ncid, *variable_id, variable_name); + if (!rank.ok()) { + return rank.status(); + } + if (*rank != 3) { + return absl::InvalidArgumentError( + absl::StrFormat("Not a rank-3 array: %s", variable_name)); + } - // figure out dimension of data, i.e., length of string - std::vector dimensions(rank, 0); - size_t total_element_count = 1; - for (int i = 0; i < rank; ++i) { - size_t dimension = 0; - CHECK_EQ(nc_inq_dimlen(ncid, dimension_ids[i], &dimension), NC_NOERR); - dimensions[i] = dimension; - total_element_count *= dimension; + absl::StatusOr > dimensions_or = + GetVariableDimensions(ncid, *variable_id, *rank, variable_name); + if (!dimensions_or.ok()) { + return dimensions_or.status(); } + auto dimensions = dimensions_or.value(); - // actually read data - std::vector read_start_indices(rank, 0); + size_t total_element_count = dimensions[0] * dimensions[1] * dimensions[2]; + std::vector read_start_indices(*rank, 0); std::vector variable_data(total_element_count, 0.0); - CHECK_EQ(nc_get_vara(ncid, variable_id, read_start_indices.data(), - dimensions.data(), variable_data.data()), - NC_NOERR); + if (nc_get_vara(ncid, *variable_id, read_start_indices.data(), + dimensions.data(), variable_data.data()) != NC_NOERR) { + return absl::InternalError( + absl::StrFormat("could not read variable '%s'", variable_name)); + } // copy from flattened vector into three-dimensional vector of vectors std::vector > > three_dimensional_data( diff --git a/src/vmecpp/cpp/util/netcdf_io/netcdf_io.h b/src/vmecpp/cpp/util/netcdf_io/netcdf_io.h index fc2fc2f67..87d189823 100644 --- a/src/vmecpp/cpp/util/netcdf_io/netcdf_io.h +++ b/src/vmecpp/cpp/util/netcdf_io/netcdf_io.h @@ -8,44 +8,48 @@ #include #include +#include "absl/status/statusor.h" + namespace netcdf_io { // Read a scalar `bool` variable in Fortran VMEC style from the (opened) NetCDF // file identified by `ncid`. It is expected that the value is stored in a // scalar `int` variable named `__logical__`. -bool NetcdfReadBool(int ncid, const std::string& variable_name); +absl::StatusOr NetcdfReadBool(int ncid, const std::string& variable_name); // Read a scalar `char` variable from the (opened) NetCDF file identified by // `ncid`. It is expected that the value is stored in a length-1 `char` array. -char NetcdfReadChar(int ncid, const std::string& variable_name); +absl::StatusOr NetcdfReadChar(int ncid, const std::string& variable_name); // Read a scalar `int` variable from the (opened) NetCDF file identified by // `ncid`. -int NetcdfReadInt(int ncid, const std::string& variable_name); +absl::StatusOr NetcdfReadInt(int ncid, const std::string& variable_name); // Read a scalar `double` variable from the (opened) NetCDF file identified by // `ncid`. -double NetcdfReadDouble(int ncid, const std::string& variable_name); +absl::StatusOr NetcdfReadDouble(int ncid, + const std::string& variable_name); // Read a string from the (opened) NetCDF file identified by `ncid`. // It is expected that the data is stored as a rank-1 `char` array. // Whitespace at the start and end of the `char` array is stripped. -std::string NetcdfReadString(int ncid, const std::string& variable_name); +absl::StatusOr NetcdfReadString(int ncid, + const std::string& variable_name); // Read a rank-1 `double` array from the (opened) NetCDF file identified by // `ncid`. -std::vector NetcdfReadArray1D(int ncid, - const std::string& variable_name); +absl::StatusOr > NetcdfReadArray1D( + int ncid, const std::string& variable_name); // Read a rank-2 `double` array from the (opened) NetCDF file identified by // `ncid`. -std::vector > NetcdfReadArray2D( +absl::StatusOr > > NetcdfReadArray2D( int ncid, const std::string& variable_name); // Read a rank-3 `double` array from the (opened) NetCDF file identified by // `ncid`. -std::vector > > NetcdfReadArray3D( - int ncid, const std::string& variable_name); +absl::StatusOr > > > +NetcdfReadArray3D(int ncid, const std::string& variable_name); } // namespace netcdf_io diff --git a/src/vmecpp/cpp/util/netcdf_io/netcdf_io_test.cc b/src/vmecpp/cpp/util/netcdf_io/netcdf_io_test.cc index 8afab1824..d23391702 100644 --- a/src/vmecpp/cpp/util/netcdf_io/netcdf_io_test.cc +++ b/src/vmecpp/cpp/util/netcdf_io/netcdf_io_test.cc @@ -9,6 +9,7 @@ #include #include +#include "absl/status/status.h" #include "gmock/gmock.h" #include "gtest/gtest.h" @@ -22,9 +23,10 @@ TEST(TestNetcdfIO, CheckReadBool) { int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - const bool lasym = NetcdfReadBool(ncid, "lasym"); + absl::StatusOr lasym = NetcdfReadBool(ncid, "lasym"); - EXPECT_FALSE(lasym); + ASSERT_TRUE(lasym.ok()); + EXPECT_FALSE(*lasym); ASSERT_EQ(nc_close(ncid), NC_NOERR); } // CheckReadBool @@ -35,9 +37,10 @@ TEST(TestNetcdfIO, CheckReadChar) { int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - const char mgrid_mode = NetcdfReadChar(ncid, "mgrid_mode"); + absl::StatusOr mgrid_mode = NetcdfReadChar(ncid, "mgrid_mode"); - EXPECT_EQ(mgrid_mode, 'R'); + ASSERT_TRUE(mgrid_mode.ok()); + EXPECT_EQ(*mgrid_mode, 'R'); ASSERT_EQ(nc_close(ncid), NC_NOERR); } // CheckReadChar @@ -48,22 +51,38 @@ TEST(TestNetcdfIO, CheckReadInt) { int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - const int nfp = NetcdfReadInt(ncid, "nfp"); + absl::StatusOr nfp = NetcdfReadInt(ncid, "nfp"); - EXPECT_EQ(nfp, 5); + ASSERT_TRUE(nfp.ok()); + EXPECT_EQ(*nfp, 5); ASSERT_EQ(nc_close(ncid), NC_NOERR); } // CheckReadInt +TEST(TestNetcdfIO, CheckReadIntMissingVariable) { + const std::string example_netcdf = "util/netcdf_io/example_netcdf.nc"; + + int ncid = 0; + ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); + + absl::StatusOr missing = NetcdfReadInt(ncid, "does_not_exist"); + + EXPECT_FALSE(missing.ok()); + EXPECT_EQ(missing.status().code(), absl::StatusCode::kNotFound); + + ASSERT_EQ(nc_close(ncid), NC_NOERR); +} // CheckReadIntMissingVariable + TEST(TestNetcdfIO, CheckReadDouble) { const std::string example_netcdf = "util/netcdf_io/example_netcdf.nc"; int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - const double ftolv = NetcdfReadDouble(ncid, "ftolv"); + absl::StatusOr ftolv = NetcdfReadDouble(ncid, "ftolv"); - EXPECT_EQ(ftolv, 1.0e-10); + ASSERT_TRUE(ftolv.ok()); + EXPECT_EQ(*ftolv, 1.0e-10); ASSERT_EQ(nc_close(ncid), NC_NOERR); } // CheckReadDouble @@ -74,9 +93,10 @@ TEST(TestNetcdfIO, CheckReadString) { int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - const std::string mgrid_file = NetcdfReadString(ncid, "mgrid_file"); + absl::StatusOr mgrid_file = NetcdfReadString(ncid, "mgrid_file"); - EXPECT_EQ(mgrid_file, "mgrid_cth_like.nc"); + ASSERT_TRUE(mgrid_file.ok()); + EXPECT_EQ(*mgrid_file, "mgrid_cth_like.nc"); ASSERT_EQ(nc_close(ncid), NC_NOERR); } // CheckReadString @@ -87,7 +107,9 @@ TEST(TestNetcdfIO, CheckReadArray1D) { int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - std::vector am = NetcdfReadArray1D(ncid, "am"); + absl::StatusOr > am = NetcdfReadArray1D(ncid, "am"); + + ASSERT_TRUE(am.ok()); // `am` is stored in the wout file with its default (maximum) length // and only the first few (relevant) entries are actually populated. @@ -96,7 +118,7 @@ TEST(TestNetcdfIO, CheckReadArray1D) { reference_am[1] = 5.0; reference_am[2] = 10.0; - EXPECT_THAT(am, ElementsAreArray(reference_am)); + EXPECT_THAT(*am, ElementsAreArray(reference_am)); ASSERT_EQ(nc_close(ncid), NC_NOERR); } // CheckReadArray1D @@ -107,14 +129,17 @@ TEST(TestNetcdfIO, CheckReadArray2D) { int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - std::vector > rmnc = NetcdfReadArray2D(ncid, "rmnc"); + absl::StatusOr > > rmnc = + NetcdfReadArray2D(ncid, "rmnc"); + + ASSERT_TRUE(rmnc.ok()); std::vector > reference_rmnc = {{0.0, 1.0, 2.0}, {0.1, 1.1, 2.1}}; - ASSERT_EQ(rmnc.size(), reference_rmnc.size()); + ASSERT_EQ(rmnc->size(), reference_rmnc.size()); for (size_t i = 0; i < reference_rmnc.size(); ++i) { - EXPECT_THAT(rmnc[i], ElementsAreArray(reference_rmnc[i])); + EXPECT_THAT((*rmnc)[i], ElementsAreArray(reference_rmnc[i])); } ASSERT_EQ(nc_close(ncid), NC_NOERR); @@ -126,9 +151,11 @@ TEST(TestNetcdfIO, CheckReadArray3D) { int ncid = 0; ASSERT_EQ(nc_open(example_netcdf.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - std::vector > > br_001 = + absl::StatusOr > > > br_001 = NetcdfReadArray3D(ncid, "br_001"); + ASSERT_TRUE(br_001.ok()); + std::vector > > reference_br_001 = { {{0.00, 0.01, 0.02, 0.03}, {0.10, 0.11, 0.12, 0.13}, @@ -137,11 +164,11 @@ TEST(TestNetcdfIO, CheckReadArray3D) { {1.10, 1.11, 1.12, 1.13}, {1.20, 1.21, 1.22, 1.23}}}; - ASSERT_EQ(br_001.size(), reference_br_001.size()); + ASSERT_EQ(br_001->size(), reference_br_001.size()); for (size_t i = 0; i < reference_br_001.size(); ++i) { - ASSERT_EQ(br_001[i].size(), reference_br_001[i].size()); + ASSERT_EQ((*br_001)[i].size(), reference_br_001[i].size()); for (size_t j = 0; j < reference_br_001[i].size(); ++j) { - EXPECT_THAT(br_001[i][j], ElementsAreArray(reference_br_001[i][j])); + EXPECT_THAT((*br_001)[i][j], ElementsAreArray(reference_br_001[i][j])); } } diff --git a/src/vmecpp/cpp/vmecpp/common/CMakeLists.txt b/src/vmecpp/cpp/vmecpp/common/CMakeLists.txt index ef739f391..8a6e325ff 100644 --- a/src/vmecpp/cpp/vmecpp/common/CMakeLists.txt +++ b/src/vmecpp/cpp/vmecpp/common/CMakeLists.txt @@ -1,6 +1,7 @@ add_subdirectory(composed_types_definition) add_subdirectory(composed_types_lib) add_subdirectory(flow_control) +add_subdirectory(fourier_basis) add_subdirectory(fourier_basis_fast_poloidal) add_subdirectory(fourier_basis_fast_toroidal) add_subdirectory(magnetic_configuration_definition) diff --git a/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.cc b/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.cc index ae577452c..70a013022 100644 --- a/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.cc +++ b/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.cc @@ -60,6 +60,7 @@ FlowControl::FlowControl(bool lfreeb, double delt, int num_grids, ijacob = 0; restart_reason = RestartReason::NO_RESTART; res0 = -1; + res1 = -1; delt0r = delt; multi_ns_grid = num_grids; neqs_old = 0; diff --git a/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h b/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h index dd4900572..ab7e2cb79 100644 --- a/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h +++ b/src/vmecpp/cpp/vmecpp/common/flow_control/flow_control.h @@ -107,7 +107,11 @@ class FlowControl { // occurred) std::vector restart_reasons; + // Running minimum of the preconditioned residual sum (fsq). double res0; + // Running minimum of the invariant residual sum (fsqr + fsqz + fsql); used + // only by the PARVMEC time-step control. + double res1; Eigen::Vector3d fResInvar; Eigen::Vector3d fResPrecd; diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis/BUILD.bazel b/src/vmecpp/cpp/vmecpp/common/fourier_basis/BUILD.bazel new file mode 100644 index 000000000..6ea1c2176 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis/BUILD.bazel @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# SPDX-License-Identifier: MIT +cc_library( + name = "fourier_basis", + srcs = ["fourier_basis.cc"], + hdrs = ["fourier_basis.h"], + visibility = ["//visibility:public"], + deps = [ + "@abseil-cpp//absl/algorithm:container", + "@abseil-cpp//absl/log:check", + "@abseil-cpp//absl/strings:str_format", + "//vmecpp/common/util:util", + "//vmecpp/common/sizes:sizes", + ], +) diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis/CMakeLists.txt b/src/vmecpp/cpp/vmecpp/common/fourier_basis/CMakeLists.txt new file mode 100644 index 000000000..360a82747 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis/CMakeLists.txt @@ -0,0 +1,5 @@ +list (APPEND vmecpp_sources + ${CMAKE_CURRENT_SOURCE_DIR}/fourier_basis.cc + ${CMAKE_CURRENT_SOURCE_DIR}/fourier_basis.h +) +set (vmecpp_sources "${vmecpp_sources}" PARENT_SCOPE) diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.cc b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.cc similarity index 70% rename from src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.cc rename to src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.cc index 40a727889..85d3b6551 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.cc +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.cc @@ -2,7 +2,7 @@ // // // SPDX-License-Identifier: MIT -#include "vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.h" +#include "vmecpp/common/fourier_basis/fourier_basis.h" #include #include @@ -13,7 +13,8 @@ namespace vmecpp { -FourierBasisFastPoloidal::FourierBasisFastPoloidal(const Sizes* s) : s_(*s) { +template +FourierBasis::FourierBasis(const Sizes* s) : s_(*s) { mscale.resize(s_.mnyq2 + 1); nscale.resize(s_.nnyq2 + 1); @@ -31,7 +32,7 @@ FourierBasisFastPoloidal::FourierBasisFastPoloidal(const Sizes* s) : s_(*s) { cosnvn.resize((s_.nnyq2 + 1) * s_.nZeta); sinnvn.resize((s_.nnyq2 + 1) * s_.nZeta); - computeFourierBasisFastPoloidal(s_.nfp); + computeFourierBasis(s_.nfp); // ----------------- @@ -51,7 +52,8 @@ FourierBasisFastPoloidal::FourierBasisFastPoloidal(const Sizes* s) : s_(*s) { s_.mnyq + 1, s_.nfp); } -void FourierBasisFastPoloidal::computeFourierBasisFastPoloidal(int nfp) { +template +void FourierBasis::computeFourierBasis(int nfp) { static constexpr double kTwoPi = 2.0 * M_PI; // Fourier transforms are always computed in VMEC @@ -77,7 +79,8 @@ void FourierBasisFastPoloidal::computeFourierBasisFastPoloidal(int nfp) { for (int l = 0; l < s_.nThetaReduced; ++l) { // need to compute theta grid using _full_ number of theta points! const double theta = kTwoPi * l / s_.nThetaEven; - const int idx_ml = m * s_.nThetaReduced + l; + const int idx_ml = + Layout::PoloidalBasisIndex(m, l, s_.mnyq2 + 1, s_.nThetaReduced); const double arg = m * theta; @@ -118,7 +121,8 @@ void FourierBasisFastPoloidal::computeFourierBasisFastPoloidal(int nfp) { for (int k = 0; k < s_.nZeta; ++k) { const double zeta = kTwoPi * k / s_.nZeta; for (int n = 0; n < s_.nnyq2 + 1; ++n) { - const int idx_kn = k * (s_.nnyq2 + 1) + n; + const int idx_kn = + Layout::ToroidalBasisIndex(n, k, s_.nnyq2 + 1, s_.nZeta); const double arg = n * zeta; @@ -134,10 +138,11 @@ void FourierBasisFastPoloidal::computeFourierBasisFastPoloidal(int nfp) { } // convert cos(xm[mn] theta - xn[mn] zeta) into 2D FC array form -int FourierBasisFastPoloidal::cos_to_cc_ss(const std::span fcCos, - std::span m_fcCC, - std::span m_fcSS, int n_size, - int m_size) const { +template +int FourierBasis::cos_to_cc_ss(const std::span fcCos, + std::span m_fcCC, + std::span m_fcSS, int n_size, + int m_size) const { // m = 0: n = 0, 1, ..., ntor --> ntor + 1 // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); @@ -155,7 +160,7 @@ int FourierBasisFastPoloidal::cos_to_cc_ss(const std::span fcCos, double normedFC = basis_norm * fcCos[mn]; - m_fcCC[m * (n_size + 1) + abs_n] += normedFC; + m_fcCC[Layout::ProductIndex(m, abs_n, m_size, n_size)] += normedFC; // no contribution to fcSS where (m == 0 || n == 0) mn++; @@ -170,9 +175,10 @@ int FourierBasisFastPoloidal::cos_to_cc_ss(const std::span fcCos, double normedFC = basis_norm * fcCos[mn]; - m_fcCC[m * (n_size + 1) + abs_n] += normedFC; + m_fcCC[Layout::ProductIndex(m, abs_n, m_size, n_size)] += normedFC; if (abs_n > 0) { - m_fcSS[m * (n_size + 1) + abs_n] += sgn_n * normedFC; + m_fcSS[Layout::ProductIndex(m, abs_n, m_size, n_size)] += + sgn_n * normedFC; } mn++; @@ -185,10 +191,11 @@ int FourierBasisFastPoloidal::cos_to_cc_ss(const std::span fcCos, return mnmax; } -int FourierBasisFastPoloidal::sin_to_sc_cs(const std::span fcSin, - std::span m_fcSC, - std::span m_fcCS, int n_size, - int m_size) const { +template +int FourierBasis::sin_to_sc_cs(const std::span fcSin, + std::span m_fcSC, + std::span m_fcCS, int n_size, + int m_size) const { // m = 0: n = 0, 1, ..., ntor --> ntor + 1 // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); @@ -209,7 +216,7 @@ int FourierBasisFastPoloidal::sin_to_sc_cs(const std::span fcSin, // no contribution to fcSC where m == 0 // check for n > 0 is redundant when starting loop at n=1 - m_fcCS[m * (n_size + 1) + abs_n] = -sgn_n * normedFC; + m_fcCS[Layout::ProductIndex(m, abs_n, m_size, n_size)] = -sgn_n * normedFC; mn++; } @@ -223,9 +230,10 @@ int FourierBasisFastPoloidal::sin_to_sc_cs(const std::span fcSin, double normedFC = basis_norm * fcSin[mn]; - m_fcSC[m * (n_size + 1) + abs_n] += normedFC; + m_fcSC[Layout::ProductIndex(m, abs_n, m_size, n_size)] += normedFC; if (abs_n > 0) { - m_fcCS[m * (n_size + 1) + abs_n] += -sgn_n * normedFC; + m_fcCS[Layout::ProductIndex(m, abs_n, m_size, n_size)] += + -sgn_n * normedFC; } mn++; @@ -238,10 +246,11 @@ int FourierBasisFastPoloidal::sin_to_sc_cs(const std::span fcSin, return mnmax; } -int FourierBasisFastPoloidal::cc_ss_to_cos(const std::span fcCC, - const std::span fcSS, - std::span m_fcCos, - int n_size, int m_size) const { +template +int FourierBasis::cc_ss_to_cos(const std::span fcCC, + const std::span fcSS, + std::span m_fcCos, int n_size, + int m_size) const { // m = 0: n = 0, 1, ..., ntor --> ntor + 1 // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); @@ -254,7 +263,7 @@ int FourierBasisFastPoloidal::cc_ss_to_cos(const std::span fcCC, for (int n = 0; n < n_size + 1; ++n) { double basis_norm = 1.0 / (mscale[m] * nscale[n]); - m_fcCos[mn] = fcCC[n] / basis_norm; + m_fcCos[mn] = fcCC[Layout::ProductIndex(m, n, m_size, n_size)] / basis_norm; mn++; } // n @@ -267,10 +276,11 @@ int FourierBasisFastPoloidal::cc_ss_to_cos(const std::span fcCC, double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); if (abs_n == 0) { - m_fcCos[mn] = fcCC[m * (n_size + 1) + abs_n] / basis_norm; + m_fcCos[mn] = + fcCC[Layout::ProductIndex(m, abs_n, m_size, n_size)] / basis_norm; } else { - double raw_cc = fcCC[m * (n_size + 1) + abs_n]; - double raw_ss = fcSS[m * (n_size + 1) + abs_n]; + double raw_cc = fcCC[Layout::ProductIndex(m, abs_n, m_size, n_size)]; + double raw_ss = fcSS[Layout::ProductIndex(m, abs_n, m_size, n_size)]; m_fcCos[mn] = 0.5 * (raw_cc + sgn_n * raw_ss) / basis_norm; } @@ -284,10 +294,11 @@ int FourierBasisFastPoloidal::cc_ss_to_cos(const std::span fcCC, return mnmax; } -int FourierBasisFastPoloidal::sc_cs_to_sin(const std::span fcSC, - const std::span fcCS, - std::span m_fcSin, - int n_size, int m_size) const { +template +int FourierBasis::sc_cs_to_sin(const std::span fcSC, + const std::span fcCS, + std::span m_fcSin, int n_size, + int m_size) const { // m = 0: n = 0, 1, ..., ntor --> ntor + 1 // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); @@ -300,7 +311,8 @@ int FourierBasisFastPoloidal::sc_cs_to_sin(const std::span fcSC, for (int n = 1; n < n_size + 1; ++n) { double basis_norm = 1.0 / (mscale[m] * nscale[n]); - m_fcSin[mn] = -fcCS[n] / basis_norm; + m_fcSin[mn] = + -fcCS[Layout::ProductIndex(m, n, m_size, n_size)] / basis_norm; mn++; } // n @@ -313,10 +325,11 @@ int FourierBasisFastPoloidal::sc_cs_to_sin(const std::span fcSC, double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); if (abs_n == 0) { - m_fcSin[mn] = fcSC[m * (n_size + 1) + abs_n] / basis_norm; + m_fcSin[mn] = + fcSC[Layout::ProductIndex(m, abs_n, m_size, n_size)] / basis_norm; } else { - double raw_sc = fcSC[m * (n_size + 1) + abs_n]; - double raw_cs = fcCS[m * (n_size + 1) + abs_n]; + double raw_sc = fcSC[Layout::ProductIndex(m, abs_n, m_size, n_size)]; + double raw_cs = fcCS[Layout::ProductIndex(m, abs_n, m_size, n_size)]; m_fcSin[mn] = 0.5 * (raw_sc - sgn_n * raw_cs) / basis_norm; } @@ -330,7 +343,8 @@ int FourierBasisFastPoloidal::sc_cs_to_sin(const std::span fcSC, return mnmax; } -int FourierBasisFastPoloidal::mnIdx(int m, int n) const { +template +int FourierBasis::mnIdx(int m, int n) const { if (m == 0) { CHECK_GE(n, 0) << "no mn index available for n < 0"; return n; @@ -342,7 +356,8 @@ int FourierBasisFastPoloidal::mnIdx(int m, int n) const { // number of unique Fourier coefficients for // m = 0, 1, ..., m_size - 1 // n = -n_size, -(n_size-1), ..., -1, 0, 1, ..., (n_size-1), n_size -int FourierBasisFastPoloidal::mnMax(int m_size, int n_size) const { +template +int FourierBasis::mnMax(int m_size, int n_size) const { // m = 0: n = 0, 1, ..., ntor --> ntor + 1 // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); @@ -350,10 +365,11 @@ int FourierBasisFastPoloidal::mnMax(int m_size, int n_size) const { return mnmax; } -void FourierBasisFastPoloidal::computeConversionIndices(Eigen::VectorXi& m_xm, - Eigen::VectorXi& m_xn, - int n_size, int m_size, - int nfp) const { +template +void FourierBasis::computeConversionIndices(Eigen::VectorXi& m_xm, + Eigen::VectorXi& m_xn, + int n_size, int m_size, + int nfp) const { const int mnmax = mnMax(m_size, n_size); int mn = 0; @@ -375,4 +391,8 @@ void FourierBasisFastPoloidal::computeConversionIndices(Eigen::VectorXi& m_xm, CHECK_EQ(mn, mnmax) << "counting error: mn=" << mn << " should be " << mnmax; } +// The two layouts that VMEC++ (theta fast) and Nestor (zeta fast) use. +template class FourierBasis; +template class FourierBasis; + } // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.h b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.h new file mode 100644 index 000000000..f4148e007 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.h @@ -0,0 +1,364 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#ifndef VMECPP_COMMON_FOURIER_BASIS_FOURIER_BASIS_H_ +#define VMECPP_COMMON_FOURIER_BASIS_FOURIER_BASIS_H_ + +#include +#include + +#include "vmecpp/common/sizes/sizes.h" + +namespace vmecpp { + +// The two data layouts differ only in which flat index a (mode, grid-point) +// pair maps to. VMEC++ iterates its basis arrays with the poloidal (theta) +// coordinate as the fast axis, while Nestor iterates with the toroidal (zeta) +// coordinate fastest, so the two store the identical basis values in transposed +// memory order. Everything else (the trigonometric arithmetic, the scaling +// factors, the combined <-> product basis conversions, the mode-number +// bookkeeping) is layout independent. The layout is therefore a compile-time +// policy supplying only the three flat index formulas; FourierBasis below is +// written once against it. +// +// FastPoloidal (m-major): the poloidal (theta) coordinate is the fast index of +// the poloidal basis arrays, and the poloidal mode m is the slow index of the +// product-basis coefficient arrays. Used by the VMEC++ core solver. +struct FourierBasisFastPoloidalLayout { + // Poloidal basis arrays, logical shape [num_m][num_l] over (mode m, theta l). + static int PoloidalBasisIndex(int m, int l, int num_m, int num_l) { + (void)num_m; + return m * num_l + l; + } + // Toroidal basis arrays, logical shape [num_k][num_n] over (zeta k, mode n). + static int ToroidalBasisIndex(int n, int k, int num_n, int num_k) { + (void)num_k; + return k * num_n + n; + } + // Product-basis coefficient arrays, logical shape [m_size][n_size + 1]. + static int ProductIndex(int m, int n, int m_size, int n_size) { + (void)m_size; + return m * (n_size + 1) + n; + } +}; + +// FastToroidal (n-major): the toroidal (zeta) coordinate is the fast index of +// the toroidal basis arrays, and the toroidal mode n is the slow index of the +// product-basis coefficient arrays. Used by Nestor / the free-boundary code. +struct FourierBasisFastToroidalLayout { + static int PoloidalBasisIndex(int m, int l, int num_m, int num_l) { + (void)num_l; + return l * num_m + m; + } + static int ToroidalBasisIndex(int n, int k, int num_n, int num_k) { + (void)num_n; + return n * num_k + k; + } + static int ProductIndex(int m, int n, int m_size, int n_size) { + (void)n_size; + return n * m_size + m; + } +}; + +// Fourier basis representation for VMEC++ spectral computations. +// +// This class provides the fundamental spectral basis for VMEC++ computations, +// representing 3D plasma quantities using Fourier decomposition in flux +// coordinates (s, \theta, \zeta) where: +// s = normalized toroidal flux (radial coordinate) +// \theta = poloidal angle +// \zeta = toroidal angle = nfp * \phi (field period toroidal angle) +// +// Physical quantities are expanded as: +// f(s,\theta,\zeta) = \sum_{m,n} f_{mn}(s) * basis_function(m*\theta, +// n*\zeta) +// +// The Layout template policy fixes the flat memory order of the basis and +// coefficient arrays (see FourierBasisFastPoloidalLayout / +// FourierBasisFastToroidalLayout). The concrete classes are the two type +// aliases at the bottom of this header: +// FourierBasisFastPoloidal (theta fast; VMEC++ core solver) +// FourierBasisFastToroidal (zeta fast; Nestor / free boundary) +template +class FourierBasis { + public: + explicit FourierBasis(const Sizes* s); + + // ============================================================================ + // FOURIER BASIS SCALING FACTORS + // ============================================================================ + + // [mnyq2+1] Poloidal mode scaling factors: sqrt(2) for m>0, 1.0 for m=0 + // Applied to cos(m*\theta) and sin(m*\theta) basis functions for DFT + // normalization Enables proper normalization: 1/\pi for m>0 modes, 1/(2\pi) + // for m=0 mode + Eigen::VectorXd mscale; + + // [nnyq2+1] Toroidal mode scaling factors: sqrt(2) for n>0, 1.0 for n=0 + // Applied to cos(n*\zeta) and sin(n*\zeta) basis functions for DFT + // normalization Enables proper normalization: 1/\pi for n>0 modes, 1/(2\pi) + // for n=0 mode + Eigen::VectorXd nscale; + + // ============================================================================ + // POLOIDAL BASIS FUNCTIONS + // ============================================================================ + // Flat index: Layout::PoloidalBasisIndex(m, l, mnyq2+1, nThetaReduced). + // \theta[l] = 2*\pi*l/nThetaEven for l=0...nThetaReduced-1 (reduced [0,\pi] + // interval). + + // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal cosine basis + // cosmu[idx(m,l)] = cos(m*\theta[l]) * mscale[m] + Eigen::VectorXd cosmu; + + // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal sine basis + // sinmu[idx(m,l)] = sin(m*\theta[l]) * mscale[m] + Eigen::VectorXd sinmu; + + // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal cosine derivative + // cosmum[idx(m,l)] = m * cos(m*\theta[l]) * mscale[m] + // Used for computing \partial/\partial\theta derivatives in force + // calculations + Eigen::VectorXd cosmum; + + // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal sine derivative + // sinmum[idx(m,l)] = -m * sin(m*\theta[l]) * mscale[m] + // Used for computing \partial/\partial\theta derivatives in force + // calculations + Eigen::VectorXd sinmum; + + // ============================================================================ + // POLOIDAL BASIS WITH INTEGRATION WEIGHTS + // ============================================================================ + + // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal cosine basis + // cosmui[idx(m,l)] = cosmu[idx(m,l)] * intNorm + // intNorm = 1/(nZeta*(nThetaReduced-1)), with boundary point factor 1/2 + Eigen::VectorXd cosmui; + + // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal sine basis + // sinmui[idx(m,l)] = sinmu[idx(m,l)] * intNorm + Eigen::VectorXd sinmui; + + // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal cosine derivative + // cosmumi[idx(m,l)] = cosmum[idx(m,l)] * intNorm + Eigen::VectorXd cosmumi; + + // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal sine derivative + // sinmumi[idx(m,l)] = sinmum[idx(m,l)] * intNorm + Eigen::VectorXd sinmumi; + + // ============================================================================ + // TOROIDAL BASIS FUNCTIONS + // ============================================================================ + // Flat index: Layout::ToroidalBasisIndex(n, k, nnyq2+1, nZeta). + // \zeta[k] = 2*\pi*k/nZeta for k=0...nZeta-1 (full [0,2\pi] interval). + + // [(nnyq2+1) * nZeta] Pre-scaled toroidal cosine basis + // cosnv[idx(n,k)] = cos(n*\zeta[k]) * nscale[n] + Eigen::VectorXd cosnv; + + // [(nnyq2+1) * nZeta] Pre-scaled toroidal sine basis + // sinnv[idx(n,k)] = sin(n*\zeta[k]) * nscale[n] + Eigen::VectorXd sinnv; + + // [(nnyq2+1) * nZeta] Pre-scaled toroidal cosine derivative with nfp factor + // cosnvn[idx(n,k)] = n*nfp * cos(n*\zeta[k]) * nscale[n] + // Factor nfp converts \partial/\partial\zeta to \partial/\partial\phi + // derivatives + Eigen::VectorXd cosnvn; + + // [(nnyq2+1) * nZeta] Pre-scaled toroidal sine derivative with nfp factor + // sinnvn[idx(n,k)] = -n*nfp * sin(n*\zeta[k]) * nscale[n] + // Factor nfp converts \partial/\partial\zeta to \partial/\partial\phi + // derivatives + Eigen::VectorXd sinnvn; + + // ============================================================================ + // FOURIER BASIS CONVERSION FUNCTIONS + // ============================================================================ + // + // These functions convert between VMEC++'s two Fourier basis representations + // using trigonometric identities and pre-computed scaling factors. + // See docs/fourier_basis_implementation.md for complete mathematical details. + // + // Two Fourier basis types: + // 1. COMBINED BASIS (External): cos(m*\theta - n*\zeta), sin(m*\theta - + // n*\zeta) + // - Used in: wout files, Python API, traditional VMEC format + // - Storage: Linear arrays indexed by mode number mn + // + // 2. PRODUCT BASIS (Internal): cos(m*\theta)*cos(n*\zeta), + // sin(m*\theta)*sin(n*\zeta), etc. + // - Used in: Internal computations with separable DFT operations + // - Storage: 2D arrays indexed by (m,n) via Layout::ProductIndex + // + // Mathematical basis function identity: + // cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + + // sin(m*\theta)*sin(n*\zeta) + + /** + * Convert coefficients from combined cosine basis to separable product basis. + * + * Basis function identity: + * cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + + * sin(m*\theta)*sin(n*\zeta) + * + * This function transforms coefficients for cos(m*\theta - n*\zeta) basis + * functions into coefficients for the separable product basis + * cos(m*\theta)*cos(n*\zeta) and sin(m*\theta)*sin(n*\zeta). The + * transformation accounts for VMEC symmetry where only n >= 0 coefficients + * are stored. + * + * Implementation uses pre-computed scaling factors (mscale, nscale) and + * handles positive/negative toroidal mode symmetry. Standalone function. + * + * Physics context: Converts external coefficient format (wout files) to + * internal product basis coefficients that enable separable DFT operations. + * + * @param fcCos [input] Coefficients for cos(m*\theta - n*\zeta) basis, size + * mnmax + * @param m_fcCC [output] Coefficients for cos(m*\theta)*cos(n*\zeta) basis, + * size m_size*(n_size+1) + * @param m_fcSS [output] Coefficients for sin(m*\theta)*sin(n*\zeta) basis, + * size m_size*(n_size+1) + * @param n_size Toroidal mode range: n in [-n_size, n_size] + * @param m_size Poloidal mode range: m in [0, m_size-1] + * @return Total number of modes processed (mnmax) + */ + int cos_to_cc_ss(const std::span fcCos, + std::span m_fcCC, std::span m_fcSS, + int n_size, int m_size) const; + + /** + * Convert coefficients from combined sine basis to separable product basis. + * + * Basis function identity: + * sin(m*\theta - n*\zeta) = sin(m*\theta)*cos(n*\zeta) - + * cos(m*\theta)*sin(n*\zeta) + * + * This function transforms coefficients for sin(m*\theta - n*\zeta) basis + * functions into coefficients for the separable product basis + * sin(m*\theta)*cos(n*\zeta) and cos(m*\theta)*sin(n*\zeta). Enforces + * sin(0*\theta - 0*\zeta) = 0 constraint. + * + * Physics context: Handles sine-parity quantities like Z coordinates (zmns) + * and \lambda angle functions (lmns coefficients). + * + * @param fcSin [input] Coefficients for sin(m*\theta - n*\zeta) basis, size + * mnmax + * @param m_fcSC [output] Coefficients for sin(m*\theta)*cos(n*\zeta) basis, + * size m_size*(n_size+1) + * @param m_fcCS [output] Coefficients for cos(m*\theta)*sin(n*\zeta) basis, + * size m_size*(n_size+1) + * @param n_size Toroidal mode range: n in [-n_size, n_size] + * @param m_size Poloidal mode range: m in [0, m_size-1] + * @return Total number of modes processed (mnmax) + */ + int sin_to_sc_cs(const std::span fcSin, + std::span m_fcSC, std::span m_fcCS, + int n_size, int m_size) const; + + /** + * Convert coefficients from separable product basis back to combined cosine + * basis. + * + * Inverse transformation using basis function identity: + * cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + + * sin(m*\theta)*sin(n*\zeta) + * + * This function reconstructs coefficients for cos(m*\theta - n*\zeta) basis + * from coefficients of the separable product basis. Handles positive/negative + * toroidal mode reconstruction and applies inverse scaling factors. + * + * Physics context: Converts internal computational results back to external + * coefficient format for wout files, Python API, and traditional VMEC output. + * + * @param fcCC [input] Coefficients for cos(m*\theta)*cos(n*\zeta) basis, size + * m_size*(n_size+1) + * @param fcSS [input] Coefficients for sin(m*\theta)*sin(n*\zeta) basis, size + * m_size*(n_size+1) + * @param m_fcCos [output] Coefficients for cos(m*\theta - n*\zeta) basis, + * size mnmax + * @param n_size Toroidal mode range: n in [-n_size, n_size] + * @param m_size Poloidal mode range: m in [0, m_size-1] + * @return Total number of modes processed (mnmax) + */ + int cc_ss_to_cos(const std::span fcCC, + const std::span fcSS, + std::span m_fcCos, int n_size, int m_size) const; + + /** + * Convert coefficients from separable product basis back to combined sine + * basis. + * + * Inverse transformation using basis function identity: + * sin(m*\theta - n*\zeta) = sin(m*\theta)*cos(n*\zeta) - + * cos(m*\theta)*sin(n*\zeta) + * + * This function reconstructs coefficients for sin(m*\theta - n*\zeta) basis + * from coefficients of the separable product basis. Enforces sin(0*\theta - + * 0*\zeta) = 0. + * + * Physics context: Converts internal results for sine-parity quantities + * back to external coefficient format for output and analysis. + * + * @param fcSC [input] Coefficients for sin(m*\theta)*cos(n*\zeta) basis, size + * m_size*(n_size+1) + * @param fcCS [input] Coefficients for cos(m*\theta)*sin(n*\zeta) basis, size + * m_size*(n_size+1) + * @param m_fcSin [output] Coefficients for sin(m*\theta - n*\zeta) basis, + * size mnmax + * @param n_size Toroidal mode range: n in [-n_size, n_size] + * @param m_size Poloidal mode range: m in [0, m_size-1] + * @return Total number of modes processed (mnmax) + */ + int sc_cs_to_sin(const std::span fcSC, + const std::span fcCS, + std::span m_fcSin, int n_size, int m_size) const; + + int mnIdx(int m, int n) const; + int mnMax(int m_size, int n_size) const; + void computeConversionIndices(Eigen::VectorXi& m_xm, Eigen::VectorXi& m_xn, + int n_size, int m_size, int nfp) const; + + // ============================================================================ + // MODE NUMBER MAPPING ARRAYS + // ============================================================================ + + // [mnmax] Poloidal mode numbers for standard resolution Fourier coefficients + // Layout: xm[mn] = poloidal mode number m for the mn-th coefficient + // Maps linear coefficient index mn to 2D mode (m,n) for spectral operations + Eigen::VectorXi xm; + + // [mnmax] Toroidal mode numbers for standard resolution Fourier coefficients + // Layout: xn[mn] = toroidal mode number n*nfp for the mn-th coefficient + // Factor nfp included to convert from field periods to geometric toroidal + // modes + Eigen::VectorXi xn; + + // [mnmax_nyq] Poloidal mode numbers for Nyquist-extended Fourier coefficients + // Layout: xm_nyq[mn] = poloidal mode number m for the mn-th Nyquist + // coefficient Extended resolution to avoid aliasing in nonlinear force + // calculations + Eigen::VectorXi xm_nyq; + + // [mnmax_nyq] Toroidal mode numbers for Nyquist-extended Fourier coefficients + // Layout: xn_nyq[mn] = toroidal mode number n*nfp for the mn-th Nyquist + // coefficient Extended resolution to avoid aliasing in nonlinear force + // calculations + Eigen::VectorXi xn_nyq; + + private: + const Sizes& s_; + + void computeFourierBasis(int nfp); +}; + +using FourierBasisFastPoloidal = FourierBasis; +using FourierBasisFastToroidal = FourierBasis; + +} // namespace vmecpp + +#endif // VMECPP_COMMON_FOURIER_BASIS_FOURIER_BASIS_H_ diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/BUILD.bazel b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/BUILD.bazel index 67b28a424..a650f0a60 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/BUILD.bazel @@ -3,15 +3,10 @@ # SPDX-License-Identifier: MIT cc_library( name = "fourier_basis_fast_poloidal", - srcs = ["fourier_basis_fast_poloidal.cc"], hdrs = ["fourier_basis_fast_poloidal.h"], visibility = ["//visibility:public"], deps = [ - "@abseil-cpp//absl/algorithm:container", - "@abseil-cpp//absl/log:check", - "@abseil-cpp//absl/strings:str_format", - "//vmecpp/common/util:util", - "//vmecpp/common/sizes:sizes", + "//vmecpp/common/fourier_basis", ], ) diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/CMakeLists.txt b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/CMakeLists.txt index bba3e7834..9c2daccb2 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/CMakeLists.txt +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/CMakeLists.txt @@ -1,5 +1,4 @@ list (APPEND vmecpp_sources - ${CMAKE_CURRENT_SOURCE_DIR}/fourier_basis_fast_poloidal.cc ${CMAKE_CURRENT_SOURCE_DIR}/fourier_basis_fast_poloidal.h ) set (vmecpp_sources "${vmecpp_sources}" PARENT_SCOPE) diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.h b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.h index f565fbc98..725216c83 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.h +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.h @@ -5,308 +5,9 @@ #ifndef VMECPP_COMMON_FOURIER_BASIS_FAST_POLOIDAL_FOURIER_BASIS_FAST_POLOIDAL_H_ #define VMECPP_COMMON_FOURIER_BASIS_FAST_POLOIDAL_FOURIER_BASIS_FAST_POLOIDAL_H_ -#include -#include - -#include "vmecpp/common/sizes/sizes.h" - -namespace vmecpp { - -// Fourier basis representation optimized for poloidal coordinate operations. -// -// This class provides the fundamental spectral basis for VMEC++ computations, -// representing 3D plasma quantities using Fourier decomposition in flux -// coordinates (s, \theta, \zeta) where: -// s = normalized toroidal flux (radial coordinate) -// \theta = poloidal angle -// \zeta = toroidal angle = nfp * \phi (field period toroidal angle) -// -// Physical quantities are expanded as: -// f(s,\theta,\zeta) = \sum_{m,n} f_{mn}(s) * basis_function(m*\theta, -// n*\zeta) -// -// The "FastPoloidal" layout stores data with poloidal (\theta) coordinate as -// the fast (innermost) loop index, optimizing for operations that iterate -// over poloidal modes. This differs from FastToroidal layout. -// -// NOTE: Nestor has its own implementation of this class because we want to be -// able to use different data layouts between VMEC++ and Nestor. -class FourierBasisFastPoloidal { - public: - explicit FourierBasisFastPoloidal(const Sizes* s); - - // ============================================================================ - // FOURIER BASIS SCALING FACTORS - // ============================================================================ - - // [mnyq2+1] Poloidal mode scaling factors: sqrt(2) for m>0, 1.0 for m=0 - // Applied to cos(m*\theta) and sin(m*\theta) basis functions for DFT - // normalization Enables proper normalization: 1/\pi for m>0 modes, 1/(2\pi) - // for m=0 mode - Eigen::VectorXd mscale; - - // [nnyq2+1] Toroidal mode scaling factors: sqrt(2) for n>0, 1.0 for n=0 - // Applied to cos(n*\zeta) and sin(n*\zeta) basis functions for DFT - // normalization Enables proper normalization: 1/\pi for n>0 modes, 1/(2\pi) - // for n=0 mode - Eigen::VectorXd nscale; - - // ============================================================================ - // POLOIDAL BASIS FUNCTIONS (m-major layout: [m][l]) - // ============================================================================ - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal cosine basis - // Layout: cosmu[m*nThetaReduced + l] = cos(m*\theta[l]) * mscale[m] - // \theta[l] = 2*\pi*l/nThetaEven for l=0...nThetaReduced-1 (reduced [0,\pi] - // interval) - Eigen::VectorXd cosmu; - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal sine basis - // Layout: sinmu[m*nThetaReduced + l] = sin(m*\theta[l]) * mscale[m] - Eigen::VectorXd sinmu; - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal cosine derivative - // Layout: cosmum[m*nThetaReduced + l] = m * cos(m*\theta[l]) * mscale[m] - // Used for computing \partial/\partial\theta derivatives in force - // calculations - Eigen::VectorXd cosmum; - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal sine derivative - // Layout: sinmum[m*nThetaReduced + l] = -m * sin(m*\theta[l]) * mscale[m] - // Used for computing \partial/\partial\theta derivatives in force - // calculations - Eigen::VectorXd sinmum; - - // ============================================================================ - // POLOIDAL BASIS WITH INTEGRATION WEIGHTS - // ============================================================================ - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal cosine basis - // Layout: cosmui[m*nThetaReduced + l] = cosmu[m*nThetaReduced + l] * intNorm - // intNorm = 1/(nZeta*(nThetaReduced-1)), with boundary point factor 1/2 - Eigen::VectorXd cosmui; - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal sine basis - // Layout: sinmui[m*nThetaReduced + l] = sinmu[m*nThetaReduced + l] * intNorm - Eigen::VectorXd sinmui; - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal cosine derivative - // Layout: cosmumi[m*nThetaReduced + l] = cosmum[m*nThetaReduced + l] * - // intNorm - Eigen::VectorXd cosmumi; - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal sine derivative - // Layout: sinmumi[m*nThetaReduced + l] = sinmum[m*nThetaReduced + l] * - // intNorm - Eigen::VectorXd sinmumi; - - // ============================================================================ - // TOROIDAL BASIS FUNCTIONS (zeta-major layout: [k][n]) - // ============================================================================ - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal cosine basis - // Layout: cosnv[k*(nnyq2+1) + n] = cos(n*\zeta[k]) * nscale[n] - // \zeta[k] = 2*\pi*k/nZeta for k=0...nZeta-1 (full [0,2\pi] interval) - Eigen::VectorXd cosnv; - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal sine basis - // Layout: sinnv[k*(nnyq2+1) + n] = sin(n*\zeta[k]) * nscale[n] - Eigen::VectorXd sinnv; - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal cosine derivative with nfp factor - // Layout: cosnvn[k*(nnyq2+1) + n] = n*nfp * cos(n*\zeta[k]) * nscale[n] - // Factor nfp converts \partial/\partial\zeta to \partial/\partial\phi - // derivatives - Eigen::VectorXd cosnvn; - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal sine derivative with nfp factor - // Layout: sinnvn[k*(nnyq2+1) + n] = -n*nfp * sin(n*\zeta[k]) * nscale[n] - // Factor nfp converts \partial/\partial\zeta to \partial/\partial\phi - // derivatives - Eigen::VectorXd sinnvn; - - // ============================================================================ - // FOURIER BASIS CONVERSION FUNCTIONS - // ============================================================================ - // - // These functions convert between VMEC++'s two Fourier basis representations - // using trigonometric identities and pre-computed scaling factors. - // See docs/fourier_basis_implementation.md for complete mathematical details. - // - // Two Fourier basis types: - // 1. COMBINED BASIS (External): cos(m*\theta - n*\zeta), sin(m*\theta - - // n*\zeta) - // - Used in: wout files, Python API, traditional VMEC format - // - Storage: Linear arrays indexed by mode number mn - // - // 2. PRODUCT BASIS (Internal): cos(m*\theta)*cos(n*\zeta), - // sin(m*\theta)*sin(n*\zeta), etc. - // - Used in: Internal computations with separable DFT operations - // - Storage: 2D arrays indexed by (m,n) separately - // - Layout: fcCC[m*(n_size+1) + n] (m-major ordering for poloidal class) - // - // Mathematical basis function identity: - // cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + - // sin(m*\theta)*sin(n*\zeta) - - /** - * Convert coefficients from combined cosine basis to separable product basis. - * - * Basis function identity: - * cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + - * sin(m*\theta)*sin(n*\zeta) - * - * This function transforms coefficients for cos(m*\theta - n*\zeta) basis - * functions into coefficients for the separable product basis - * cos(m*\theta)*cos(n*\zeta) and sin(m*\theta)*sin(n*\zeta). The - * transformation accounts for VMEC symmetry where only n >= 0 coefficients - * are stored. - * - * Implementation uses pre-computed scaling factors (mscale, nscale) and - * handles positive/negative toroidal mode symmetry. Standalone function. - * - * Physics context: Converts external coefficient format (wout files) to - * internal product basis coefficients that enable separable DFT operations. - * - * @param fcCos [input] Coefficients for cos(m*\theta - n*\zeta) basis, size - * mnmax - * @param m_fcCC [output] Coefficients for cos(m*\theta)*cos(n*\zeta) basis, - * size m_size*(n_size+1) - * @param m_fcSS [output] Coefficients for sin(m*\theta)*sin(n*\zeta) basis, - * size m_size*(n_size+1) - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int cos_to_cc_ss(const std::span fcCos, - std::span m_fcCC, std::span m_fcSS, - int n_size, int m_size) const; - - /** - * Convert coefficients from combined sine basis to separable product basis. - * - * Basis function identity: - * sin(m*\theta - n*\zeta) = sin(m*\theta)*cos(n*\zeta) - - * cos(m*\theta)*sin(n*\zeta) - * - * This function transforms coefficients for sin(m*\theta - n*\zeta) basis - * functions into coefficients for the separable product basis - * sin(m*\theta)*cos(n*\zeta) and cos(m*\theta)*sin(n*\zeta). Enforces - * sin(0*\theta - 0*\zeta) = 0 constraint. - * - * Physics context: Handles sine-parity quantities like Z coordinates (zmns) - * and \lambda angle functions (lmns coefficients). - * - * @param fcSin [input] Coefficients for sin(m*\theta - n*\zeta) basis, size - * mnmax - * @param m_fcSC [output] Coefficients for sin(m*\theta)*cos(n*\zeta) basis, - * size m_size*(n_size+1) - * @param m_fcCS [output] Coefficients for cos(m*\theta)*sin(n*\zeta) basis, - * size m_size*(n_size+1) - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int sin_to_sc_cs(const std::span fcSin, - std::span m_fcSC, std::span m_fcCS, - int n_size, int m_size) const; - - /** - * Convert coefficients from separable product basis back to combined cosine - * basis. - * - * Inverse transformation using basis function identity: - * cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + - * sin(m*\theta)*sin(n*\zeta) - * - * This function reconstructs coefficients for cos(m*\theta - n*\zeta) basis - * from coefficients of the separable product basis. Handles positive/negative - * toroidal mode reconstruction and applies inverse scaling factors. - * - * Physics context: Converts internal computational results back to external - * coefficient format for wout files, Python API, and traditional VMEC output. - * - * @param fcCC [input] Coefficients for cos(m*\theta)*cos(n*\zeta) basis, size - * m_size*(n_size+1) - * @param fcSS [input] Coefficients for sin(m*\theta)*sin(n*\zeta) basis, size - * m_size*(n_size+1) - * @param m_fcCos [output] Coefficients for cos(m*\theta - n*\zeta) basis, - * size mnmax - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int cc_ss_to_cos(const std::span fcCC, - const std::span fcSS, - std::span m_fcCos, int n_size, int m_size) const; - - /** - * Convert coefficients from separable product basis back to combined sine - * basis. - * - * Inverse transformation using basis function identity: - * sin(m*\theta - n*\zeta) = sin(m*\theta)*cos(n*\zeta) - - * cos(m*\theta)*sin(n*\zeta) - * - * This function reconstructs coefficients for sin(m*\theta - n*\zeta) basis - * from coefficients of the separable product basis. Enforces sin(0*\theta - - * 0*\zeta) = 0. - * - * Physics context: Converts internal results for sine-parity quantities - * back to external coefficient format for output and analysis. - * - * @param fcSC [input] Coefficients for sin(m*\theta)*cos(n*\zeta) basis, size - * m_size*(n_size+1) - * @param fcCS [input] Coefficients for cos(m*\theta)*sin(n*\zeta) basis, size - * m_size*(n_size+1) - * @param m_fcSin [output] Coefficients for sin(m*\theta - n*\zeta) basis, - * size mnmax - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int sc_cs_to_sin(const std::span fcSC, - const std::span fcCS, - std::span m_fcSin, int n_size, int m_size) const; - - int mnIdx(int m, int n) const; - int mnMax(int m_size, int n_size) const; - void computeConversionIndices(Eigen::VectorXi& m_xm, Eigen::VectorXi& m_xn, - int n_size, int m_size, int nfp) const; - - // ============================================================================ - // MODE NUMBER MAPPING ARRAYS - // ============================================================================ - - // [mnmax] Poloidal mode numbers for standard resolution Fourier coefficients - // Layout: xm[mn] = poloidal mode number m for the mn-th coefficient - // Maps linear coefficient index mn to 2D mode (m,n) for spectral operations - Eigen::VectorXi xm; - - // [mnmax] Toroidal mode numbers for standard resolution Fourier coefficients - // Layout: xn[mn] = toroidal mode number n*nfp for the mn-th coefficient - // Factor nfp included to convert from field periods to geometric toroidal - // modes - Eigen::VectorXi xn; - - // [mnmax_nyq] Poloidal mode numbers for Nyquist-extended Fourier coefficients - // Layout: xm_nyq[mn] = poloidal mode number m for the mn-th Nyquist - // coefficient Extended resolution to avoid aliasing in nonlinear force - // calculations - Eigen::VectorXi xm_nyq; - - // [mnmax_nyq] Toroidal mode numbers for Nyquist-extended Fourier coefficients - // Layout: xn_nyq[mn] = toroidal mode number n*nfp for the mn-th Nyquist - // coefficient Extended resolution to avoid aliasing in nonlinear force - // calculations - Eigen::VectorXi xn_nyq; - - private: - const Sizes& s_; - - void computeFourierBasisFastPoloidal(int nfp); -}; - -} // namespace vmecpp +// FourierBasisFastPoloidal is the theta-fast (m-major) layout of the shared +// FourierBasis template, which now holds the single implementation. This header +// is retained as a stable include path for existing call sites. +#include "vmecpp/common/fourier_basis/fourier_basis.h" // IWYU pragma: export #endif // VMECPP_COMMON_FOURIER_BASIS_FAST_POLOIDAL_FOURIER_BASIS_FAST_POLOIDAL_H_ diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/BUILD.bazel b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/BUILD.bazel index 6e267502c..0d641c945 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/BUILD.bazel @@ -3,15 +3,10 @@ # SPDX-License-Identifier: MIT cc_library( name = "fourier_basis_fast_toroidal", - srcs = ["fourier_basis_fast_toroidal.cc"], hdrs = ["fourier_basis_fast_toroidal.h"], visibility = ["//visibility:public"], deps = [ - "@abseil-cpp//absl/algorithm:container", - "@abseil-cpp//absl/log:check", - "@abseil-cpp//absl/strings:str_format", - "//vmecpp/common/util:util", - "//vmecpp/common/sizes:sizes", + "//vmecpp/common/fourier_basis", ], ) diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/CMakeLists.txt b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/CMakeLists.txt index e0511fbe0..8bf47828d 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/CMakeLists.txt +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/CMakeLists.txt @@ -1,5 +1,4 @@ list (APPEND vmecpp_sources - ${CMAKE_CURRENT_SOURCE_DIR}/fourier_basis_fast_toroidal.cc ${CMAKE_CURRENT_SOURCE_DIR}/fourier_basis_fast_toroidal.h ) set (vmecpp_sources "${vmecpp_sources}" PARENT_SCOPE) diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.cc b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.cc deleted file mode 100644 index 40384f4f6..000000000 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.cc +++ /dev/null @@ -1,381 +0,0 @@ -// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH -// -// -// SPDX-License-Identifier: MIT -#include "vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.h" - -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/log/check.h" -#include "vmecpp/common/util/util.h" - -namespace vmecpp { - -FourierBasisFastToroidal::FourierBasisFastToroidal(const Sizes* s) : s_(*s) { - mscale.resize(s_.mnyq2 + 1); - nscale.resize(s_.nnyq2 + 1); - - cosmu.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - sinmu.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - cosmum.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - sinmum.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - cosmui.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - sinmui.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - cosmumi.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - sinmumi.resize(s_.nThetaReduced * (s_.mnyq2 + 1)); - - cosnv.resize((s_.nnyq2 + 1) * s_.nZeta); - sinnv.resize((s_.nnyq2 + 1) * s_.nZeta); - cosnvn.resize((s_.nnyq2 + 1) * s_.nZeta); - sinnvn.resize((s_.nnyq2 + 1) * s_.nZeta); - - computeFourierBasisFastToroidal(s_.nfp); - - // ----------------- - - xm.resize(s_.mnmax); - xm.setZero(); - xn.resize(s_.mnmax); - xn.setZero(); - - computeConversionIndices(/*m_xm=*/xm, /*m_xn=*/xn, s_.ntor, s_.mpol, s_.nfp); - - xm_nyq.resize(s_.mnmax_nyq); - xm_nyq.setZero(); - xn_nyq.resize(s_.mnmax_nyq); - xn_nyq.setZero(); - - computeConversionIndices(/*m_xm=*/xm_nyq, /*m_xn=*/xn_nyq, s_.nnyq, - s_.mnyq + 1, s_.nfp); -} - -void FourierBasisFastToroidal::computeFourierBasisFastToroidal(int nfp) { - static constexpr double kTwoPi = 2.0 * M_PI; - - // Fourier transforms are always computed in VMEC - // over the reduced theta interval from [0, pi]. - // Thus, need a fixed normalization factor (cannot use dnorm3 or wInt in - // Sizes) here. - const double intNorm = 1.0 / (s_.nZeta * (s_.nThetaReduced - 1)); - - // poloidal - for (int m = 0; m < s_.mnyq2 + 1; ++m) { - // DFTs for m>0 need 1/pi==2/(2pi) normalization factor - // vs. 1/(2pi) for the cos(m=0)-mode. - // --> introduce one sqrt(2) in fwd-DFT (geometry-into-realspace) - // and one sqrt(2) into inv-DFT (forces-into-Fourier) via mscale - if (m == 0) { - mscale[m] = 1.0; - } else { - mscale[m] = std::numbers::sqrt2; - } - } // m - - for (int l = 0; l < s_.nThetaReduced; ++l) { - // need to compute theta grid using _full_ number of theta points! - const double theta = kTwoPi * l / s_.nThetaEven; - for (int m = 0; m < s_.mnyq2 + 1; ++m) { - const int idx_lm = l * (s_.mnyq2 + 1) + m; - - const double arg = m * theta; - - // poloidal Fourier basis - cosmu[idx_lm] = std::cos(arg) * mscale[m]; - sinmu[idx_lm] = std::sin(arg) * mscale[m]; - - // integration - cosmui[idx_lm] = cosmu[idx_lm] * intNorm; - sinmui[idx_lm] = sinmu[idx_lm] * intNorm; - - if (l == 0 || l == s_.nThetaReduced - 1) { - cosmui[idx_lm] /= 2.0; - } - - // poloidal derivatives - cosmum[idx_lm] = m * cosmu[idx_lm]; - sinmum[idx_lm] = -m * sinmu[idx_lm]; - - cosmumi[idx_lm] = m * cosmui[idx_lm]; - sinmumi[idx_lm] = -m * sinmui[idx_lm]; - } // m - } // l - - // toroidal - for (int n = 0; n < s_.nnyq2 + 1; ++n) { - // DFTs for m>0 need 1/pi==2/(2pi) normalization factor - // vs. 1/(2pi) for the cos(m=0)-mode. - // --> introduce one sqrt(2) in fwd-DFT (geometry-into-realspace) - // and one sqrt(2) into inv-DFT (forces-into-Fourier) via nscale - if (n == 0) { - nscale[n] = 1.0; - } else { - nscale[n] = std::numbers::sqrt2; - } - } // n - - for (int k = 0; k < s_.nZeta; ++k) { - const double zeta = kTwoPi * k / s_.nZeta; - for (int n = 0; n < s_.nnyq2 + 1; ++n) { - const int idx_nk = n * s_.nZeta + k; - - const double arg = n * zeta; - - // toroidal Fourier basis - cosnv[idx_nk] = std::cos(arg) * nscale[n]; - sinnv[idx_nk] = std::sin(arg) * nscale[n]; - - // toroidal derivatives - cosnvn[idx_nk] = n * nfp * cosnv[idx_nk]; - sinnvn[idx_nk] = -n * nfp * sinnv[idx_nk]; - } // n - } // k -} - -// convert cos(xm[mn] theta - xn[mn] zeta) into 2D FC array form -int FourierBasisFastToroidal::cos_to_cc_ss(const std::span fcCos, - std::span m_fcCC, - std::span m_fcSS, int n_size, - int m_size) const { - // m = 0: n = 0, 1, ..., ntor --> ntor + 1 - // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) - int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); - - absl::c_fill_n(m_fcCC, m_size * (n_size + 1), 0); - absl::c_fill_n(m_fcSS, m_size * (n_size + 1), 0); - - int mn = 0; - - int m = 0; - for (int n = 0; n < n_size + 1; ++n) { - int abs_n = abs(n); - - double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); - - double normedFC = basis_norm * fcCos[mn]; - - m_fcCC[abs_n * m_size + m] += normedFC; - // no contribution to fcSS where (m == 0 || n == 0) - - mn++; - } - - for (m = 1; m < m_size; ++m) { - for (int n = -n_size; n < n_size + 1; ++n) { - int abs_n = abs(n); - int sgn_n = signum(n); - - double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); - - double normedFC = basis_norm * fcCos[mn]; - - m_fcCC[abs_n * m_size + m] += normedFC; - if (abs_n > 0) { - m_fcSS[abs_n * m_size + m] += sgn_n * normedFC; - } - - mn++; - } // n - } // m - - CHECK_EQ(mn, mnmax) << "counting error: mn=" << mn << " should be " << mnmax - << " in cos_to_cc_ss"; - - return mnmax; -} - -int FourierBasisFastToroidal::sin_to_sc_cs(const std::span fcSin, - std::span m_fcSC, - std::span m_fcCS, int n_size, - int m_size) const { - // m = 0: n = 0, 1, ..., ntor --> ntor + 1 - // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) - int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); - - absl::c_fill_n(m_fcSC, m_size * (n_size + 1), 0); - absl::c_fill_n(m_fcCS, m_size * (n_size + 1), 0); - - int mn = 1; - - int m = 0; - for (int n = 1; n < n_size + 1; ++n) { - int abs_n = abs(n); - int sgn_n = signum(n); - - double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); - - double normedFC = basis_norm * fcSin[mn]; - - // no contribution to fcSC where m == 0 - // check for n > 0 is redundant when starting loop at n=1 - m_fcCS[abs_n * m_size + m] = -sgn_n * normedFC; - - mn++; - } - - for (m = 1; m < m_size; ++m) { - for (int n = -n_size; n < n_size + 1; ++n) { - int abs_n = abs(n); - int sgn_n = signum(n); - - double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); - - double normedFC = basis_norm * fcSin[mn]; - - m_fcSC[abs_n * m_size + m] += normedFC; - if (abs_n > 0) { - m_fcCS[abs_n * m_size + m] += -sgn_n * normedFC; - } - - mn++; - } // n - } // m - - CHECK_EQ(mn, mnmax) << "counting error: mn=" << mn << " should be " << mnmax - << " in sin_to_sc_cs"; - - return mnmax; -} - -int FourierBasisFastToroidal::cc_ss_to_cos(const std::span fcCC, - const std::span fcSS, - std::span m_fcCos, - int n_size, int m_size) const { - // m = 0: n = 0, 1, ..., ntor --> ntor + 1 - // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) - int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); - - absl::c_fill_n(m_fcCos, mnmax, 0); - - int mn = 0; - - int m = 0; - for (int n = 0; n < n_size + 1; ++n) { - int abs_n = abs(n); - - double basis_norm = 1.0 / (mscale[m] * nscale[n]); - - m_fcCos[mn] = fcCC[abs_n * m_size + m] / basis_norm; - - mn++; - } // n - - for (m = 1; m < m_size; ++m) { - for (int n = -n_size; n < n_size + 1; ++n) { - int abs_n = abs(n); - int sgn_n = signum(n); - - double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); - - if (abs_n == 0) { - m_fcCos[mn] = fcCC[abs_n * m_size + m] / basis_norm; - } else { - double raw_cc = fcCC[abs_n * m_size + m]; - double raw_ss = fcSS[abs_n * m_size + m]; - m_fcCos[mn] = 0.5 * (raw_cc + sgn_n * raw_ss) / basis_norm; - } - - mn++; - } // n - } // m - - CHECK_EQ(mn, mnmax) << "counting error: mn=" << mn << " should be " << mnmax - << " in cc_ss_to_cos"; - - return mnmax; -} - -int FourierBasisFastToroidal::sc_cs_to_sin(const std::span fcSC, - const std::span fcCS, - std::span m_fcSin, - int n_size, int m_size) const { - // m = 0: n = 0, 1, ..., ntor --> ntor + 1 - // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) - int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); - - absl::c_fill_n(m_fcSin, mnmax, 0); - - int mn = 1; - - int m = 0; - for (int n = 1; n < n_size + 1; ++n) { - int abs_n = abs(n); - double basis_norm = 1.0 / (mscale[m] * nscale[n]); - - m_fcSin[mn] = -fcCS[abs_n * m_size + m] / basis_norm; - - mn++; - } // n - - for (m = 1; m < m_size; ++m) { - for (int n = -n_size; n < n_size + 1; ++n) { - int abs_n = abs(n); - int sgn_n = signum(n); - - double basis_norm = 1.0 / (mscale[m] * nscale[abs_n]); - - if (abs_n == 0) { - m_fcSin[mn] = fcSC[abs_n * m_size + m] / basis_norm; - } else { - double raw_sc = fcSC[abs_n * m_size + m]; - double raw_cs = fcCS[abs_n * m_size + m]; - m_fcSin[mn] = 0.5 * (raw_sc - sgn_n * raw_cs) / basis_norm; - } - - mn++; - } // n - } // m - - CHECK_EQ(mn, mnmax) << "counting error: mn=" << mn << " should be " << mnmax - << " in sc_cs_to_sin"; - - return mnmax; -} - -int FourierBasisFastToroidal::mnIdx(int m, int n) const { - if (m == 0) { - CHECK_GE(n, 0) << "no mn index available for n < 0"; - return n; - } else { - return (s_.ntor + 1) + (m - 1) * (2 * s_.ntor + 1) + (n + s_.ntor); - } -} - -// number of unique Fourier coefficients for -// m = 0, 1, ..., m_size - 1 -// n = -n_size, -(n_size-1), ..., -1, 0, 1, ..., (n_size-1), n_size -int FourierBasisFastToroidal::mnMax(int m_size, int n_size) const { - // m = 0: n = 0, 1, ..., ntor --> ntor + 1 - // m > 0: n = -ntor, ..., ntor --> (mpol - 1) * (2 * ntor + 1) - int mnmax = (n_size + 1) + (m_size - 1) * (2 * n_size + 1); - - return mnmax; -} - -void FourierBasisFastToroidal::computeConversionIndices(Eigen::VectorXi& m_xm, - Eigen::VectorXi& m_xn, - int n_size, int m_size, - int nfp) const { - const int mnmax = mnMax(m_size, n_size); - int mn = 0; - - int m = 0; - for (int n = 0; n < n_size + 1; ++n) { - m_xm[mn] = m; - m_xn[mn] = n * nfp; - mn++; - } - - for (m = 1; m < m_size; ++m) { - for (int n = -n_size; n < n_size + 1; ++n) { - m_xm[mn] = m; - m_xn[mn] = n * nfp; - mn++; - } - } - - CHECK_EQ(mn, mnmax) << "counting error: mn=" << mn << " should be " << mnmax; -} - -} // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.h b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.h index 8f1598c4e..383af2a3a 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.h +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.h @@ -5,307 +5,9 @@ #ifndef VMECPP_COMMON_FOURIER_BASIS_FAST_TOROIDAL_FOURIER_BASIS_FAST_TOROIDAL_H_ #define VMECPP_COMMON_FOURIER_BASIS_FAST_TOROIDAL_FOURIER_BASIS_FAST_TOROIDAL_H_ -#include -#include - -#include "vmecpp/common/sizes/sizes.h" - -namespace vmecpp { - -// Fourier basis representation optimized for toroidal coordinate operations. -// -// This class provides the fundamental spectral basis for VMEC++ computations, -// representing 3D plasma quantities using Fourier decomposition in flux -// coordinates (s, \theta, \zeta) where: -// s = normalized toroidal flux (radial coordinate) -// \theta = poloidal angle -// \zeta = toroidal angle = nfp * \phi (field period toroidal angle) -// -// Physical quantities are expanded as: -// f(s,\theta,\zeta) = \sum_{m,n} f_{mn}(s) * basis_function(m*\theta, -// n*\zeta) -// -// The "FastToroidal" layout stores data with toroidal (\zeta) coordinate as -// the fast (innermost) loop index, optimizing for operations that iterate -// over toroidal modes. This differs from FastPoloidal layout. -// -// NOTE: Nestor has its own implementation of this class because we want to be -// able to use different data layouts between VMEC++ and Nestor. -// TODO(eguiraud) reduce overall code duplication -class FourierBasisFastToroidal { - public: - explicit FourierBasisFastToroidal(const Sizes* s); - - // ============================================================================ - // FOURIER BASIS SCALING FACTORS - // ============================================================================ - - // [mnyq2+1] Poloidal mode scaling factors: sqrt(2) for m>0, 1.0 for m=0 - // Applied to cos(m*\theta) and sin(m*\theta) basis functions for DFT - // normalization Enables proper normalization: 1/\pi for m>0 modes, 1/(2\pi) - // for m=0 mode - Eigen::VectorXd mscale; - - // [nnyq2+1] Toroidal mode scaling factors: sqrt(2) for n>0, 1.0 for n=0 - // Applied to cos(n*\zeta) and sin(n*\zeta) basis functions for DFT - // normalization Enables proper normalization: 1/\pi for n>0 modes, 1/(2\pi) - // for n=0 mode - Eigen::VectorXd nscale; - - // ============================================================================ - // POLOIDAL BASIS FUNCTIONS (l-major layout: [l][m]) - // ============================================================================ - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal cosine basis - // Layout: cosmu[l*(mnyq2+1) + m] = cos(m*\theta[l]) * mscale[m] - // \theta[l] = 2*\pi*l/nThetaEven for l=0...nThetaReduced-1 (reduced [0,\pi] - // interval) - Eigen::VectorXd cosmu; - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal sine basis - // Layout: sinmu[l*(mnyq2+1) + m] = sin(m*\theta[l]) * mscale[m] - Eigen::VectorXd sinmu; - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal cosine derivative - // Layout: cosmum[l*(mnyq2+1) + m] = m * cos(m*\theta[l]) * mscale[m] - // Used for computing \partial/\partial\theta derivatives in force - // calculations - Eigen::VectorXd cosmum; - - // [nThetaReduced * (mnyq2+1)] Pre-scaled poloidal sine derivative - // Layout: sinmum[l*(mnyq2+1) + m] = -m * sin(m*\theta[l]) * mscale[m] - // Used for computing \partial/\partial\theta derivatives in force - // calculations - Eigen::VectorXd sinmum; - - // ============================================================================ - // POLOIDAL BASIS WITH INTEGRATION WEIGHTS - // ============================================================================ - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal cosine basis - // Layout: cosmui[l*(mnyq2+1) + m] = cosmu[l*(mnyq2+1) + m] * intNorm - // intNorm = 1/(nZeta*(nThetaReduced-1)), with boundary point factor 1/2 - Eigen::VectorXd cosmui; - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal sine basis - // Layout: sinmui[l*(mnyq2+1) + m] = sinmu[l*(mnyq2+1) + m] * intNorm - Eigen::VectorXd sinmui; - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal cosine derivative - // Layout: cosmumi[l*(mnyq2+1) + m] = cosmum[l*(mnyq2+1) + m] * intNorm - Eigen::VectorXd cosmumi; - - // [nThetaReduced * (mnyq2+1)] Integration-weighted poloidal sine derivative - // Layout: sinmumi[l*(mnyq2+1) + m] = sinmum[l*(mnyq2+1) + m] * intNorm - Eigen::VectorXd sinmumi; - - // ============================================================================ - // TOROIDAL BASIS FUNCTIONS (n-major layout: [n][k]) - // ============================================================================ - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal cosine basis - // Layout: cosnv[n*nZeta + k] = cos(n*\zeta[k]) * nscale[n] - // \zeta[k] = 2*\pi*k/nZeta for k=0...nZeta-1 (full [0,2\pi] interval) - Eigen::VectorXd cosnv; - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal sine basis - // Layout: sinnv[n*nZeta + k] = sin(n*\zeta[k]) * nscale[n] - Eigen::VectorXd sinnv; - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal cosine derivative with nfp factor - // Layout: cosnvn[n*nZeta + k] = n*nfp * cos(n*\zeta[k]) * nscale[n] - // Factor nfp converts \partial/\partial\zeta to \partial/\partial\phi - // derivatives - Eigen::VectorXd cosnvn; - - // [(nnyq2+1) * nZeta] Pre-scaled toroidal sine derivative with nfp factor - // Layout: sinnvn[n*nZeta + k] = -n*nfp * sin(n*\zeta[k]) * nscale[n] - // Factor nfp converts \partial/\partial\zeta to \partial/\partial\phi - // derivatives - Eigen::VectorXd sinnvn; - - // ============================================================================ - // FOURIER BASIS CONVERSION FUNCTIONS - // ============================================================================ - // - // These functions convert between VMEC++'s two Fourier basis representations - // using trigonometric identities and pre-computed scaling factors. - // See docs/fourier_basis_implementation.md for complete mathematical details. - // - // Two Fourier basis types: - // 1. COMBINED BASIS (External): cos(m*\theta - n*\zeta), sin(m*\theta - - // n*\zeta) - // - Used in: wout files, Python API, traditional VMEC format - // - Storage: Linear arrays indexed by mode number mn - // - // 2. PRODUCT BASIS (Internal): cos(m*\theta)*cos(n*\zeta), - // sin(m*\theta)*sin(n*\zeta), etc. - // - Used in: Internal computations with separable DFT operations - // - Storage: 2D arrays indexed by (m,n) separately - // - Layout: fcCC[n*m_size + m] (n-major ordering for toroidal class) - // - // Mathematical basis function identity: - // cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + - // sin(m*\theta)*sin(n*\zeta) - - /** - * Convert coefficients from combined cosine basis to separable product basis. - * - * Basis function identity: - * cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + - * sin(m*\theta)*sin(n*\zeta) - * - * This function transforms coefficients for cos(m*\theta - n*\zeta) basis - * functions into coefficients for the separable product basis - * cos(m*\theta)*cos(n*\zeta) and sin(m*\theta)*sin(n*\zeta). The - * transformation accounts for VMEC symmetry where only n >= 0 coefficients - * are stored. - * - * Implementation uses pre-computed scaling factors (mscale, nscale) and - * handles positive/negative toroidal mode symmetry. Standalone function. - * - * Physics context: Converts external coefficient format (wout files) to - * internal product basis coefficients that enable separable DFT operations. - * - * @param fcCos [input] Coefficients for cos(m*\theta - n*\zeta) basis, size - * mnmax - * @param m_fcCC [output] Coefficients for cos(m*\theta)*cos(n*\zeta) basis, - * size m_size*(n_size+1) - * @param m_fcSS [output] Coefficients for sin(m*\theta)*sin(n*\zeta) basis, - * size m_size*(n_size+1) - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int cos_to_cc_ss(const std::span fcCos, - std::span m_fcCC, std::span m_fcSS, - int n_size, int m_size) const; - - /** - * Convert coefficients from combined sine basis to separable product basis. - * - * Basis function identity: - * sin(m*\theta - n*\zeta) = sin(m*\theta)*cos(n*\zeta) - - * cos(m*\theta)*sin(n*\zeta) - * - * This function transforms coefficients for sin(m*\theta - n*\zeta) basis - * functions into coefficients for the separable product basis - * sin(m*\theta)*cos(n*\zeta) and cos(m*\theta)*sin(n*\zeta). Enforces - * sin(0*\theta - 0*\zeta) = 0 constraint. - * - * Physics context: Handles sine-parity quantities like Z coordinates (zmns) - * and \lambda angle functions (lmns coefficients). - * - * @param fcSin [input] Coefficients for sin(m*\theta - n*\zeta) basis, size - * mnmax - * @param m_fcSC [output] Coefficients for sin(m*\theta)*cos(n*\zeta) basis, - * size m_size*(n_size+1) - * @param m_fcCS [output] Coefficients for cos(m*\theta)*sin(n*\zeta) basis, - * size m_size*(n_size+1) - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int sin_to_sc_cs(const std::span fcSin, - std::span m_fcSC, std::span m_fcCS, - int n_size, int m_size) const; - - /** - * Convert coefficients from separable product basis back to combined cosine - * basis. - * - * Inverse transformation using basis function identity: - * cos(m*\theta - n*\zeta) = cos(m*\theta)*cos(n*\zeta) + - * sin(m*\theta)*sin(n*\zeta) - * - * This function reconstructs coefficients for cos(m*\theta - n*\zeta) basis - * from coefficients of the separable product basis. Handles positive/negative - * toroidal mode reconstruction and applies inverse scaling factors. - * - * Physics context: Converts internal computational results back to external - * coefficient format for wout files, Python API, and traditional VMEC output. - * - * @param fcCC [input] Coefficients for cos(m*\theta)*cos(n*\zeta) basis, size - * m_size*(n_size+1) - * @param fcSS [input] Coefficients for sin(m*\theta)*sin(n*\zeta) basis, size - * m_size*(n_size+1) - * @param m_fcCos [output] Coefficients for cos(m*\theta - n*\zeta) basis, - * size mnmax - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int cc_ss_to_cos(const std::span fcCC, - const std::span fcSS, - std::span m_fcCos, int n_size, int m_size) const; - - /** - * Convert coefficients from separable product basis back to combined sine - * basis. - * - * Inverse transformation using basis function identity: - * sin(m*\theta - n*\zeta) = sin(m*\theta)*cos(n*\zeta) - - * cos(m*\theta)*sin(n*\zeta) - * - * This function reconstructs coefficients for sin(m*\theta - n*\zeta) basis - * from coefficients of the separable product basis. Enforces sin(0*\theta - - * 0*\zeta) = 0. - * - * Physics context: Converts internal results for sine-parity quantities - * back to external coefficient format for output and analysis. - * - * @param fcSC [input] Coefficients for sin(m*\theta)*cos(n*\zeta) basis, size - * m_size*(n_size+1) - * @param fcCS [input] Coefficients for cos(m*\theta)*sin(n*\zeta) basis, size - * m_size*(n_size+1) - * @param m_fcSin [output] Coefficients for sin(m*\theta - n*\zeta) basis, - * size mnmax - * @param n_size Toroidal mode range: n in [-n_size, n_size] - * @param m_size Poloidal mode range: m in [0, m_size-1] - * @return Total number of modes processed (mnmax) - */ - int sc_cs_to_sin(const std::span fcSC, - const std::span fcCS, - std::span m_fcSin, int n_size, int m_size) const; - - int mnIdx(int m, int n) const; - int mnMax(int m_size, int n_size) const; - void computeConversionIndices(Eigen::VectorXi& m_xm, Eigen::VectorXi& m_xn, - int n_size, int m_size, int nfp) const; - - // ============================================================================ - // MODE NUMBER MAPPING ARRAYS - // ============================================================================ - - // [mnmax] Poloidal mode numbers for standard resolution Fourier coefficients - // Layout: xm[mn] = poloidal mode number m for the mn-th coefficient - // Maps linear coefficient index mn to 2D mode (m,n) for spectral operations - Eigen::VectorXi xm; - - // [mnmax] Toroidal mode numbers for standard resolution Fourier coefficients - // Layout: xn[mn] = toroidal mode number n*nfp for the mn-th coefficient - // Factor nfp included to convert from field periods to geometric toroidal - // modes - Eigen::VectorXi xn; - - // [mnmax_nyq] Poloidal mode numbers for Nyquist-extended Fourier coefficients - // Layout: xm_nyq[mn] = poloidal mode number m for the mn-th Nyquist - // coefficient Extended resolution to avoid aliasing in nonlinear force - // calculations - Eigen::VectorXi xm_nyq; - - // [mnmax_nyq] Toroidal mode numbers for Nyquist-extended Fourier coefficients - // Layout: xn_nyq[mn] = toroidal mode number n*nfp for the mn-th Nyquist - // coefficient Extended resolution to avoid aliasing in nonlinear force - // calculations - Eigen::VectorXi xn_nyq; - - private: - const Sizes& s_; - - void computeFourierBasisFastToroidal(int nfp); -}; - -} // namespace vmecpp +// FourierBasisFastToroidal is the zeta-fast (n-major) layout of the shared +// FourierBasis template, which now holds the single implementation. This header +// is retained as a stable include path for existing call sites. +#include "vmecpp/common/fourier_basis/fourier_basis.h" // IWYU pragma: export #endif // VMECPP_COMMON_FOURIER_BASIS_FAST_TOROIDAL_FOURIER_BASIS_FAST_TOROIDAL_H_ diff --git a/src/vmecpp/cpp/vmecpp/common/makegrid_lib/makegrid_lib_test.cc b/src/vmecpp/cpp/vmecpp/common/makegrid_lib/makegrid_lib_test.cc index ed2952b5f..d328ee5fa 100644 --- a/src/vmecpp/cpp/vmecpp/common/makegrid_lib/makegrid_lib_test.cc +++ b/src/vmecpp/cpp/vmecpp/common/makegrid_lib/makegrid_lib_test.cc @@ -509,19 +509,23 @@ TEST_P(CheckComputeMagneticFieldResponseTable, MatchesFortranReference) { int ncid = 0; ASSERT_EQ(nc_open(p.reference_nc_file.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - EXPECT_EQ(NetcdfReadInt(ncid, "nfp"), + EXPECT_EQ(NetcdfReadInt(ncid, "nfp").value(), makegrid_parameters.number_of_field_periods); - EXPECT_EQ(NetcdfReadInt(ncid, "ir"), + EXPECT_EQ(NetcdfReadInt(ncid, "ir").value(), makegrid_parameters.number_of_r_grid_points); - EXPECT_EQ(NetcdfReadDouble(ncid, "rmin"), makegrid_parameters.r_grid_minimum); - EXPECT_EQ(NetcdfReadDouble(ncid, "rmax"), makegrid_parameters.r_grid_maximum); - EXPECT_EQ(NetcdfReadInt(ncid, "jz"), + EXPECT_EQ(NetcdfReadDouble(ncid, "rmin").value(), + makegrid_parameters.r_grid_minimum); + EXPECT_EQ(NetcdfReadDouble(ncid, "rmax").value(), + makegrid_parameters.r_grid_maximum); + EXPECT_EQ(NetcdfReadInt(ncid, "jz").value(), makegrid_parameters.number_of_z_grid_points); - EXPECT_EQ(NetcdfReadDouble(ncid, "zmin"), makegrid_parameters.z_grid_minimum); - EXPECT_EQ(NetcdfReadDouble(ncid, "zmax"), makegrid_parameters.z_grid_maximum); - EXPECT_EQ(NetcdfReadInt(ncid, "kp"), + EXPECT_EQ(NetcdfReadDouble(ncid, "zmin").value(), + makegrid_parameters.z_grid_minimum); + EXPECT_EQ(NetcdfReadDouble(ncid, "zmax").value(), + makegrid_parameters.z_grid_maximum); + EXPECT_EQ(NetcdfReadInt(ncid, "kp").value(), makegrid_parameters.number_of_phi_grid_points); - EXPECT_EQ(NetcdfReadInt(ncid, "nextcur"), number_of_serial_circuits); + EXPECT_EQ(NetcdfReadInt(ncid, "nextcur").value(), number_of_serial_circuits); for (int circuit_index = 0; circuit_index < number_of_serial_circuits; ++circuit_index) { @@ -537,11 +541,14 @@ TEST_P(CheckComputeMagneticFieldResponseTable, MatchesFortranReference) { // load mgrid data from NetCDF file std::vector>> b_r_contribution = - NetcdfReadArray3D(ncid, absl::StrFormat("br_%03d", circuit_index + 1)); + NetcdfReadArray3D(ncid, absl::StrFormat("br_%03d", circuit_index + 1)) + .value(); std::vector>> b_p_contribution = - NetcdfReadArray3D(ncid, absl::StrFormat("bp_%03d", circuit_index + 1)); + NetcdfReadArray3D(ncid, absl::StrFormat("bp_%03d", circuit_index + 1)) + .value(); std::vector>> b_z_contribution = - NetcdfReadArray3D(ncid, absl::StrFormat("bz_%03d", circuit_index + 1)); + NetcdfReadArray3D(ncid, absl::StrFormat("bz_%03d", circuit_index + 1)) + .value(); // perform comparison of points that are not explicitly excluded from the // comparison @@ -647,31 +654,31 @@ TEST_P(CheckComputeVectorPotentialCache, MatchesFortranReference) { int ncid = 0; ASSERT_EQ(nc_open(p.reference_nc_file.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - const int nfp = NetcdfReadInt(ncid, "nfp"); + const int nfp = NetcdfReadInt(ncid, "nfp").value(); EXPECT_EQ(nfp, makegrid_parameters.number_of_field_periods); - const int numR = NetcdfReadInt(ncid, "ir"); + const int numR = NetcdfReadInt(ncid, "ir").value(); EXPECT_EQ(numR, makegrid_parameters.number_of_r_grid_points); - const double minR = NetcdfReadDouble(ncid, "rmin"); + const double minR = NetcdfReadDouble(ncid, "rmin").value(); EXPECT_EQ(minR, makegrid_parameters.r_grid_minimum); - const double maxR = NetcdfReadDouble(ncid, "rmax"); + const double maxR = NetcdfReadDouble(ncid, "rmax").value(); EXPECT_EQ(maxR, makegrid_parameters.r_grid_maximum); - const int numZ = NetcdfReadInt(ncid, "jz"); + const int numZ = NetcdfReadInt(ncid, "jz").value(); EXPECT_EQ(numZ, makegrid_parameters.number_of_z_grid_points); - const double minZ = NetcdfReadDouble(ncid, "zmin"); + const double minZ = NetcdfReadDouble(ncid, "zmin").value(); EXPECT_EQ(minZ, makegrid_parameters.z_grid_minimum); - const double maxZ = NetcdfReadDouble(ncid, "zmax"); + const double maxZ = NetcdfReadDouble(ncid, "zmax").value(); EXPECT_EQ(maxZ, makegrid_parameters.z_grid_maximum); - const int numPhi = NetcdfReadInt(ncid, "kp"); + const int numPhi = NetcdfReadInt(ncid, "kp").value(); EXPECT_EQ(numPhi, makegrid_parameters.number_of_phi_grid_points); - const int nextcur = NetcdfReadInt(ncid, "nextcur"); + const int nextcur = NetcdfReadInt(ncid, "nextcur").value(); EXPECT_EQ(nextcur, number_of_serial_circuits); for (int circuit_index = 0; circuit_index < number_of_serial_circuits; @@ -688,11 +695,14 @@ TEST_P(CheckComputeVectorPotentialCache, MatchesFortranReference) { // load mgrid data from NetCDF file std::vector>> a_r_contribution = - NetcdfReadArray3D(ncid, absl::StrFormat("ar_%03d", circuit_index + 1)); + NetcdfReadArray3D(ncid, absl::StrFormat("ar_%03d", circuit_index + 1)) + .value(); std::vector>> a_p_contribution = - NetcdfReadArray3D(ncid, absl::StrFormat("ap_%03d", circuit_index + 1)); + NetcdfReadArray3D(ncid, absl::StrFormat("ap_%03d", circuit_index + 1)) + .value(); std::vector>> a_z_contribution = - NetcdfReadArray3D(ncid, absl::StrFormat("az_%03d", circuit_index + 1)); + NetcdfReadArray3D(ncid, absl::StrFormat("az_%03d", circuit_index + 1)) + .value(); // perform comparison of points that are not explicitly excluded from the // comparison diff --git a/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc b/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc index 86d15b98e..22fef164f 100644 --- a/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc +++ b/src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc @@ -1460,13 +1460,14 @@ absl::Status IsConsistent(const VmecINDATA& vmec_indata, // nothing to check here: lforbal can be true or false and both are valid... // iteration_style - // For the current state of the code, we only accept VMEC_8_52, - // but in the future [TODO(jons)] also all other (implemented) enum values - // are valid. - if (vmec_indata.iteration_style != IterationStyle::VMEC_8_52) { - return absl::InvalidArgumentError(absl::StrFormat( - "input variable 'iteration_style' must be 'vmec_8_52', but is %s\n", - ToString(vmec_indata.iteration_style))); + // VMEC_8_52 and PARVMEC are both implemented in Vmec::SolveEquilibriumLoop. + if (vmec_indata.iteration_style != IterationStyle::VMEC_8_52 && + vmec_indata.iteration_style != IterationStyle::PARVMEC) { + return absl::InvalidArgumentError( + absl::StrFormat("input variable 'iteration_style' must be 'vmec_8_52' " + "or 'parvmec', but " + "is %s\n", + ToString(vmec_indata.iteration_style))); } // return_outputs_even_if_not_converged diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/mgrid_provider/mgrid_provider.cc b/src/vmecpp/cpp/vmecpp/free_boundary/mgrid_provider/mgrid_provider.cc index e6a02b5a2..b1291d81e 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/mgrid_provider/mgrid_provider.cc +++ b/src/vmecpp/cpp/vmecpp/free_boundary/mgrid_provider/mgrid_provider.cc @@ -26,6 +26,40 @@ using netcdf_io::NetcdfReadDouble; using netcdf_io::NetcdfReadInt; using netcdf_io::NetcdfReadString; +namespace { + +absl::Status ValidateFieldContributionShape( + const std::vector > >& field_contribution, + const std::string& variable_name, int num_phi, int num_z, int num_r) { + if (field_contribution.size() != static_cast(num_phi)) { + return absl::InvalidArgumentError( + absl::StrFormat("Variable '%s' has %d phi slices, expected %d.", + variable_name, field_contribution.size(), num_phi)); + } + + for (int index_phi = 0; index_phi < num_phi; ++index_phi) { + if (field_contribution[index_phi].size() != static_cast(num_z)) { + return absl::InvalidArgumentError(absl::StrFormat( + "Variable '%s' has %d z slices at phi index %d, expected %d.", + variable_name, field_contribution[index_phi].size(), index_phi, + num_z)); + } + for (int index_z = 0; index_z < num_z; ++index_z) { + if (field_contribution[index_phi][index_z].size() != + static_cast(num_r)) { + return absl::InvalidArgumentError(absl::StrFormat( + "Variable '%s' has %d r points at phi index %d and z index %d, " + "expected %d.", + variable_name, field_contribution[index_phi][index_z].size(), + index_phi, index_z, num_r)); + } + } + } + return absl::OkStatus(); +} + +} // namespace + MGridProvider::MGridProvider() { nfp = -1; @@ -73,30 +107,71 @@ absl::Status MGridProvider::LoadFile(const std::filesystem::path& filename, filename.string())); } - // TODO(jurasic) All of these should be handled with abseil status, but - // terminate on error with absl::CHECK instead. - nfp = NetcdfReadInt(ncid, "nfp"); + // Reads below return absl::Status on failure (e.g. a missing variable), + // which we propagate to the caller instead of aborting the process. + auto with_context = [&filename](const absl::Status& s) { + return absl::Status(s.code(), + absl::StrFormat("While reading mgrid file '%s': %s", + filename.string(), s.message())); + }; + + absl::StatusOr nfp_or = NetcdfReadInt(ncid, "nfp"); + + absl::StatusOr num_r_or = NetcdfReadInt(ncid, "ir"); + absl::StatusOr min_r_or = NetcdfReadDouble(ncid, "rmin"); + absl::StatusOr max_r_or = NetcdfReadDouble(ncid, "rmax"); + + absl::StatusOr num_z_or = NetcdfReadInt(ncid, "jz"); + absl::StatusOr min_z_or = NetcdfReadDouble(ncid, "zmin"); + absl::StatusOr max_z_or = NetcdfReadDouble(ncid, "zmax"); + + absl::StatusOr num_phi_or = NetcdfReadInt(ncid, "kp"); + + absl::StatusOr nextcur_or = NetcdfReadInt(ncid, "nextcur"); + + absl::StatusOr mgrid_mode_or = + NetcdfReadString(ncid, "mgrid_mode"); + + absl::Status read_status; + read_status.Update(nfp_or.status()); + read_status.Update(num_r_or.status()); + read_status.Update(min_r_or.status()); + read_status.Update(max_r_or.status()); + read_status.Update(num_z_or.status()); + read_status.Update(min_z_or.status()); + read_status.Update(max_z_or.status()); + read_status.Update(num_phi_or.status()); + read_status.Update(nextcur_or.status()); + read_status.Update(mgrid_mode_or.status()); + if (!read_status.ok()) { + nc_close(ncid); + return with_context(read_status); + } + + nfp = *nfp_or; - numR = NetcdfReadInt(ncid, "ir"); - minR = NetcdfReadDouble(ncid, "rmin"); - maxR = NetcdfReadDouble(ncid, "rmax"); + numR = *num_r_or; + minR = *min_r_or; + maxR = *max_r_or; deltaR = (maxR - minR) / (numR - 1.0); - numZ = NetcdfReadInt(ncid, "jz"); - minZ = NetcdfReadDouble(ncid, "zmin"); - maxZ = NetcdfReadDouble(ncid, "zmax"); + numZ = *num_z_or; + minZ = *min_z_or; + maxZ = *max_z_or; deltaZ = (maxZ - minZ) / (numZ - 1.0); - numPhi = NetcdfReadInt(ncid, "kp"); + numPhi = *num_phi_or; - nextcur = NetcdfReadInt(ncid, "nextcur"); + nextcur = *nextcur_or; if (coil_currents.size() != nextcur) { + nc_close(ncid); return absl::InvalidArgumentError( absl::StrFormat("Number of currents %d does not match number of mgrid " "coil fields nextcur=%d.", coil_currents.size(), nextcur)); } - mgrid_mode = NetcdfReadString(ncid, "mgrid_mode"); + + mgrid_mode = *mgrid_mode_or; // Resize and make sure that the accumulation arrays are reset to zeros // if they contained previous contents from an earlier call to this routine. @@ -111,16 +186,44 @@ absl::Status MGridProvider::LoadFile(const std::filesystem::path& filename, // from i=1, 2, ..., nextcur std::string br_variable = absl::StrFormat("br_%03d", i + 1); - std::vector > > b_r_contribution = - NetcdfReadArray3D(ncid, br_variable); + absl::StatusOr > > > + b_r_contribution_or = NetcdfReadArray3D(ncid, br_variable); std::string bp_variable = absl::StrFormat("bp_%03d", i + 1); - std::vector > > b_p_contribution = - NetcdfReadArray3D(ncid, bp_variable); + absl::StatusOr > > > + b_p_contribution_or = NetcdfReadArray3D(ncid, bp_variable); std::string bz_variable = absl::StrFormat("bz_%03d", i + 1); + absl::StatusOr > > > + b_z_contribution_or = NetcdfReadArray3D(ncid, bz_variable); + + absl::Status contribution_status; + contribution_status.Update(b_r_contribution_or.status()); + contribution_status.Update(b_p_contribution_or.status()); + contribution_status.Update(b_z_contribution_or.status()); + if (!contribution_status.ok()) { + nc_close(ncid); + return with_context(contribution_status); + } + + std::vector > > b_r_contribution = + std::move(*b_r_contribution_or); + std::vector > > b_p_contribution = + std::move(*b_p_contribution_or); std::vector > > b_z_contribution = - NetcdfReadArray3D(ncid, bz_variable); + std::move(*b_z_contribution_or); + + absl::Status shape_status; + shape_status.Update(ValidateFieldContributionShape( + b_r_contribution, br_variable, numPhi, numZ, numR)); + shape_status.Update(ValidateFieldContributionShape( + b_p_contribution, bp_variable, numPhi, numZ, numR)); + shape_status.Update(ValidateFieldContributionShape( + b_z_contribution, bz_variable, numPhi, numZ, numR)); + if (!shape_status.ok()) { + nc_close(ncid); + return with_context(shape_status); + } for (int index_phi = 0; index_phi < numPhi; ++index_phi) { for (int index_z = 0; index_z < numZ; ++index_z) { diff --git a/src/vmecpp/cpp/vmecpp/test_data/parvmec_cth_like_fixed_bdy_force_trace.csv b/src/vmecpp/cpp/vmecpp/test_data/parvmec_cth_like_fixed_bdy_force_trace.csv new file mode 100644 index 000000000..57c74bc5a --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/test_data/parvmec_cth_like_fixed_bdy_force_trace.csv @@ -0,0 +1,125 @@ +# Per-iteration force residuals (fsqr, fsqz, fsql) from ORNL-Fusion/PARVMEC +# (PARVMEC / VMEC2000, version 10.0) for input.cth_like_fixed_bdy at ns=25, +# ftol=1e-6, single radial grid. Pins the native parvmec iteration style to the +# independent Fortran implementation. See test_parvmec_follows_ornl_parvmec_trace. +iteration,fsqr,fsqz,fsql +1,2.4089006554949474e-01,1.8189242874036761e-02,5.3233806831922206e-02 +2,3.6681393963863773e-02,2.9257078637605996e-03,2.0822882493332223e-02 +3,2.7779199334269252e-02,3.1344672791890939e-03,7.8400861928866594e-03 +4,4.5783030023920118e-02,4.9600076740968823e-03,1.1774119147540077e-02 +5,5.1533978999608258e-02,4.2652972219809670e-03,1.7993385241339890e-02 +6,5.7018736996709410e-02,4.6757668824585806e-03,1.7574015932580466e-02 +7,4.3109330291012618e-02,3.7895229757328045e-03,1.4163129388007296e-02 +8,3.5282867818118714e-02,9.3090538363044869e-03,1.4740903073900581e-02 +9,6.7512273795342706e-02,1.0963609530785181e-02,1.2272875086431585e-02 +10,2.7731565647694276e-02,5.4971626708317858e-03,7.0968082879577458e-03 +11,1.8717883628984681e-02,7.9409877631612287e-03,7.7416656727391040e-03 +12,2.5547247380531962e-02,5.1424085690450535e-03,9.7358643588169400e-03 +13,2.8975258027980587e-02,4.8055608935682372e-03,9.0112340635940102e-03 +14,2.4286783100253121e-02,5.4042619245261983e-03,6.0005852113188245e-03 +15,2.4623084138782503e-02,5.0198621981307860e-03,4.0754030267348669e-03 +16,2.5198600190923900e-02,4.5242919425412254e-03,5.0709829354750732e-03 +17,2.9336677006333411e-02,5.5715757575800348e-03,6.5578684491351099e-03 +18,3.4250965694387604e-02,5.8341711391764889e-03,5.0424903948993807e-03 +19,1.4890461945285590e-02,4.0595884422078218e-03,2.6995555783096421e-03 +20,1.0960954530469752e-02,3.6337629330389833e-03,2.4430551668922938e-03 +21,1.4952388678969699e-02,3.7276251824249918e-03,3.1916765493162943e-03 +22,1.9072944625196837e-02,4.4314470737497752e-03,3.3141794810079331e-03 +23,1.7087456203280744e-02,4.4834765275410032e-03,2.4722947963312925e-03 +24,1.0740281367609138e-02,3.0186777469101622e-03,1.9368825172277929e-03 +25,1.2840504816407781e-02,2.6206319864030322e-03,2.3728354525726779e-03 +26,2.0227894874843701e-02,3.5341442533740690e-03,2.8087096115284446e-03 +27,9.1854065263426831e-03,3.1106250041615323e-03,1.6839643143148856e-03 +28,8.1929531223409232e-03,2.4021569329149587e-03,1.4245228455584296e-03 +29,9.9063901357674174e-03,2.1780369906444432e-03,1.8151202969998671e-03 +30,8.4477400213595205e-03,2.5177223881340822e-03,1.6931479868236103e-03 +31,7.3751233951450832e-03,2.4661877740712592e-03,1.1036347771694944e-03 +32,6.5140312161359199e-03,1.6737840375596028e-03,9.0185973458750408e-04 +33,4.8665154345577666e-03,1.7609872012311961e-03,1.1547369581585613e-03 +34,5.9785995061611315e-03,2.1594629331621465e-03,1.2380368015390739e-03 +35,4.7391882750653551e-03,1.4704845093805891e-03,9.3795350370076387e-04 +36,2.7872560872948013e-03,9.4778206777475726e-04,7.4110996319325822e-04 +37,3.1287264530982461e-03,1.1977431712061541e-03,9.1157208890838633e-04 +38,3.8888019508756789e-03,1.2647909226180031e-03,1.0006383836476552e-03 +39,2.9780133015655354e-03,9.2515682752382995e-04,6.9478974931391551e-04 +40,1.7669386983827143e-03,6.4915889677512589e-04,5.1798859153854974e-04 +41,2.1048767614347910e-03,6.2215028317566545e-04,7.4684884344045977e-04 +42,2.8273637192167468e-03,7.2494305230094203e-04,8.0638310097971377e-04 +43,2.0184520142943643e-03,6.8451266247316429e-04,5.0095729470111562e-04 +44,1.2320260438745210e-03,4.5200918435776139e-04,3.4306803363647693e-04 +45,1.5633870223955200e-03,3.4189828310606429e-04,4.5906579613510926e-04 +46,1.7533603612993433e-03,4.4243567846981725e-04,4.6074948824261012e-04 +47,1.2762851718127464e-03,4.6209761253911624e-04,3.2269268416470722e-04 +48,1.0360272162643995e-03,3.1282784108317382e-04,2.8070184347688833e-04 +49,1.2739551513741842e-03,2.4669424019462221e-04,2.7789810658378851e-04 +50,1.2118834611176486e-03,2.8490012492095611e-04,2.1730078826885340e-04 +51,9.2143954616726780e-04,2.4898249862150251e-04,1.8560199983736977e-04 +52,7.7453959063069624e-04,1.7700436246373088e-04,1.9295616902615287e-04 +53,8.2803332123056606e-04,1.8260210751007069e-04,1.8170107182194307e-04 +54,8.8077280714586408e-04,2.0111185390611454e-04,1.4721096556302981e-04 +55,6.3953808033151731e-04,1.4786771753090290e-04,1.0979691075920456e-04 +56,4.6613720004616343e-04,1.0764713748773060e-04,1.1428979781776200e-04 +57,5.3783246418038480e-04,1.5356273011109318e-04,1.3334839660373454e-04 +58,5.8737697605561038e-04,1.6421612478678200e-04,1.0317001510912678e-04 +59,4.3041340767152906e-04,9.2315441238424853e-05,6.6155471587168573e-05 +60,2.7401718837970337e-04,8.4200187951912430e-05,8.4187425665565430e-05 +61,2.9377728556985981e-04,1.2112390177755989e-04,9.6228748378130450e-05 +62,3.2046565443466556e-04,8.6245427680554625e-05,5.7998596467513980e-05 +63,2.6635492112791093e-04,5.7496717496859106e-05,3.7055868372997421e-05 +64,1.8325185763783775e-04,7.6560955876649520e-05,5.7826003418891111e-05 +65,1.7193748555547836e-04,6.8973295525026561e-05,6.1390482949141541e-05 +66,1.9149982902609764e-04,4.6508665609842873e-05,3.8555178754464992e-05 +67,1.8521029634335786e-04,5.0715556091417724e-05,3.0938658432605484e-05 +68,1.4951149711365264e-04,5.6654581065966857e-05,3.6153171650008441e-05 +69,9.5066730814980337e-05,4.1790947417216183e-05,3.4730935250607754e-05 +70,1.1522448533967846e-04,3.2755314977542116e-05,3.2110748970670513e-05 +71,1.3215060028283192e-04,4.2282367000290102e-05,2.7918495552995390e-05 +72,9.5251933020723005e-05,4.0178042359465423e-05,2.3133253778786973e-05 +73,7.3717973702780459e-05,2.7737771916274311e-05,2.6037061070436951e-05 +74,7.6840052138871972e-05,3.2061872281407278e-05,2.6424430581407548e-05 +75,8.5632192316582225e-05,3.2596710010716295e-05,1.8626669170519188e-05 +76,6.6980285358956732e-05,2.2420356055828839e-05,1.7171008868068926e-05 +77,5.1597980251584379e-05,2.4106560232429674e-05,2.2873255937497923e-05 +78,4.8460133817226004e-05,2.4896744131165897e-05,2.0195807226981035e-05 +79,5.1534105677286906e-05,1.9479617993564232e-05,1.2693547255740077e-05 +80,4.7965676036113929e-05,1.8275208850858940e-05,1.3344494556237672e-05 +81,3.5000987072786026e-05,1.6746671093858607e-05,1.6774582185270918e-05 +82,3.2020293817422061e-05,1.4045305643726987e-05,1.3752620552954986e-05 +83,3.4571481241582132e-05,1.3222284564468306e-05,8.9849791597011068e-06 +84,3.6445289666553020e-05,1.3490211875571353e-05,7.6715832688247774e-06 +85,2.6836598884880545e-05,1.2228820709752092e-05,7.1047069567287085e-06 +86,2.1103703548172135e-05,9.2461388576749575e-06,6.7335117314720831e-06 +87,2.3045574529612850e-05,8.2783972746246295e-06,6.6274790604370499e-06 +88,2.0370542800839981e-05,8.7002943805473195e-06,5.2252480315205973e-06 +89,1.7977669721492016e-05,7.8672988854320640e-06,3.7936399705513430e-06 +90,1.3430524155009589e-05,6.9087201322425244e-06,3.2787527308475006e-06 +91,1.1181120201723126e-05,5.4025043317155100e-06,3.0809032723935404e-06 +92,1.3559887209373415e-05,4.2014661938091536e-06,3.4608729407792537e-06 +93,1.2746061880701487e-05,5.1789278661230076e-06,3.5352969772966303e-06 +94,9.4217905731992719e-06,4.9835523918641098e-06,2.2751170648794482e-06 +95,7.7366127289888493e-06,3.0759568914052026e-06,1.5196010009991481e-06 +96,8.3603327348265684e-06,2.9096250885540066e-06,2.1524084232659252e-06 +97,8.6955511799238599e-06,3.2835936714182820e-06,2.1679155065792652e-06 +98,6.9502126484870322e-06,2.5593521364206751e-06,1.3457562878104268e-06 +99,5.9463989880255614e-06,2.2659281857306051e-06,1.2728569415488959e-06 +100,5.5438900783405735e-06,2.2731303145953447e-06,1.5086087107421557e-06 +101,5.0693593578870755e-06,1.7912387403060682e-06,1.0948523289948467e-06 +102,4.8241421692710616e-06,1.3960215997118064e-06,8.2136369354832465e-07 +103,4.5679347253786716e-06,1.3256231291206880e-06,1.0512473810209935e-06 +104,4.1855115377915792e-06,1.2833385807810773e-06,1.0037105339563841e-06 +105,3.1489417069699166e-06,1.1025622870406821e-06,6.5276467896386802e-07 +106,3.2098909735179557e-06,9.3601268365983292e-07,5.3290041464301652e-07 +107,3.2205881914185780e-06,8.1964263628709349e-07,5.5528555685447722e-07 +108,2.6896939040987748e-06,8.9278193965618380e-07,4.9394894423187446e-07 +109,2.3377317959143821e-06,8.5659235693917731e-07,4.3174297964198452e-07 +110,1.9391127073720760e-06,6.3856890892719942e-07,3.9759929688946925e-07 +111,1.9175605674709595e-06,6.4151677568300660e-07,3.6914280286888421e-07 +112,1.5977791741740498e-06,6.3857749544403374e-07,3.5828333026535638e-07 +113,1.4921799483519888e-06,5.1292691554744418e-07,3.1973232345980027e-07 +114,1.4023514921744768e-06,4.8103063561005090e-07,2.6861565212357621e-07 +115,1.3156479553749446e-06,4.8034511015759358e-07,2.7239727527235290e-07 +116,1.4063099029231428e-06,4.6454601027385580e-07,2.6220737517779700e-07 +117,1.1147425009542418e-06,4.2478177802809816e-07,1.9162536410246975e-07 +118,1.0646051378747073e-06,3.6844629041079497e-07,1.8411084320916765e-07 +119,1.0231548922946661e-06,3.3310841632637473e-07,2.5467984006468637e-07 +120,9.3467676899647807e-07,3.1138879132238200e-07,2.6141971225702011e-07 diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel index e44538496..7f5c618d6 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel @@ -120,18 +120,19 @@ cc_binary( name = "fft_toroidal_bench", srcs = ["fft_toroidal_bench.cc"], deps = [ - ":dft_toroidal", - ":fft_toroidal", ":ideal_mhd_model", "//vmecpp/common/flow_control:flow_control", "//vmecpp/common/fourier_basis_fast_poloidal", "//vmecpp/common/sizes:sizes", + "//vmecpp/common/util:util", "//vmecpp/common/vmec_indata:vmec_indata", "//vmecpp/vmec/fourier_forces:fourier_forces", "//vmecpp/vmec/fourier_geometry:fourier_geometry", "//vmecpp/vmec/handover_storage:handover_storage", "//vmecpp/vmec/radial_partitioning:radial_partitioning", "//vmecpp/vmec/radial_profiles:radial_profiles", + "//vmecpp/vmec/thread_local_storage:thread_local_storage", + "//vmecpp/vmec/vmec_constants:vmec_constants", "@eigen", "@google_benchmark//:benchmark_main", ], 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/fft_toroidal_bench.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal_bench.cc index 3fc304976..d67b631dc 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal_bench.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/fft_toroidal_bench.cc @@ -3,98 +3,97 @@ // // SPDX-License-Identifier: MIT -// Microbenchmarks for the toroidal FFT hot loop. -// Covers both directions (spectral->real, real->spectral) at four -// representative resolutions used in typical VMEC runs. +// Microbenchmark for the toroidal transform hot loop, across a spread of +// resolutions covering both the FFT and DFT dispatch paths. // -// The parallel benchmark (BM_FourierToReal_Parallel) matches the actual VMEC -// call pattern: N threads simultaneously call -// FourierToReal3DSymmFastPoloidalFft on their own radial slice, sharing only -// the read-only ToroidalFftPlans. +// This calls IdealMhdModel::dft_FourierToReal_3d_symm() / +// dft_ForcesToFourier_3d_symm() -- the same dispatcher the real solver calls +// every iteration -- rather than the underlying FFT/DFT kernels directly. +// That dispatcher internally picks the FFTX path when a precompiled codelet +// exists for the resolution's (nZeta, 12*mpol) shape, and falls back to the +// plain DFT otherwise (see kernels_available() in fft_toroidal.h). Measuring +// through the dispatcher means a single named series here transparently +// reflects whichever path is actually active for that size, instead of +// requiring separate Dft/Fft series that must be read together. -#include - -#include #include #include -#include #include "Eigen/Dense" #include "benchmark/benchmark.h" #include "vmecpp/common/flow_control/flow_control.h" #include "vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal.h" #include "vmecpp/common/sizes/sizes.h" +#include "vmecpp/common/util/util.h" #include "vmecpp/common/vmec_indata/vmec_indata.h" #include "vmecpp/vmec/fourier_forces/fourier_forces.h" #include "vmecpp/vmec/fourier_geometry/fourier_geometry.h" #include "vmecpp/vmec/handover_storage/handover_storage.h" -#include "vmecpp/vmec/ideal_mhd_model/dft_data.h" -#include "vmecpp/vmec/ideal_mhd_model/fft_toroidal.h" #include "vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h" #include "vmecpp/vmec/radial_partitioning/radial_partitioning.h" #include "vmecpp/vmec/radial_profiles/radial_profiles.h" +#include "vmecpp/vmec/thread_local_storage/thread_local_storage.h" +#include "vmecpp/vmec/vmec_constants/vmec_constants.h" namespace vmecpp { namespace { -// ns = 51 is representative of a medium-resolution VMEC run. constexpr int kNs = 51; +// label nfp mpol ntor nZeta batch FFTX codelet? +// 4x4 1 4 4 12 48 no (small DFT baseline) +// cma_5x6 5 5 6 16 60 no (real cma config, DFT) +// 6x8 5 8 6 16 96 yes (real cma_6x8 config, FFT) +// w7x_12x12 5 12 12 28 144 yes (flagship W7-X, FFT) +// 12x13 5 12 13 30 144 no (large DFT; same batch as +// w7x_12x12 but no codelet) +struct ResParams { + int nfp; + int mpol; + int ntor; + const char* label; +}; + +constexpr ResParams kResolutions[] = { + {1, 4, 4, "4x4"}, {5, 5, 6, "cma_5x6"}, {5, 8, 6, "6x8"}, + {5, 12, 12, "w7x_12x12"}, {5, 12, 13, "12x13"}, +}; + // ---------------------------------------------------------------------------- -// Shared fixture data, built once per (nfp, mpol, ntor) combination. -// -// RadialProfiles stores pointers/references to its constructor arguments, so -// all dependencies (indata, handover, fc) must outlive the RadialProfiles -// object. We heap-allocate everything here. +// Fixture: builds a fixed-boundary IdealMhdModel and the FourierGeometry +// input it operates on. FreeBoundaryBase* is null: dft_FourierToReal_3d_symm +// and dft_ForcesToFourier_3d_symm never touch it (only referenced by the +// free-boundary vacuum path in update()/computeBContra(), which this +// benchmark never calls), and the constructor only requires a non-null +// FreeBoundaryBase when FlowControl::lfreeb is true. // ---------------------------------------------------------------------------- struct BenchFixture { - // Dependencies that RadialProfiles points into -- must be declared first. Sizes s; RadialPartitioning rp; FourierBasisFastPoloidal fb; - ToroidalFftPlans plans; - Eigen::VectorXd xmpq; + VmecConstants constants; + ThreadLocalStorage ls; std::unique_ptr indata; std::unique_ptr handover; std::unique_ptr fc; std::unique_ptr rprof; + VacuumPressureState vacuum_pressure_state = VacuumPressureState::kOff; + std::unique_ptr model; - // Forward-transform inputs. std::unique_ptr phys_x; + std::unique_ptr phys_f; - // Forward-transform output storage. - std::vector r1_e, r1_o, ru_e, ru_o, rv_e, rv_o; - std::vector z1_e, z1_o, zu_e, zu_o, zv_e, zv_o; - std::vector lu_e, lu_o, lv_e, lv_o; - std::vector rCon, zCon; - RealSpaceGeometry geom; - - // Inverse-transform inputs. - std::vector armn_e, armn_o, azmn_e, azmn_o; - std::vector blmn_e, blmn_o, brmn_e, brmn_o, bzmn_e, bzmn_o; - std::vector clmn_e, clmn_o, crmn_e, crmn_o, czmn_e, czmn_o; - std::vector frcon_e, frcon_o, fzcon_e, fzcon_o; - RealSpaceForces forces; - - // Inverse-transform output. - std::unique_ptr ff; - - explicit BenchFixture(int nfp, int mpol, int ntor, int ntheta = 0, - int nzeta = 0) - : s(/*lasym=*/false, nfp, mpol, ntor, ntheta, nzeta), + explicit BenchFixture(int nfp, int mpol, int ntor) + : s(/*lasym=*/false, nfp, mpol, ntor, /*ntheta=*/0, /*nzeta=*/0), fb(&s), - plans(s.nZeta, s.nfp, s.mpol), + ls(&s), indata(std::make_unique()), handover(std::make_unique(&s)), fc(std::make_unique(/*lfreeb=*/false, /*delt=*/0.9, /*num_grids=*/1)) { rp.adjustRadialPartitioning(/*num_threads=*/1, /*thread_id=*/0, kNs, /*lfreeb=*/false, /*printout=*/false); - - xmpq.resize(s.mpol); - for (int m = 0; m < s.mpol; ++m) xmpq[m] = m * (m - 1); - fc->ns = kNs; rprof = std::make_unique(&rp, handover.get(), indata.get(), @@ -107,8 +106,12 @@ struct BenchFixture { std::sqrt(0.05 + 0.9 * j / (nsurf > 1 ? nsurf - 1 : 1)); } - phys_x = std::make_unique(&s, &rp, kNs); + model = std::make_unique( + fc.get(), &s, &fb, rprof.get(), &constants, &ls, handover.get(), &rp, + /*m_fb=*/nullptr, /*signOfJacobian=*/-1, /*nvacskip=*/0, + &vacuum_pressure_state); + phys_x = std::make_unique(&s, &rp, kNs); std::mt19937 rng(42); std::uniform_real_distribution dist(-1.0, 1.0); auto rfill = [&](std::span sp) { @@ -121,430 +124,48 @@ struct BenchFixture { rfill(phys_x->lmnsc); rfill(phys_x->lmncs); - // Allocate forward-transform output. - const int nrzt1 = s.nZnT * (rp.nsMaxF1 - rp.nsMinF1); - const int nrzt_con = s.nZnT * (rp.nsMaxFIncludingLcfs - rp.nsMinF); - auto alloc = [](int n) { return std::vector(n, 0.0); }; - r1_e = alloc(nrzt1); - r1_o = alloc(nrzt1); - ru_e = alloc(nrzt1); - ru_o = alloc(nrzt1); - rv_e = alloc(nrzt1); - rv_o = alloc(nrzt1); - z1_e = alloc(nrzt1); - z1_o = alloc(nrzt1); - zu_e = alloc(nrzt1); - zu_o = alloc(nrzt1); - zv_e = alloc(nrzt1); - zv_o = alloc(nrzt1); - lu_e = alloc(nrzt1); - lu_o = alloc(nrzt1); - lv_e = alloc(nrzt1); - lv_o = alloc(nrzt1); - rCon = alloc(nrzt_con); - zCon = alloc(nrzt_con); - geom = - RealSpaceGeometry{r1_e, r1_o, ru_e, ru_o, rv_e, rv_o, z1_e, z1_o, zu_e, - zu_o, zv_e, zv_o, lu_e, lu_o, lv_e, lv_o, rCon, zCon}; - - // Allocate inverse-transform input. - const int nrzt = s.nZnT * (rp.nsMaxF - rp.nsMinF); - const int nrzt_lcfs = s.nZnT * (rp.nsMaxFIncludingLcfs - rp.nsMinF); - auto rvec = [&](int n) { - std::vector v(n); - for (double& x : v) x = dist(rng); - return v; - }; - armn_e = rvec(nrzt); - armn_o = rvec(nrzt); - azmn_e = rvec(nrzt); - azmn_o = rvec(nrzt); - blmn_e = rvec(nrzt_lcfs); - blmn_o = rvec(nrzt_lcfs); - brmn_e = rvec(nrzt); - brmn_o = rvec(nrzt); - bzmn_e = rvec(nrzt); - bzmn_o = rvec(nrzt); - clmn_e = rvec(nrzt_lcfs); - clmn_o = rvec(nrzt_lcfs); - crmn_e = rvec(nrzt); - crmn_o = rvec(nrzt); - czmn_e = rvec(nrzt); - czmn_o = rvec(nrzt); - frcon_e = rvec(nrzt); - frcon_o = rvec(nrzt); - fzcon_e = rvec(nrzt); - fzcon_o = rvec(nrzt); - forces = RealSpaceForces{armn_e, armn_o, azmn_e, azmn_o, blmn_e, - blmn_o, brmn_e, brmn_o, bzmn_e, bzmn_o, - clmn_e, clmn_o, crmn_e, crmn_o, czmn_e, - czmn_o, frcon_e, frcon_o, fzcon_e, fzcon_o}; - - ff = std::make_unique(&s, &rp, kNs); + phys_f = std::make_unique(&s, &rp, kNs); } }; -// ---------------------------------------------------------------------------- -// Benchmark helpers -// ---------------------------------------------------------------------------- - -// (nfp, mpol, ntor) pairs for the four benchmark resolutions. -struct ResParams { - int nfp; - int mpol; - int ntor; - const char* label; -}; - -constexpr ResParams kResolutions[] = { - {1, 4, 4, "4x4"}, - {1, 7, 1, "7x1"}, - {5, 12, 12, "12x12"}, - {5, 16, 18, "16x18"}, -}; - -// Templated benchmarks parameterised by a ResParams index so GBench can name -// them clearly. The fixture is a function-local static so it is built exactly -// once per process (not once per benchmark iteration). - -template -void BM_FourierToReal(benchmark::State& state) { - static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, - kResolutions[kIdx].ntor); - // The FFT path is only valid when vendored FFTX codelets exist for this - // (nZeta, 12*mpol) shape; otherwise fftx_full_c2r_run is null and calling it - // would dereference a null function pointer. Skip cleanly in that case (the - // real solver falls back to the DFT path via kernels_available()). - if (!fx.plans.kernels_available()) { - state.SkipWithError("no FFTX codelet for this resolution"); - return; - } - for (auto _ : state) { - FourierToReal3DSymmFastPoloidalFft(*fx.phys_x, fx.xmpq, fx.rp, fx.s, - *fx.rprof, fx.fb, fx.plans, fx.geom); - benchmark::ClobberMemory(); - } - state.SetLabel(kResolutions[kIdx].label); -} - template -void BM_ForcesToFourier(benchmark::State& state) { +void BM_ToroidalTransform_FourierToReal(benchmark::State& state) { static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, kResolutions[kIdx].ntor); - if (!fx.plans.kernels_available()) { - state.SkipWithError("no FFTX codelet for this resolution"); - return; - } for (auto _ : state) { - ForcesToFourier3DSymmFastPoloidalFft(fx.forces, fx.xmpq, fx.rp, *fx.fc, - fx.s, fx.fb, fx.plans, - VacuumPressureState::kOff, *fx.ff); + fx.model->dft_FourierToReal_3d_symm(*fx.phys_x); benchmark::ClobberMemory(); } state.SetLabel(kResolutions[kIdx].label); } template -void BM_DftFourierToReal(benchmark::State& state) { +void BM_ToroidalTransform_ForcesToFourier(benchmark::State& state) { static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, kResolutions[kIdx].ntor); + // dft_ForcesToFourier_3d_symm reads IdealMhdModel's private real-space + // force members (armn_e, blmn_e, ...), which this fixture never populates + // -- they stay zero-initialized. That's fine for a timing benchmark: the + // transform cost doesn't depend on the input values, only its size. for (auto _ : state) { - FourierToReal3DSymmFastPoloidal(*fx.phys_x, fx.xmpq, fx.rp, fx.s, *fx.rprof, - fx.fb, fx.geom); + fx.model->dft_ForcesToFourier_3d_symm(*fx.phys_f); benchmark::ClobberMemory(); } state.SetLabel(kResolutions[kIdx].label); } -BENCHMARK_TEMPLATE(BM_DftFourierToReal, 0)->Name("DftFourierToReal/4x4"); -BENCHMARK_TEMPLATE(BM_FourierToReal, 0)->Name("FftFourierToReal/4x4"); -BENCHMARK_TEMPLATE(BM_DftFourierToReal, 1)->Name("DftFourierToReal/7x1"); -BENCHMARK_TEMPLATE(BM_FourierToReal, 1)->Name("FftFourierToReal/7x1"); -BENCHMARK_TEMPLATE(BM_DftFourierToReal, 2)->Name("DftFourierToReal/12x12"); -BENCHMARK_TEMPLATE(BM_FourierToReal, 2)->Name("FftFourierToReal/12x12"); -BENCHMARK_TEMPLATE(BM_DftFourierToReal, 3)->Name("DftFourierToReal/16x18"); -BENCHMARK_TEMPLATE(BM_FourierToReal, 3)->Name("FftFourierToReal/16x18"); - -// ---------------------------------------------------------------------------- -// Real-space resolution sweep at fixed spectral resolution (mpol=12, ntor=12, -// nfp=5). Default real-space grid is (ntheta=2*mpol+6=30, nzeta=2*ntor+4=28 -// when nfp=1; for nfp=5 the toroidal grid extends accordingly). We vary -// ntheta and nzeta independently to isolate poloidal-AXPY cost (ntheta) from -// toroidal-FFT cost (nzeta). -// ---------------------------------------------------------------------------- - -struct RealSpaceParams { - int nfp; - int mpol; - int ntor; - int ntheta; // 0 = use Sizes default - int nzeta; // 0 = use Sizes default - const char* label; -}; - -constexpr RealSpaceParams kRealSpaceSweep[] = { - // Vary nzeta (toroidal real-space grid) at fixed ntheta default. - {5, 12, 12, 0, 28, "12x12_ntheta-default_nzeta-28"}, // baseline default - {5, 12, 12, 0, 56, "12x12_ntheta-default_nzeta-56"}, // 2x toroidal - {5, 12, 12, 0, 84, "12x12_ntheta-default_nzeta-84"}, // 3x toroidal - {5, 12, 12, 0, 112, "12x12_ntheta-default_nzeta-112"}, // 4x toroidal - // Vary ntheta (poloidal real-space grid) at fixed nzeta default. - {5, 12, 12, 30, 0, "12x12_ntheta-30_nzeta-default"}, // baseline default - {5, 12, 12, 60, 0, "12x12_ntheta-60_nzeta-default"}, // 2x poloidal - {5, 12, 12, 90, 0, "12x12_ntheta-90_nzeta-default"}, // 3x poloidal - {5, 12, 12, 120, 0, "12x12_ntheta-120_nzeta-default"}, // 4x poloidal -}; - -template -void BM_FftFourierToReal_RealSpace(benchmark::State& state) { - static BenchFixture fx(kRealSpaceSweep[kIdx].nfp, kRealSpaceSweep[kIdx].mpol, - kRealSpaceSweep[kIdx].ntor, - kRealSpaceSweep[kIdx].ntheta, - kRealSpaceSweep[kIdx].nzeta); - if (!fx.plans.kernels_available()) { - state.SkipWithError("no FFTX codelet for this resolution"); - return; - } - for (auto _ : state) { - FourierToReal3DSymmFastPoloidalFft(*fx.phys_x, fx.xmpq, fx.rp, fx.s, - *fx.rprof, fx.fb, fx.plans, fx.geom); - benchmark::ClobberMemory(); - } - state.SetLabel(kRealSpaceSweep[kIdx].label); -} - -template -void BM_DftFourierToReal_RealSpace(benchmark::State& state) { - static BenchFixture fx(kRealSpaceSweep[kIdx].nfp, kRealSpaceSweep[kIdx].mpol, - kRealSpaceSweep[kIdx].ntor, - kRealSpaceSweep[kIdx].ntheta, - kRealSpaceSweep[kIdx].nzeta); - for (auto _ : state) { - FourierToReal3DSymmFastPoloidal(*fx.phys_x, fx.xmpq, fx.rp, fx.s, *fx.rprof, - fx.fb, fx.geom); - benchmark::ClobberMemory(); - } - state.SetLabel(kRealSpaceSweep[kIdx].label); -} +#define REGISTER_RES(IDX, LABEL) \ + BENCHMARK_TEMPLATE(BM_ToroidalTransform_FourierToReal, IDX) \ + ->Name("ToroidalFourierToReal/" LABEL); \ + BENCHMARK_TEMPLATE(BM_ToroidalTransform_ForcesToFourier, IDX) \ + ->Name("ToroidalForcesToFourier/" LABEL) -#define REGISTER_RS(I, NAME) \ - BENCHMARK_TEMPLATE(BM_DftFourierToReal_RealSpace, I)->Name("Dft/" NAME); \ - BENCHMARK_TEMPLATE(BM_FftFourierToReal_RealSpace, I)->Name("Fft/" NAME) +REGISTER_RES(0, "4x4"); +REGISTER_RES(1, "6x8"); +REGISTER_RES(2, "12x12"); +REGISTER_RES(3, "12x13"); -REGISTER_RS(0, "12x12_nzeta-28"); -REGISTER_RS(1, "12x12_nzeta-56"); -REGISTER_RS(2, "12x12_nzeta-84"); -REGISTER_RS(3, "12x12_nzeta-112"); -REGISTER_RS(4, "12x12_ntheta-30"); -REGISTER_RS(5, "12x12_ntheta-60"); -REGISTER_RS(6, "12x12_ntheta-90"); -REGISTER_RS(7, "12x12_ntheta-120"); - -#undef REGISTER_RS - -// ---------------------------------------------------------------------------- -// Parallel fixture: one BenchFixture per thread, each covering its own radial -// slice, sharing the same ToroidalFftPlans (read-only during the hot loop). -// This matches the real VMEC calling pattern exactly: -// #pragma omp parallel -// { models[omp_get_thread_num()].geometryFromFourier(phys_x); } -// ---------------------------------------------------------------------------- - -struct ParallelBenchFixture { - // Shared across threads (read-only during hot loop). - Sizes s; - FourierBasisFastPoloidal fb; - ToroidalFftPlans plans; - Eigen::VectorXd xmpq; - - // Per-thread state: each thread gets its own radial slice. - struct ThreadSlice { - RadialPartitioning rp; - std::unique_ptr indata; - std::unique_ptr handover; - std::unique_ptr fc; - std::unique_ptr rprof; - std::unique_ptr phys_x; - // Output buffers (each thread writes only to its own slice). - std::vector r1_e, r1_o, ru_e, ru_o, rv_e, rv_o; - std::vector z1_e, z1_o, zu_e, zu_o, zv_e, zv_o; - std::vector lu_e, lu_o, lv_e, lv_o; - std::vector rCon, zCon; - RealSpaceGeometry geom; - }; - - std::vector threads; - - explicit ParallelBenchFixture(int nfp, int mpol, int ntor, int num_threads) - : s(/*lasym=*/false, nfp, mpol, ntor, /*ntheta=*/0, /*nzeta=*/0), - fb(&s), - plans(s.nZeta, s.nfp, s.mpol), - threads(num_threads) { - xmpq.resize(s.mpol); - for (int m = 0; m < s.mpol; ++m) xmpq[m] = m * (m - 1); - - std::mt19937 rng(42); - std::uniform_real_distribution dist(-1.0, 1.0); - - for (int t = 0; t < num_threads; ++t) { - ThreadSlice& sl = threads[t]; - sl.rp.adjustRadialPartitioning(num_threads, t, kNs, - /*lfreeb=*/false, /*printout=*/false); - sl.indata = std::make_unique(); - sl.handover = std::make_unique(&s); - sl.fc = std::make_unique(/*lfreeb=*/false, /*delt=*/0.9, - /*num_grids=*/1); - sl.fc->ns = kNs; - sl.rprof = std::make_unique( - &sl.rp, sl.handover.get(), sl.indata.get(), sl.fc.get(), - /*signOfJacobian=*/-1, /*pDamp=*/0.05); - const int nsurf = sl.rp.nsMaxF1 - sl.rp.nsMinF1; - sl.rprof->sqrtSF.resize(nsurf); - for (int j = 0; j < nsurf; ++j) { - sl.rprof->sqrtSF[j] = - std::sqrt(0.05 + 0.9 * j / (nsurf > 1 ? nsurf - 1 : 1)); - } - - sl.phys_x = std::make_unique(&s, &sl.rp, kNs); - auto rfill = [&](std::span sp) { - for (double& x : sp) x = dist(rng); - }; - rfill(sl.phys_x->rmncc); - rfill(sl.phys_x->rmnss); - rfill(sl.phys_x->zmnsc); - rfill(sl.phys_x->zmncs); - rfill(sl.phys_x->lmnsc); - rfill(sl.phys_x->lmncs); - - const int nrzt1 = s.nZnT * (sl.rp.nsMaxF1 - sl.rp.nsMinF1); - const int nrzt_con = s.nZnT * (sl.rp.nsMaxFIncludingLcfs - sl.rp.nsMinF); - auto alloc = [](int n) { return std::vector(n, 0.0); }; - sl.r1_e = alloc(nrzt1); - sl.r1_o = alloc(nrzt1); - sl.ru_e = alloc(nrzt1); - sl.ru_o = alloc(nrzt1); - sl.rv_e = alloc(nrzt1); - sl.rv_o = alloc(nrzt1); - sl.z1_e = alloc(nrzt1); - sl.z1_o = alloc(nrzt1); - sl.zu_e = alloc(nrzt1); - sl.zu_o = alloc(nrzt1); - sl.zv_e = alloc(nrzt1); - sl.zv_o = alloc(nrzt1); - sl.lu_e = alloc(nrzt1); - sl.lu_o = alloc(nrzt1); - sl.lv_e = alloc(nrzt1); - sl.lv_o = alloc(nrzt1); - sl.rCon = alloc(nrzt_con); - sl.zCon = alloc(nrzt_con); - sl.geom = RealSpaceGeometry{sl.r1_e, sl.r1_o, sl.ru_e, sl.ru_o, sl.rv_e, - sl.rv_o, sl.z1_e, sl.z1_o, sl.zu_e, sl.zu_o, - sl.zv_e, sl.zv_o, sl.lu_e, sl.lu_o, sl.lv_e, - sl.lv_o, sl.rCon, sl.zCon}; - } - } -}; - -// w7x at 4 threads: matches the real VMEC calling pattern exactly. -// -// VMEC keeps a single #pragma omp parallel team alive for the entire solver -// run. Each thread owns its own IdealMhdModel (with its own RadialPartitioning -// slice) and calls dft_FourierToReal_3d_symm directly from within that -// persistent team -- there is no nested fork/join per iteration, only -// #pragma omp barrier between phases. -// -// We replicate that by opening one persistent team and looping inside it. -// Thread 0 drives the Google Benchmark state machine; the others mirror it -// via a shared atomic flag. Each call is bracketed by barriers so all threads -// start and finish together, and the wall-clock time reflects the slowest. -void BM_FourierToReal_Parallel_W7x_4t(benchmark::State& state) { - constexpr int kNumThreads = 6; - static ParallelBenchFixture fx(/*nfp=*/5, /*mpol=*/12, /*ntor=*/12, - kNumThreads); - if (!fx.plans.kernels_available()) { - state.SkipWithError("no FFTX codelet for this resolution"); - return; - } - - std::atomic keep_going{true}; - -#pragma omp parallel num_threads(kNumThreads) - { - const int tid = omp_get_thread_num(); - ParallelBenchFixture::ThreadSlice& sl = fx.threads[tid]; - - // Warmup: each thread faults in its own plan's twiddle pages and sizes - // the thread_local scratch buffer before timing starts. - FourierToReal3DSymmFastPoloidalFft(*sl.phys_x, fx.xmpq, sl.rp, fx.s, - *sl.rprof, fx.fb, fx.plans, sl.geom); -#pragma omp barrier - - if (tid == 0) { - for (auto _ : state) { -#pragma omp barrier - FourierToReal3DSymmFastPoloidalFft(*sl.phys_x, fx.xmpq, sl.rp, fx.s, - *sl.rprof, fx.fb, fx.plans, sl.geom); -#pragma omp barrier - } - keep_going.store(false, std::memory_order_release); - // Final barrier so workers see the flag and exit cleanly. -#pragma omp barrier - } else { - while (true) { -#pragma omp barrier - if (!keep_going.load(std::memory_order_acquire)) break; - FourierToReal3DSymmFastPoloidalFft(*sl.phys_x, fx.xmpq, sl.rp, fx.s, - *sl.rprof, fx.fb, fx.plans, sl.geom); -#pragma omp barrier - } - } - } - - state.SetLabel("12x12 6-thread parallel fft"); -} -BENCHMARK(BM_FourierToReal_Parallel_W7x_4t); - -// DFT parallel: same structure, no FFT plans needed. -void BM_DftFourierToReal_Parallel_W7x_4t(benchmark::State& state) { - constexpr int kNumThreads = 6; - static ParallelBenchFixture fx(/*nfp=*/5, /*mpol=*/12, /*ntor=*/12, - kNumThreads); - - std::atomic keep_going{true}; - -#pragma omp parallel num_threads(kNumThreads) - { - const int tid = omp_get_thread_num(); - ParallelBenchFixture::ThreadSlice& sl = fx.threads[tid]; - - // Warmup: fault in thread_local memory / cache lines before timing. - FourierToReal3DSymmFastPoloidal(*sl.phys_x, fx.xmpq, sl.rp, fx.s, *sl.rprof, - fx.fb, sl.geom); -#pragma omp barrier - - if (tid == 0) { - for (auto _ : state) { -#pragma omp barrier - FourierToReal3DSymmFastPoloidal(*sl.phys_x, fx.xmpq, sl.rp, fx.s, - *sl.rprof, fx.fb, sl.geom); -#pragma omp barrier - } - keep_going.store(false, std::memory_order_release); -#pragma omp barrier - } else { - while (true) { -#pragma omp barrier - if (!keep_going.load(std::memory_order_acquire)) break; - FourierToReal3DSymmFastPoloidal(*sl.phys_x, fx.xmpq, sl.rp, fx.s, - *sl.rprof, fx.fb, sl.geom); -#pragma omp barrier - } - } - } - - state.SetLabel("12x12 6-thread parallel dft"); -} -BENCHMARK(BM_DftFourierToReal_Parallel_W7x_4t); +#undef REGISTER_RES } // namespace } // namespace vmecpp 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/output_quantities/output_quantities_test.cc b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_test.cc index e696b52e4..44ecb7f86 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_test.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_test.cc @@ -88,15 +88,15 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { int ncid; ASSERT_EQ(nc_open(filename.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - EXPECT_EQ(wout.signgs, NetcdfReadInt(ncid, "signgs")); + EXPECT_EQ(wout.signgs, NetcdfReadInt(ncid, "signgs").value()); - EXPECT_EQ(wout.gamma, NetcdfReadDouble(ncid, "gamma")); + EXPECT_EQ(wout.gamma, NetcdfReadDouble(ncid, "gamma").value()); - EXPECT_EQ(wout.pcurr_type, NetcdfReadString(ncid, "pcurr_type")); - EXPECT_EQ(wout.pmass_type, NetcdfReadString(ncid, "pmass_type")); - EXPECT_EQ(wout.piota_type, NetcdfReadString(ncid, "piota_type")); + EXPECT_EQ(wout.pcurr_type, NetcdfReadString(ncid, "pcurr_type").value()); + EXPECT_EQ(wout.pmass_type, NetcdfReadString(ncid, "pmass_type").value()); + EXPECT_EQ(wout.piota_type, NetcdfReadString(ncid, "piota_type").value()); - std::vector reference_am = NetcdfReadArray1D(ncid, "am"); + std::vector reference_am = NetcdfReadArray1D(ncid, "am").value(); // remove zero-padding at end reference_am.resize(wout.am.size()); EXPECT_THAT(wout.am, ElementsAreArray(reference_am)); @@ -105,37 +105,37 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { if (vmec_indata->ncurr == 0) { // constrained-iota; ignore current profile coefficients // TODO(jons): check for spline profiles -> need to check ai_aux_* - std::vector reference_ai = NetcdfReadArray1D(ncid, "ai"); + std::vector reference_ai = NetcdfReadArray1D(ncid, "ai").value(); // remove zero-padding at end reference_ai.resize(wout.ai.size()); EXPECT_THAT(wout.ai, ElementsAreArray(reference_ai)); } else { // constrained-current // TODO(jons): check for spline profiles -> need to check ac_aux_* - std::vector reference_ac = NetcdfReadArray1D(ncid, "ac"); + std::vector reference_ac = NetcdfReadArray1D(ncid, "ac").value(); reference_ac.resize(wout.ac.size()); EXPECT_THAT(wout.ac, ElementsAreArray(reference_ac)); if (wout.ai.size() > 0) { // iota profile (if present) taken as initial guess for first iteration // TODO(jons): check for spline profiles -> need to check ai_aux_* - std::vector reference_ai = NetcdfReadArray1D(ncid, "ai"); + std::vector reference_ai = NetcdfReadArray1D(ncid, "ai").value(); // remove zero-padding at end reference_ai.resize(wout.ai.size()); EXPECT_THAT(wout.ai, ElementsAreArray(reference_ai)); } } - EXPECT_EQ(wout.nfp, NetcdfReadInt(ncid, "nfp")); - EXPECT_EQ(wout.mpol, NetcdfReadInt(ncid, "mpol")); - EXPECT_EQ(wout.ntor, NetcdfReadInt(ncid, "ntor")); - EXPECT_EQ(wout.lasym, NetcdfReadBool(ncid, "lasym")); + EXPECT_EQ(wout.nfp, NetcdfReadInt(ncid, "nfp").value()); + EXPECT_EQ(wout.mpol, NetcdfReadInt(ncid, "mpol").value()); + EXPECT_EQ(wout.ntor, NetcdfReadInt(ncid, "ntor").value()); + EXPECT_EQ(wout.lasym, NetcdfReadBool(ncid, "lasym").value()); - EXPECT_EQ(wout.ns, NetcdfReadInt(ncid, "ns")); - EXPECT_EQ(wout.ftolv, NetcdfReadDouble(ncid, "ftolv")); - EXPECT_EQ(wout.niter, NetcdfReadInt(ncid, "niter")); + EXPECT_EQ(wout.ns, NetcdfReadInt(ncid, "ns").value()); + EXPECT_EQ(wout.ftolv, NetcdfReadDouble(ncid, "ftolv").value()); + EXPECT_EQ(wout.niter, NetcdfReadInt(ncid, "niter").value()); - EXPECT_EQ(wout.lfreeb, NetcdfReadBool(ncid, "lfreeb")); + EXPECT_EQ(wout.lfreeb, NetcdfReadBool(ncid, "lfreeb").value()); if (wout.lfreeb) { // The reference data is generated using educational_VMEC, // which is run from within //vmecpp/test_data. @@ -147,88 +147,99 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { // //vmecpp/test_data/regenerate_test_data.sh. EXPECT_EQ(wout.mgrid_file, absl::StrCat("vmecpp/test_data/", - NetcdfReadString(ncid, "mgrid_file"))); - std::vector reference_extcur = NetcdfReadArray1D(ncid, "extcur"); + NetcdfReadString(ncid, "mgrid_file").value())); + std::vector reference_extcur = + NetcdfReadArray1D(ncid, "extcur").value(); EXPECT_THAT(wout.extcur, ElementsAreArray(reference_extcur)); EXPECT_EQ(wout.nextcur, static_cast(reference_extcur.size())); } - EXPECT_EQ(wout.mgrid_mode, NetcdfReadString(ncid, "mgrid_mode")); + EXPECT_EQ(wout.mgrid_mode, NetcdfReadString(ncid, "mgrid_mode").value()); // ------------------- // scalar quantities - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "wb"), wout.wb, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "wp"), wout.wp, tolerance)); + EXPECT_TRUE( + IsCloseRelAbs(NetcdfReadDouble(ncid, "wb").value(), wout.wb, tolerance)); + EXPECT_TRUE( + IsCloseRelAbs(NetcdfReadDouble(ncid, "wp").value(), wout.wp, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "rmax_surf"), wout.rmax_surf, - tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "rmin_surf"), wout.rmin_surf, - tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "zmax_surf"), wout.zmax_surf, - tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "rmax_surf").value(), + wout.rmax_surf, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "rmin_surf").value(), + wout.rmin_surf, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "zmax_surf").value(), + wout.zmax_surf, tolerance)); - EXPECT_EQ(wout.mnmax, NetcdfReadInt(ncid, "mnmax")); - EXPECT_EQ(wout.mnmax_nyq, NetcdfReadInt(ncid, "mnmax_nyq")); + EXPECT_EQ(wout.mnmax, NetcdfReadInt(ncid, "mnmax").value()); + EXPECT_EQ(wout.mnmax_nyq, NetcdfReadInt(ncid, "mnmax_nyq").value()); - EXPECT_EQ(wout.ier_flag, NetcdfReadInt(ncid, "ier_flag")); + EXPECT_EQ(wout.ier_flag, NetcdfReadInt(ncid, "ier_flag").value()); - EXPECT_TRUE( - IsCloseRelAbs(NetcdfReadDouble(ncid, "aspect"), wout.aspect, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "aspect").value(), + wout.aspect, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betatotal"), wout.betatotal, - tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betapol"), wout.betapol, - tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betator"), wout.betator, - tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betaxis"), wout.betaxis, - tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betatotal").value(), + wout.betatotal, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betapol").value(), + wout.betapol, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betator").value(), + wout.betator, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "betaxis").value(), + wout.betaxis, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "b0"), wout.b0, tolerance)); - - EXPECT_TRUE( - IsCloseRelAbs(NetcdfReadDouble(ncid, "rbtor0"), wout.rbtor0, tolerance)); EXPECT_TRUE( - IsCloseRelAbs(NetcdfReadDouble(ncid, "rbtor"), wout.rbtor, tolerance)); + IsCloseRelAbs(NetcdfReadDouble(ncid, "b0").value(), wout.b0, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "IonLarmor"), wout.IonLarmor, + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "rbtor0").value(), + wout.rbtor0, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "rbtor").value(), wout.rbtor, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "volavgB"), wout.volavgB, + + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "IonLarmor").value(), + wout.IonLarmor, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "volavgB").value(), + wout.volavgB, tolerance)); + + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "ctor").value(), wout.ctor, tolerance)); - EXPECT_TRUE( - IsCloseRelAbs(NetcdfReadDouble(ncid, "ctor"), wout.ctor, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "Aminor_p").value(), + wout.Aminor_p, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "Rmajor_p").value(), + wout.Rmajor_p, tolerance)); + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "volume_p").value(), + wout.volume, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "Aminor_p"), wout.Aminor_p, + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "fsqr").value(), wout.fsqr, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "Rmajor_p"), wout.Rmajor_p, + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "fsqz").value(), wout.fsqz, tolerance)); - EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "volume_p"), wout.volume, + EXPECT_TRUE(IsCloseRelAbs(NetcdfReadDouble(ncid, "fsql").value(), wout.fsql, tolerance)); - EXPECT_TRUE( - IsCloseRelAbs(NetcdfReadDouble(ncid, "fsqr"), wout.fsqr, tolerance)); - EXPECT_TRUE( - IsCloseRelAbs(NetcdfReadDouble(ncid, "fsqz"), wout.fsqz, tolerance)); - EXPECT_TRUE( - IsCloseRelAbs(NetcdfReadDouble(ncid, "fsql"), wout.fsql, tolerance)); - // ------------------- // one-dimensional array quantities - std::vector reference_iota_full = NetcdfReadArray1D(ncid, "iotaf"); + std::vector reference_iota_full = + NetcdfReadArray1D(ncid, "iotaf").value(); std::vector reference_safety_factor = - NetcdfReadArray1D(ncid, "q_factor"); + NetcdfReadArray1D(ncid, "q_factor").value(); std::vector reference_pressure_full = - NetcdfReadArray1D(ncid, "presf"); - std::vector reference_toroidal_flux = NetcdfReadArray1D(ncid, "phi"); - std::vector reference_poloidal_flux = NetcdfReadArray1D(ncid, "chi"); - std::vector reference_phipf = NetcdfReadArray1D(ncid, "phipf"); - std::vector reference_chipf = NetcdfReadArray1D(ncid, "chipf"); - std::vector reference_jcuru = NetcdfReadArray1D(ncid, "jcuru"); - std::vector reference_jcurv = NetcdfReadArray1D(ncid, "jcurv"); + NetcdfReadArray1D(ncid, "presf").value(); + std::vector reference_toroidal_flux = + NetcdfReadArray1D(ncid, "phi").value(); + std::vector reference_poloidal_flux = + NetcdfReadArray1D(ncid, "chi").value(); + std::vector reference_phipf = + NetcdfReadArray1D(ncid, "phipf").value(); + std::vector reference_chipf = + NetcdfReadArray1D(ncid, "chipf").value(); + std::vector reference_jcuru = + NetcdfReadArray1D(ncid, "jcuru").value(); + std::vector reference_jcurv = + NetcdfReadArray1D(ncid, "jcurv").value(); std::vector reference_spectral_width = - NetcdfReadArray1D(ncid, "specw"); + NetcdfReadArray1D(ncid, "specw").value(); for (int jF = 0; jF < fc.ns; ++jF) { EXPECT_TRUE( IsCloseRelAbs(reference_iota_full[jF], wout.iotaf[jF], tolerance)); @@ -248,15 +259,21 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { IsCloseRelAbs(reference_spectral_width[jF], wout.specw[jF], tolerance)); } // jF - std::vector reference_iota_half = NetcdfReadArray1D(ncid, "iotas"); - std::vector reference_mass_half = NetcdfReadArray1D(ncid, "mass"); - std::vector reference_pressure_half = NetcdfReadArray1D(ncid, "pres"); - std::vector reference_beta = NetcdfReadArray1D(ncid, "beta_vol"); - std::vector reference_buco = NetcdfReadArray1D(ncid, "buco"); - std::vector reference_bvco = NetcdfReadArray1D(ncid, "bvco"); - std::vector reference_dVds = NetcdfReadArray1D(ncid, "vp"); - std::vector reference_phips = NetcdfReadArray1D(ncid, "phips"); - std::vector reference_overr = NetcdfReadArray1D(ncid, "over_r"); + std::vector reference_iota_half = + NetcdfReadArray1D(ncid, "iotas").value(); + std::vector reference_mass_half = + NetcdfReadArray1D(ncid, "mass").value(); + std::vector reference_pressure_half = + NetcdfReadArray1D(ncid, "pres").value(); + std::vector reference_beta = + NetcdfReadArray1D(ncid, "beta_vol").value(); + std::vector reference_buco = NetcdfReadArray1D(ncid, "buco").value(); + std::vector reference_bvco = NetcdfReadArray1D(ncid, "bvco").value(); + std::vector reference_dVds = NetcdfReadArray1D(ncid, "vp").value(); + std::vector reference_phips = + NetcdfReadArray1D(ncid, "phips").value(); + std::vector reference_overr = + NetcdfReadArray1D(ncid, "over_r").value(); for (int jF = 0; jF < fc.ns; ++jF) { EXPECT_TRUE( IsCloseRelAbs(reference_iota_half[jF], wout.iotas[jF], tolerance)); @@ -273,10 +290,12 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { EXPECT_TRUE(IsCloseRelAbs(reference_overr[jF], wout.over_r[jF], tolerance)); } // jF - std::vector reference_jdotb = NetcdfReadArray1D(ncid, "jdotb"); - std::vector reference_bdotb = NetcdfReadArray1D(ncid, "bdotb"); + std::vector reference_jdotb = + NetcdfReadArray1D(ncid, "jdotb").value(); + std::vector reference_bdotb = + NetcdfReadArray1D(ncid, "bdotb").value(); std::vector reference_bdotgradv = - NetcdfReadArray1D(ncid, "bdotgradv"); + NetcdfReadArray1D(ncid, "bdotgradv").value(); for (int jF = 0; jF < fc.ns; ++jF) { EXPECT_TRUE( IsCloseRelAbs(reference_jdotb[jF], wout.jdotb[jF], 10 * tolerance)); @@ -285,11 +304,16 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { IsCloseRelAbs(reference_bdotgradv[jF], wout.bdotgradv[jF], tolerance)); } // jF - std::vector reference_DMerc = NetcdfReadArray1D(ncid, "DMerc"); - std::vector reference_Dshear = NetcdfReadArray1D(ncid, "DShear"); - std::vector reference_Dwell = NetcdfReadArray1D(ncid, "DWell"); - std::vector reference_Dcurr = NetcdfReadArray1D(ncid, "DCurr"); - std::vector reference_Dgeod = NetcdfReadArray1D(ncid, "DGeod"); + std::vector reference_DMerc = + NetcdfReadArray1D(ncid, "DMerc").value(); + std::vector reference_Dshear = + NetcdfReadArray1D(ncid, "DShear").value(); + std::vector reference_Dwell = + NetcdfReadArray1D(ncid, "DWell").value(); + std::vector reference_Dcurr = + NetcdfReadArray1D(ncid, "DCurr").value(); + std::vector reference_Dgeod = + NetcdfReadArray1D(ncid, "DGeod").value(); for (int jF = 0; jF < fc.ns; ++jF) { EXPECT_TRUE(IsCloseRelAbs(reference_DMerc[jF], wout.DMerc[jF], tolerance)); EXPECT_TRUE( @@ -299,7 +323,8 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { EXPECT_TRUE(IsCloseRelAbs(reference_Dgeod[jF], wout.DGeod[jF], tolerance)); } // jF - std::vector reference_equif = NetcdfReadArray1D(ncid, "equif"); + std::vector reference_equif = + NetcdfReadArray1D(ncid, "equif").value(); for (int jF = 0; jF < fc.ns; ++jF) { EXPECT_TRUE(IsCloseRelAbs(reference_equif[jF], wout.equif[jF], tolerance)); } @@ -309,15 +334,17 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { // ------------------- // mode numbers for Fourier coefficient arrays below - std::vector reference_xm = NetcdfReadArray1D(ncid, "xm"); - std::vector reference_xn = NetcdfReadArray1D(ncid, "xn"); + std::vector reference_xm = NetcdfReadArray1D(ncid, "xm").value(); + std::vector reference_xn = NetcdfReadArray1D(ncid, "xn").value(); for (int mn = 0; mn < wout.mnmax; ++mn) { EXPECT_EQ(wout.xm[mn], reference_xm[mn]); EXPECT_EQ(wout.xn[mn], reference_xn[mn]); } // mn - std::vector reference_xm_nyq = NetcdfReadArray1D(ncid, "xm_nyq"); - std::vector reference_xn_nyq = NetcdfReadArray1D(ncid, "xn_nyq"); + std::vector reference_xm_nyq = + NetcdfReadArray1D(ncid, "xm_nyq").value(); + std::vector reference_xn_nyq = + NetcdfReadArray1D(ncid, "xn_nyq").value(); for (int mn_nyq = 0; mn_nyq < wout.mnmax_nyq; ++mn_nyq) { EXPECT_EQ(wout.xm_nyq[mn_nyq], reference_xm_nyq[mn_nyq]); EXPECT_EQ(wout.xn_nyq[mn_nyq], reference_xn_nyq[mn_nyq]); @@ -326,8 +353,10 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { // ------------------- // stellarator-symmetric Fourier coefficients - std::vector reference_raxis_cc = NetcdfReadArray1D(ncid, "raxis_cc"); - std::vector reference_zaxis_cs = NetcdfReadArray1D(ncid, "zaxis_cs"); + std::vector reference_raxis_cc = + NetcdfReadArray1D(ncid, "raxis_cc").value(); + std::vector reference_zaxis_cs = + NetcdfReadArray1D(ncid, "zaxis_cs").value(); for (int n = 0; n <= wout.ntor; ++n) { EXPECT_TRUE( IsCloseRelAbs(reference_raxis_cc[n], wout.raxis_cc[n], tolerance)); @@ -336,9 +365,9 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { } // n std::vector> reference_rmnc = - NetcdfReadArray2D(ncid, "rmnc"); + NetcdfReadArray2D(ncid, "rmnc").value(); std::vector> reference_zmns = - NetcdfReadArray2D(ncid, "zmns"); + NetcdfReadArray2D(ncid, "zmns").value(); for (int jF = 0; jF < fc.ns; ++jF) { for (int mn = 0; mn < s.mnmax; ++mn) { EXPECT_TRUE( @@ -349,7 +378,7 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { } // jF std::vector> reference_lmns = - NetcdfReadArray2D(ncid, "lmns"); + NetcdfReadArray2D(ncid, "lmns").value(); for (int jF = 0; jF < fc.ns; ++jF) { for (int mn = 0; mn < s.mnmax; ++mn) { EXPECT_TRUE( @@ -358,19 +387,19 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { } // jF std::vector> reference_gmnc = - NetcdfReadArray2D(ncid, "gmnc"); + NetcdfReadArray2D(ncid, "gmnc").value(); std::vector> reference_bmnc = - NetcdfReadArray2D(ncid, "bmnc"); + NetcdfReadArray2D(ncid, "bmnc").value(); std::vector> reference_bsubumnc = - NetcdfReadArray2D(ncid, "bsubumnc"); + NetcdfReadArray2D(ncid, "bsubumnc").value(); std::vector> reference_bsubvmnc = - NetcdfReadArray2D(ncid, "bsubvmnc"); + NetcdfReadArray2D(ncid, "bsubvmnc").value(); std::vector> reference_bsubsmns = - NetcdfReadArray2D(ncid, "bsubsmns"); + NetcdfReadArray2D(ncid, "bsubsmns").value(); std::vector> reference_bsupumnc = - NetcdfReadArray2D(ncid, "bsupumnc"); + NetcdfReadArray2D(ncid, "bsupumnc").value(); std::vector> reference_bsupvmnc = - NetcdfReadArray2D(ncid, "bsupvmnc"); + NetcdfReadArray2D(ncid, "bsupvmnc").value(); for (int jF = 0; jF < fc.ns; ++jF) { for (int mn_nyq = 0; mn_nyq < s.mnmax_nyq; ++mn_nyq) { EXPECT_TRUE(IsCloseRelAbs(reference_gmnc[jF][mn_nyq], @@ -395,9 +424,9 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { if (s.lasym) { std::vector reference_raxis_cs = - NetcdfReadArray1D(ncid, "raxis_cs"); + NetcdfReadArray1D(ncid, "raxis_cs").value(); std::vector reference_zaxis_cc = - NetcdfReadArray1D(ncid, "zaxis_cc"); + NetcdfReadArray1D(ncid, "zaxis_cc").value(); for (int n = 0; n <= wout.ntor; ++n) { EXPECT_TRUE( IsCloseRelAbs(reference_raxis_cs[n], wout.raxis_cs[n], tolerance)); @@ -556,16 +585,17 @@ TEST(SplineProfileEquilibrium, CthLikeCubicSplinePressureMatchesFortranGolden) { }; // scalar physics quantities - scalar("volume_p", NetcdfReadDouble(ncid, "volume_p"), wout.volume); - scalar("betatotal", NetcdfReadDouble(ncid, "betatotal"), wout.betatotal); - scalar("aspect", NetcdfReadDouble(ncid, "aspect"), wout.aspect); - scalar("b0", NetcdfReadDouble(ncid, "b0"), wout.b0); - scalar("wp", NetcdfReadDouble(ncid, "wp"), wout.wp); - scalar("wb", NetcdfReadDouble(ncid, "wb"), wout.wb); - scalar("rbtor", NetcdfReadDouble(ncid, "rbtor"), wout.rbtor); - scalar("ctor", NetcdfReadDouble(ncid, "ctor"), wout.ctor); - scalar("Aminor_p", NetcdfReadDouble(ncid, "Aminor_p"), wout.Aminor_p); - scalar("Rmajor_p", NetcdfReadDouble(ncid, "Rmajor_p"), wout.Rmajor_p); + scalar("volume_p", NetcdfReadDouble(ncid, "volume_p").value(), wout.volume); + scalar("betatotal", NetcdfReadDouble(ncid, "betatotal").value(), + wout.betatotal); + scalar("aspect", NetcdfReadDouble(ncid, "aspect").value(), wout.aspect); + scalar("b0", NetcdfReadDouble(ncid, "b0").value(), wout.b0); + scalar("wp", NetcdfReadDouble(ncid, "wp").value(), wout.wp); + scalar("wb", NetcdfReadDouble(ncid, "wb").value(), wout.wb); + scalar("rbtor", NetcdfReadDouble(ncid, "rbtor").value(), wout.rbtor); + scalar("ctor", NetcdfReadDouble(ncid, "ctor").value(), wout.ctor); + scalar("Aminor_p", NetcdfReadDouble(ncid, "Aminor_p").value(), wout.Aminor_p); + scalar("Rmajor_p", NetcdfReadDouble(ncid, "Rmajor_p").value(), wout.Rmajor_p); // 1D radial profiles, including the spline-driven pressure std::vector wpresf(fc.ns), wpres(fc.ns), wmass(fc.ns), wiotaf(fc.ns), @@ -578,38 +608,38 @@ TEST(SplineProfileEquilibrium, CthLikeCubicSplinePressureMatchesFortranGolden) { wiotas[jF] = wout.iotas[jF]; wjcurv[jF] = wout.jcurv[jF]; } - compare("presf", NetcdfReadArray1D(ncid, "presf"), wpresf); - compare("pres", NetcdfReadArray1D(ncid, "pres"), wpres); - compare("mass", NetcdfReadArray1D(ncid, "mass"), wmass); - compare("iotaf", NetcdfReadArray1D(ncid, "iotaf"), wiotaf); - compare("iotas", NetcdfReadArray1D(ncid, "iotas"), wiotas); - compare("jcurv", NetcdfReadArray1D(ncid, "jcurv"), wjcurv); + compare("presf", NetcdfReadArray1D(ncid, "presf").value(), wpresf); + compare("pres", NetcdfReadArray1D(ncid, "pres").value(), wpres); + compare("mass", NetcdfReadArray1D(ncid, "mass").value(), wmass); + compare("iotaf", NetcdfReadArray1D(ncid, "iotaf").value(), wiotaf); + compare("iotas", NetcdfReadArray1D(ncid, "iotas").value(), wiotas); + compare("jcurv", NetcdfReadArray1D(ncid, "jcurv").value(), wjcurv); // flux-surface geometry auto [r_ref, r_val] = - flatten(NetcdfReadArray2D(ncid, "rmnc"), fc.ns, s.mnmax, + flatten(NetcdfReadArray2D(ncid, "rmnc").value(), fc.ns, s.mnmax, [&](int mn, int jF) { return wout.rmnc(mn, jF); }); compare("rmnc", r_ref, r_val); auto [z_ref, z_val] = - flatten(NetcdfReadArray2D(ncid, "zmns"), fc.ns, s.mnmax, + flatten(NetcdfReadArray2D(ncid, "zmns").value(), fc.ns, s.mnmax, [&](int mn, int jF) { return wout.zmns(mn, jF); }); compare("zmns", z_ref, z_val); auto [l_ref, l_val] = - flatten(NetcdfReadArray2D(ncid, "lmns"), fc.ns, s.mnmax, + flatten(NetcdfReadArray2D(ncid, "lmns").value(), fc.ns, s.mnmax, [&](int mn, int jF) { return wout.lmns(mn, jF); }); compare("lmns", l_ref, l_val); // magnetic field on the Nyquist mode set auto [b_ref, b_val] = - flatten(NetcdfReadArray2D(ncid, "bmnc"), fc.ns, s.mnmax_nyq, + flatten(NetcdfReadArray2D(ncid, "bmnc").value(), fc.ns, s.mnmax_nyq, [&](int mn, int jF) { return wout.bmnc(mn, jF); }); compare("bmnc", b_ref, b_val); auto [bu_ref, bu_val] = - flatten(NetcdfReadArray2D(ncid, "bsubumnc"), fc.ns, s.mnmax_nyq, + flatten(NetcdfReadArray2D(ncid, "bsubumnc").value(), fc.ns, s.mnmax_nyq, [&](int mn, int jF) { return wout.bsubumnc(mn, jF); }); compare("bsubumnc", bu_ref, bu_val); auto [bv_ref, bv_val] = - flatten(NetcdfReadArray2D(ncid, "bsubvmnc"), fc.ns, s.mnmax_nyq, + flatten(NetcdfReadArray2D(ncid, "bsubvmnc").value(), fc.ns, s.mnmax_nyq, [&](int mn, int jF) { return wout.bsubvmnc(mn, jF); }); compare("bsubvmnc", bv_ref, bv_val); @@ -704,17 +734,18 @@ TEST(SolovevFreeBoundary, MatchesEducationalVmecGolden) { }; // scalar physics quantities - scalar("volume_p", NetcdfReadDouble(ncid, "volume_p"), wout.volume); - scalar("betatotal", NetcdfReadDouble(ncid, "betatotal"), wout.betatotal); - scalar("aspect", NetcdfReadDouble(ncid, "aspect"), wout.aspect); - scalar("b0", NetcdfReadDouble(ncid, "b0"), wout.b0); - scalar("wp", NetcdfReadDouble(ncid, "wp"), wout.wp); - scalar("wb", NetcdfReadDouble(ncid, "wb"), wout.wb); - scalar("rbtor", NetcdfReadDouble(ncid, "rbtor"), wout.rbtor); - scalar("ctor", NetcdfReadDouble(ncid, "ctor"), wout.ctor); - scalar("Aminor_p", NetcdfReadDouble(ncid, "Aminor_p"), wout.Aminor_p); - scalar("Rmajor_p", NetcdfReadDouble(ncid, "Rmajor_p"), wout.Rmajor_p); - scalar("volavgB", NetcdfReadDouble(ncid, "volavgB"), wout.volavgB); + scalar("volume_p", NetcdfReadDouble(ncid, "volume_p").value(), wout.volume); + scalar("betatotal", NetcdfReadDouble(ncid, "betatotal").value(), + wout.betatotal); + scalar("aspect", NetcdfReadDouble(ncid, "aspect").value(), wout.aspect); + scalar("b0", NetcdfReadDouble(ncid, "b0").value(), wout.b0); + scalar("wp", NetcdfReadDouble(ncid, "wp").value(), wout.wp); + scalar("wb", NetcdfReadDouble(ncid, "wb").value(), wout.wb); + scalar("rbtor", NetcdfReadDouble(ncid, "rbtor").value(), wout.rbtor); + scalar("ctor", NetcdfReadDouble(ncid, "ctor").value(), wout.ctor); + scalar("Aminor_p", NetcdfReadDouble(ncid, "Aminor_p").value(), wout.Aminor_p); + scalar("Rmajor_p", NetcdfReadDouble(ncid, "Rmajor_p").value(), wout.Rmajor_p); + scalar("volavgB", NetcdfReadDouble(ncid, "volavgB").value(), wout.volavgB); // 1D radial profiles (pressure and rotational transform) std::vector wpresf(fc.ns), wpres(fc.ns), wiotaf(fc.ns), wiotas(fc.ns); @@ -724,36 +755,36 @@ TEST(SolovevFreeBoundary, MatchesEducationalVmecGolden) { wiotaf[jF] = wout.iotaf[jF]; wiotas[jF] = wout.iotas[jF]; } - compare("presf", NetcdfReadArray1D(ncid, "presf"), wpresf); - compare("pres", NetcdfReadArray1D(ncid, "pres"), wpres); - compare("iotaf", NetcdfReadArray1D(ncid, "iotaf"), wiotaf); - compare("iotas", NetcdfReadArray1D(ncid, "iotas"), wiotas); + compare("presf", NetcdfReadArray1D(ncid, "presf").value(), wpresf); + compare("pres", NetcdfReadArray1D(ncid, "pres").value(), wpres); + compare("iotaf", NetcdfReadArray1D(ncid, "iotaf").value(), wiotaf); + compare("iotas", NetcdfReadArray1D(ncid, "iotas").value(), wiotas); // flux-surface geometry auto [r_ref, r_val] = - flatten(NetcdfReadArray2D(ncid, "rmnc"), fc.ns, s.mnmax, + flatten(NetcdfReadArray2D(ncid, "rmnc").value(), fc.ns, s.mnmax, [&](int mn, int jF) { return wout.rmnc(mn, jF); }); compare("rmnc", r_ref, r_val); auto [z_ref, z_val] = - flatten(NetcdfReadArray2D(ncid, "zmns"), fc.ns, s.mnmax, + flatten(NetcdfReadArray2D(ncid, "zmns").value(), fc.ns, s.mnmax, [&](int mn, int jF) { return wout.zmns(mn, jF); }); compare("zmns", z_ref, z_val); auto [l_ref, l_val] = - flatten(NetcdfReadArray2D(ncid, "lmns"), fc.ns, s.mnmax, + flatten(NetcdfReadArray2D(ncid, "lmns").value(), fc.ns, s.mnmax, [&](int mn, int jF) { return wout.lmns(mn, jF); }); compare("lmns", l_ref, l_val); // magnetic field on the Nyquist mode set auto [b_ref, b_val] = - flatten(NetcdfReadArray2D(ncid, "bmnc"), fc.ns, s.mnmax_nyq, + flatten(NetcdfReadArray2D(ncid, "bmnc").value(), fc.ns, s.mnmax_nyq, [&](int mn, int jF) { return wout.bmnc(mn, jF); }); compare("bmnc", b_ref, b_val); auto [bu_ref, bu_val] = - flatten(NetcdfReadArray2D(ncid, "bsubumnc"), fc.ns, s.mnmax_nyq, + flatten(NetcdfReadArray2D(ncid, "bsubumnc").value(), fc.ns, s.mnmax_nyq, [&](int mn, int jF) { return wout.bsubumnc(mn, jF); }); compare("bsubumnc", bu_ref, bu_val); auto [bv_ref, bv_val] = - flatten(NetcdfReadArray2D(ncid, "bsubvmnc"), fc.ns, s.mnmax_nyq, + flatten(NetcdfReadArray2D(ncid, "bsubvmnc").value(), fc.ns, s.mnmax_nyq, [&](int mn, int jF) { return wout.bsubvmnc(mn, jF); }); compare("bsubvmnc", bv_ref, bv_val); @@ -764,8 +795,8 @@ TEST(SolovevFreeBoundary, MatchesEducationalVmecGolden) { wjcuru[jF] = wout.jcuru[jF]; wjcurv[jF] = wout.jcurv[jF]; } - compare("jcuru", NetcdfReadArray1D(ncid, "jcuru"), wjcuru); - compare("jcurv", NetcdfReadArray1D(ncid, "jcurv"), wjcurv); + compare("jcuru", NetcdfReadArray1D(ncid, "jcuru").value(), wjcuru); + compare("jcurv", NetcdfReadArray1D(ncid, "jcurv").value(), wjcurv); ASSERT_EQ(nc_close(ncid), NC_NOERR); 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/cpp/vmecpp/vmec/vmec/vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc index 9f6c1ea96..4915f753f 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc @@ -429,6 +429,7 @@ bool Vmec::InitializeRadial( fc_.ijacob = 0; fc_.restart_reason = RestartReason::NO_RESTART; fc_.res0 = -1; + fc_.res1 = -1; m_delt0 = indata_.delt; // INITIALIZE MESH-DEPENDENT SCALARS @@ -963,9 +964,34 @@ absl::StatusOr Vmec::SolveEquilibriumLoop( // res0 is the best force residual we got so far fc_.res0 = std::min(fc_.res0, fc_.fsq); + + // PARVMEC additionally tracks the invariant residual minimum res1. Keep + // it (and its inputs) off the vmec_8_52 path so the default control stays + // byte-for-byte unchanged. + if (indata_.iteration_style == IterationStyle::PARVMEC) { + const double fsq_invariant = fc_.fsqr + fc_.fsqz + fc_.fsql; + if (iter2 == iter1_ || fc_.res1 == -1) { + fc_.res1 = fsq_invariant; + } + fc_.res1 = std::min(fc_.res1, fsq_invariant); + } } - if (fc_.fsq <= fc_.res0 && (iter2 - iter1_) > 10) { + if (indata_.iteration_style == IterationStyle::PARVMEC) { + // PARVMEC control: store when both residual minima improve; revert via + // BAD_PROGRESS (delt0r /= 1.03, no ijacob) when either exceeds 1e4 * its + // minimum after 10 steps. + const double fsq_invariant = fc_.fsqr + fc_.fsqz + fc_.fsql; + if (fc_.fsq <= fc_.res0 && fsq_invariant <= fc_.res1) { + RestartIteration(fc_.delt0r, thread_id); + } else if ((iter2 - iter1_) > 10 && (fc_.fsq > 1.0e4 * fc_.res0 || + fsq_invariant > 1.0e4 * fc_.res1)) { +#ifdef _OPENMP +#pragma omp single +#endif // _OPENMP + fc_.restart_reason = RestartReason::BAD_PROGRESS; + } + } else if (fc_.fsq <= fc_.res0 && (iter2 - iter1_) > 10) { // Store current state (restart_reason=NO_RESTART) // --> was able to reduce force consistenly over at least 10 iterations RestartIteration(fc_.delt0r, thread_id); diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/free_boundary/mgrid_provider/mgrid_provider_test.cc b/src/vmecpp/cpp/vmecpp_large_cpp_tests/free_boundary/mgrid_provider/mgrid_provider_test.cc index aa52e884b..5cc996a2c 100644 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/free_boundary/mgrid_provider/mgrid_provider_test.cc +++ b/src/vmecpp/cpp/vmecpp_large_cpp_tests/free_boundary/mgrid_provider/mgrid_provider_test.cc @@ -104,21 +104,21 @@ TEST_P(LoadMGridTest, CheckLoadMGrid) { ASSERT_EQ(nc_open(vmec_indata->mgrid_file.c_str(), NC_NOWRITE, &ncid), NC_NOERR); - const int number_of_field_periods = NetcdfReadInt(ncid, "nfp"); + const int number_of_field_periods = NetcdfReadInt(ncid, "nfp").value(); - const int number_of_r_grid_points = NetcdfReadInt(ncid, "ir"); - const double r_grid_minimum = NetcdfReadDouble(ncid, "rmin"); - const double r_grid_maximum = NetcdfReadDouble(ncid, "rmax"); + const int number_of_r_grid_points = NetcdfReadInt(ncid, "ir").value(); + const double r_grid_minimum = NetcdfReadDouble(ncid, "rmin").value(); + const double r_grid_maximum = NetcdfReadDouble(ncid, "rmax").value(); const double r_grid_increment = (r_grid_maximum - r_grid_minimum) / (number_of_r_grid_points - 1.0); - const int number_of_z_grid_points = NetcdfReadInt(ncid, "jz"); - const double z_grid_minimum = NetcdfReadDouble(ncid, "zmin"); - const double z_grid_maximum = NetcdfReadDouble(ncid, "zmax"); + const int number_of_z_grid_points = NetcdfReadInt(ncid, "jz").value(); + const double z_grid_minimum = NetcdfReadDouble(ncid, "zmin").value(); + const double z_grid_maximum = NetcdfReadDouble(ncid, "zmax").value(); const double z_grid_increment = (z_grid_maximum - z_grid_minimum) / (number_of_z_grid_points - 1.0); - const int number_of_phi_grid_points = NetcdfReadInt(ncid, "kp"); + const int number_of_phi_grid_points = NetcdfReadInt(ncid, "kp").value(); const double phi_grid_increment = 2.0 * M_PI / (number_of_phi_grid_points * number_of_field_periods); diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/output_quantities/output_quantities_test.cc b/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/output_quantities/output_quantities_test.cc index 68b062c89..87cd81582 100644 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/output_quantities/output_quantities_test.cc +++ b/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/output_quantities/output_quantities_test.cc @@ -1358,9 +1358,9 @@ TEST_P(CurrentDensityTest, CheckCurrentDensityFourierCoefficients) { // Read 2D arrays: Fortran stores as (ns, mnmax_nyq), C++ as (mnmax_nyq, ns) const std::vector> ref_currumnc = - NetcdfReadArray2D(ncid, "currumnc"); + NetcdfReadArray2D(ncid, "currumnc").value(); const std::vector> ref_currvmnc = - NetcdfReadArray2D(ncid, "currvmnc"); + NetcdfReadArray2D(ncid, "currvmnc").value(); nc_close(ncid); const int ns = static_cast(ref_currumnc.size()); 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_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 index a682721e3..5e20b7fa0 100644 --- a/tests/test_external_optimizers.py +++ b/tests/test_external_optimizers.py @@ -2,12 +2,13 @@ # # # SPDX-License-Identifier: MIT -"""External optimizers reach the same equilibrium as the native solver. +"""External optimizers reach force balance on axisymmetric and 3D cases. -The raw internal-basis force (gradient of VMEC's augmented functional) is the residual -F(x); F(x) = 0 at equilibrium. Both a native-style preconditioned descent and a -Jacobian-free Newton-Krylov solver drive it to zero and recover the native solver's -converged state and energy. +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 @@ -21,44 +22,86 @@ make_model, reference_equilibrium, residual, + solve_newton_hvp, solve_newton_krylov, + solve_newton_krylov_preconditioned, solve_preconditioned_descent, ) -CMA = ( - Path(__file__).resolve().parents[1] - / "src" - / "vmecpp" - / "cpp" - / "vmecpp" - / "test_data" - / "cma.json" +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(): - return reference_equilibrium() +def reference_2d(): + return reference_equilibrium(SOLOVEV, ns=11) -@pytest.mark.parametrize("solver", [solve_preconditioned_descent, solve_newton_krylov]) -def test_optimizer_reaches_equilibrium(solver, reference): - x_star, w_star = reference - x, result = solver() - # Force balance achieved. - assert result.residual_norm < 1e-7 - # Same equilibrium as the native solver. - assert abs(result.energy - w_star) < 1e-8 - assert np.linalg.norm(x - x_star) < 1e-5 +@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(): - # cma.json is a 3D stellarator (nfp=2, ntor=6) that ships no magnetic axis - # (raxis/zaxis all zero), so the initial geometry has a singular Jacobian. - # make_model reguesses the axis like the native solver, after which the raw - # internal-basis force is well defined on the non-axisymmetric force chain. 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_init.py b/tests/test_init.py index 0d223c8ec..10604d6a5 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -51,8 +51,7 @@ def test_run(max_threads, input_file, verbose): [ ("cma.json", RuntimeError), # Invalid netcdf ("does_not_exist", RuntimeError), - # TODO(jurasic) Enable test after switching netcdf_io to absl::Status - # ("wout_cma.nc", RuntimeError), # Valid netcdf, but invalid mgrid + ("wout_cma.nc", RuntimeError), # Valid netcdf, but invalid mgrid ], ) def test_raise_invalid_mgrid(mgrid_path: str, expected_exception): @@ -615,8 +614,6 @@ def test_vmec_input_validation(): # The test_file json may exclude fields that have default values, # while the parsed versions should have all fields populated. indata_dict_from_json = json.loads(vmec_input._to_cpp_vmecindata().to_json()) - # TODO(jurasic): iteration_style is not yet present in VmecInput, since there's only one option atm. - del indata_dict_from_json["iteration_style"] vmec_input_dict_from_json = json.loads(vmec_input.model_dump_json()) if not vmec_input.lasym: 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_iteration.py b/tests/test_iteration.py index 8932ab092..f8699daca 100644 --- a/tests/test_iteration.py +++ b/tests/test_iteration.py @@ -208,6 +208,172 @@ def test_alternative_styles_converge_cma(): assert results["parvmec"].restarts == 0 +def test_native_parvmec_matches_python_parvmec(): + """The native C++ PARVMEC control reproduces the ported Python PARVMEC control. + + Vmec::SolveEquilibriumLoop runs the PARVMEC time-step control when + indata.iteration_style == PARVMEC; it must match the Python parvmec loop on the + same forward model step for step, the analog of + test_python_iteration_matches_cpp_restart_path for the default style. + """ + cpp_indata = _single_resolution_indata("cma", 72, ftol=1.0e-16, niter=200) + cpp_indata.iteration_style = _vmecpp.IterationStyle.PARVMEC + + reference = _vmecpp.VmecModel.create(cpp_indata, 72) + reference.solve() + + model = _vmecpp.VmecModel.create(cpp_indata, 72) + result = vmecpp.solve_equilibrium(model, style="parvmec") + + assert not result.failed + np.testing.assert_array_equal( + np.asarray(result.restart_reasons), np.asarray(reference.restart_reasons) + ) + cpp_r = np.asarray(reference.force_residual_r) + py_r = np.asarray(result.force_residual_r) + # The two loops make identical control decisions (restart_reasons match + # exactly above), so the force-residual traces agree up to floating-point + # accumulation of the control arithmetic: ~1e-9 relative early, growing to + # only a few 1e-9 by the deep-convergence tail (max abs ~6e-13). + np.testing.assert_allclose(py_r[:50], cpp_r[:50], rtol=1.0e-9, atol=1e-15) + np.testing.assert_allclose(py_r, cpp_r, rtol=1.0e-8, atol=1e-12) + + +def test_run_honors_iteration_style_flag(): + """vmecpp.run() honors the iteration_style input flag through the native solver. + + The two styles take different iteration paths to the same equilibrium, so the flag + must survive the VmecInput -> C++ round-trip and reach Vmec::SolveEquilibriumLoop. + """ + base = vmecpp.VmecInput.from_file(TEST_DATA / "cma.json").model_copy( + update={ + "ns_array": np.array([51], dtype=np.int64), + "ftol_array": np.array([1.0e-12]), + "niter_array": np.array([3000], dtype=np.int64), + } + ) + outputs = {} + for style in ("vmec_8_52", "parvmec"): + inp = base.model_copy(update={"iteration_style": style}) + assert inp.iteration_style == style + assert inp._to_cpp_vmecindata().iteration_style == getattr( + _vmecpp.IterationStyle, style.upper() + ) + outputs[style] = vmecpp.run(inp, max_threads=1, verbose=False) + # Same physics regardless of the iteration scheme: the two styles take + # different paths to the same equilibrium, so global geometry, pressure, and + # magnetic-field quantities must agree, not just the volume. (Local profiles + # such as the iota profile are path-sensitive at finite ftol, so they are not + # asserted here; the exact-reference check against PARVMEC covers those.) + ref = outputs["vmec_8_52"].wout + par = outputs["parvmec"].wout + assert par.volume_p == pytest.approx(ref.volume_p, rel=1.0e-9) # geometry + assert par.aspect == pytest.approx(ref.aspect, rel=1.0e-9) # geometry + assert par.betatotal == pytest.approx(ref.betatotal, rel=1.0e-9) # beta + assert par.wp == pytest.approx(ref.wp, rel=1.0e-9) # pressure energy + assert par.wb == pytest.approx(ref.wb, rel=1.0e-5) # magnetic energy + + +@pytest.mark.parametrize("case", ["cth_like_fixed_bdy", "solovev"]) +def test_parvmec_matches_parvmec_reference(case): + """The PARVMEC iteration style reproduces the ORNL-Fusion/PARVMEC wout. + + The committed reference wouts match fresh output from the Fortran ORNL- + Fusion/PARVMEC to machine precision (volume/aspect ~1e-15, geometry and iota ~1e-7 + for cth_like, ~0 for solovev), so this pins the new iteration style to the + independent parallel implementation, not only to the vmec_8_52 control. + """ + reference = vmecpp.VmecWOut.from_wout_file(TEST_DATA / f"wout_{case}.nc") + base = vmecpp.VmecInput.from_file(TEST_DATA / f"{case}.json") + result = vmecpp.run( + base.model_copy(update={"iteration_style": "parvmec"}), + max_threads=1, + verbose=False, + ) + w = result.wout + + assert w.volume_p == pytest.approx(reference.volume_p, rel=1.0e-9) + assert w.aspect == pytest.approx(reference.aspect, rel=1.0e-9) + np.testing.assert_allclose(w.iotaf, reference.iotaf, rtol=1.0e-5, atol=1.0e-6) + np.testing.assert_allclose(w.rmnc, reference.rmnc, rtol=1.0e-5, atol=1.0e-6) + np.testing.assert_allclose(w.zmns, reference.zmns, rtol=1.0e-5, atol=1.0e-6) + np.testing.assert_allclose(w.bmnc, reference.bmnc, rtol=1.0e-5, atol=1.0e-6) + + +def test_parvmec_follows_ornl_parvmec_trace(): + """The native parvmec control follows ORNL PARVMEC's per-iteration residual trace. + + cth_like_fixed_bdy has a well-posed linear guess (no cold-start axis reguess) and + converges without a restart, so the iteration is numerically stable: the native + parvmec time-step control reproduces the committed ORNL-Fusion/PARVMEC force + residuals step-for-step. The first several steps agree to machine precision; the + difference then settles into a bounded ~1e-4 relative drift (floating-point + accumulation between two independent implementations) that does not grow, and both + reach force balance in the same number of steps. This pins the control to the + Fortran PARVMEC itself, not only to the vmec_8_52 baseline or the Python port. + + A step-for-step match is only meaningful on a non-chaotic case. Inputs that trip the + parvmec-specific restart / reguess logic have violent transients (fsqr ~ 1e4) on + which the ~1e-13 floating-point difference between any two implementations amplifies + to order unity within a couple of steps; those converge to the same equilibrium but + cannot be compared trace-for-trace. The divergence between the two styles on such a + case is covered by test_iteration_styles_diverge_on_stiff_case. + """ + ref = np.loadtxt( + TEST_DATA / "parvmec_cth_like_fixed_bdy_force_trace.csv", + delimiter=",", + skiprows=5, + ) + ref_fsqr, ref_fsqz, ref_fsql = ref[:, 1], ref[:, 2], ref[:, 3] + + cpp_indata = _single_resolution_indata("cth_like_fixed_bdy", 25, 1.0e-6, 25000) + cpp_indata.iteration_style = _vmecpp.IterationStyle.PARVMEC + model = _vmecpp.VmecModel.create(cpp_indata, 25) + model.solve() + fsqr = np.asarray(model.force_residual_r) + fsqz = np.asarray(model.force_residual_z) + fsql = np.asarray(model.force_residual_lambda) + + # Same number of steps to force balance (up to the last step at finite ftol). + assert abs(len(fsqr) - len(ref_fsqr)) <= 2 + n = min(len(fsqr), len(ref_fsqr)) + + # Bit-faithful for the first several steps. + np.testing.assert_allclose(fsqr[:6], ref_fsqr[:6], rtol=1.0e-9, atol=1.0e-14) + np.testing.assert_allclose(fsqz[:6], ref_fsqz[:6], rtol=1.0e-9, atol=1.0e-14) + np.testing.assert_allclose(fsql[:6], ref_fsql[:6], rtol=1.0e-9, atol=1.0e-14) + + # Tracks ORNL PARVMEC over the whole solve; the bounded drift never turns into a + # flow-control divergence (which would show up as a discrete jump, not a drift). + np.testing.assert_allclose(fsqr[:n], ref_fsqr[:n], rtol=3.0e-3, atol=1.0e-9) + np.testing.assert_allclose(fsqz[:n], ref_fsqz[:n], rtol=3.0e-3, atol=1.0e-9) + np.testing.assert_allclose(fsql[:n], ref_fsql[:n], rtol=3.0e-3, atol=1.0e-9) + + +def test_iteration_styles_diverge_on_stiff_case(): + """vmec_8_52 and parvmec take measurably different paths on a restart-triggering + case. + + cma at ns=72 has a cold-start axis reguess and a violent initial transient that + trips the time-step-control restart logic. The two styles' restart bookkeeping and + force-residual progressions differ -- the reason the parvmec control exists -- even + though both converge to the same equilibrium. This is the flow-control difference + that the trace test cannot pin to ORNL PARVMEC directly (the case is chaotic). + """ + traces = {} + for style in ("vmec_8_52", "parvmec"): + cpp_indata = _single_resolution_indata("cma", 72, 1.0e-11, 200) + cpp_indata.iteration_style = getattr(_vmecpp.IterationStyle, style.upper()) + model = _vmecpp.VmecModel.create(cpp_indata, 72) + model.solve() + traces[style] = np.asarray(model.force_residual_r) + + r852, rpar = traces["vmec_8_52"], traces["parvmec"] + n = min(len(r852), len(rpar)) + # The two controls share only a prefix, then take genuinely different paths. + assert not np.allclose(r852[:n], rpar[:n], rtol=1.0e-6, atol=1.0e-12) + + def test_callback_records_iteration_state(): """The per-iteration callback fires once per recorded iteration with a consistent IterationState snapshot of the convergence / flow-control state. 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)