diff --git a/.github/actions/run-benchmarks/action.yml b/.github/actions/run-benchmarks/action.yml index 981b9d354..2d24521f3 100644 --- a/.github/actions/run-benchmarks/action.yml +++ b/.github/actions/run-benchmarks/action.yml @@ -1,11 +1,12 @@ name: Run VMEC++ benchmarks description: > - Install VMEC++ and run the pytest benchmark suite in a fixed environment, - then store or compare the results. This is the single source of truth for the - benchmark environment so the pull-request and main/docs runs cannot drift - apart; the only intended difference between callers is whether results are - written back to the benchmark-runs branch (auto-push) and whether a regression - is commented on the PR (comment-on-alert). + Install VMEC++ and run the benchmark suites in a fixed environment, then store + or compare the results. Two suites run here: the Python pytest end-to-end + benchmarks and the C++ Google Benchmark function-level microbenchmarks. This is + the single source of truth for the benchmark environment so the pull-request + and main/docs runs cannot drift apart; the only intended difference between + callers is whether results are written back to the benchmark-runs branch + (auto-push) and whether a regression is commented on the PR (comment-on-alert). inputs: auto-push: @@ -32,19 +33,35 @@ runs: sudo apt-get update sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev libhdf5-dev + - name: Mount bazel cache + uses: actions/cache@v4 + with: + path: "~/.cache/bazel" + key: bazel-bench + - name: Install vmecpp with benchmark dependencies shell: bash run: | pip install uv uv pip install --system -v .[benchmark] - - name: Run benchmarks + # --- Python end-to-end benchmarks (pytest-benchmark) --- + + - name: Run Python benchmarks shell: bash run: | export OMP_WAIT_POLICY=passive pytest benchmarks/test_benchmarks.py --benchmark-json=benchmark_results.json - - name: Store / compare benchmark result + # 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 # by integration'. The benchmarks still run above; only the write-back is @@ -52,10 +69,118 @@ 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%" + fail-on-alert: false + auto-push: ${{ inputs.auto-push }} + comment-on-alert: ${{ inputs.comment-on-alert }} + github-token: ${{ inputs.github-token }} + + # --- C++ function-level microbenchmarks (Google Benchmark) --- + # + # 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 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. + + - name: Run C++ microbenchmarks + shell: bash + run: | + set -euo pipefail + + TARGETS=( + "//vmecpp/vmec/ideal_mhd_model:fft_toroidal_bench" + "//vmecpp/vmec/ideal_mhd_model:dealias_constraint_force_bench" + "//vmecpp/free_boundary/laplace_solver:laplace_solver_bench" + "//vmecpp/vmec/output_quantities:output_quantities_bench" + ) + + OUT_FILE="${GITHUB_WORKSPACE}/cpp_benchmark_results.json" + TMP_DIR="$(mktemp -d)" + trap 'rm -rf "${TMP_DIR}"' EXIT + + cd src/vmecpp/cpp + + # --config=perf: optimized build with frame pointers and symbols + # retained, matching the existing .bazelrc convention for + # performance work. + bazel build --config=perf -- "${TARGETS[@]}" + + # Run each target. `bazel run` so runfiles (e.g. test_data for the + # output-quantities benchmark) are available in the working + # directory. --benchmark_min_time keeps CI runtime bounded while + # still averaging enough iterations for a stable measurement. + for target in "${TARGETS[@]}"; do + name="${target##*:}" + bazel run --config=perf -- "${target}" \ + --benchmark_min_time=0.2s \ + --benchmark_format=json \ + --benchmark_out="${TMP_DIR}/${name}.json" + done + + python3 - "${TMP_DIR}" "${OUT_FILE}" <<'PY' + import glob + import json + import os + import sys + + tmp_dir, out_file = sys.argv[1], sys.argv[2] + merged = {"context": None, "benchmarks": []} + skipped = 0 + for path in sorted(glob.glob(os.path.join(tmp_dir, "*.json"))): + with open(path) as f: + data = json.load(f) + if merged["context"] is None: + merged["context"] = data.get("context") + for bench in data.get("benchmarks", []): + if bench.get("error_occurred"): + skipped += 1 + continue + merged["benchmarks"].append(bench) + + with open(out_file, "w") as f: + json.dump(merged, f, indent=2) + + print(f"Merged {len(merged['benchmarks'])} benchmark entries " + 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. + # + # Runs sequentially after the Python store step (composite-action steps + # are serial), so the two auto-push writes to benchmark-runs never race. + # A distinct suite name keeps the C++ results under their own key in + # data.js (window.BENCHMARK_DATA.entries) so they never overwrite the + # Python "Benchmark" suite, even though both share benchmark-data-dir-path. + # + # skip-fetch-gh-pages: the Python store step above already fetched the + # benchmark-runs branch into a local ref and (in compare mode) advanced + # it, so a second `git fetch benchmark-runs:benchmark-runs` here is + # rejected as non-fast-forward. Reuse the already-fetched local ref + # instead of re-fetching. + 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: + name: "C++ Microbenchmarks" + 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 alert-threshold: "150%" fail-on-alert: false auto-push: ${{ inputs.auto-push }} diff --git a/.github/workflows/clang_tidy.yaml b/.github/workflows/clang_tidy.yaml index 4297a6164..72d944cd7 100644 --- a/.github/workflows/clang_tidy.yaml +++ b/.github/workflows/clang_tidy.yaml @@ -31,9 +31,11 @@ jobs: apt_packages: libhdf5-dev, liblapack-dev, libnetcdf-dev, gfortran, python3-dev, libomp-dev build_dir: build cmake_command: cmake -B build && cmake -DCMAKE_EXPORT_COMPILE_COMMANDS:STRING=ON build - # bazel-built test files have no cmake compile_commands entry, so - # clang-tidy cannot resolve ; exclude them from review - exclude: "./src/vmecpp/cpp/bazel-*,./src/vmecpp/cpp/external,./src/vmecpp/cpp/third_party,*_test.cc" + # bazel-built test and benchmark files have no cmake compile_commands + # entry, so clang-tidy cannot resolve / + # (both are Bazel-only deps); exclude them + # from review + exclude: "./src/vmecpp/cpp/bazel-*,./src/vmecpp/cpp/external,./src/vmecpp/cpp/third_party,*_test.cc,*_bench.cc" # turn off "LGTM" comments: we only want pings about warnings/errors lgtm_comment_body: "" # If there are any comments, fail the check 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 1b8e46003..19754aa72 100644 --- a/docs/_static/benchmarks.js +++ b/docs/_static/benchmarks.js @@ -37,36 +37,27 @@ "#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; container.innerHTML = ""; - // Collect benchmarks per test case + // Collect benchmarks per test case, grouped by suite (the `name` given to + // each github-action-benchmark step, e.g. "Benchmark" vs + // "C++ Microbenchmarks") so suites render as separate sections instead of + // interleaving in one flat list. var entries = data.entries; - var allBenches = new Map(); + var colorIdx = 0; + var theme = getThemeColors(); - Object.keys(entries).forEach(function (name) { - entries[name].forEach(function (entry) { + Object.keys(entries).forEach(function (suiteName) { + var suiteBenches = new Map(); + entries[suiteName].forEach(function (entry) { entry.benches.forEach(function (bench) { - var arr = allBenches.get(bench.name); + var arr = suiteBenches.get(bench.name); if (!arr) { arr = []; - allBenches.set(bench.name, arr); + suiteBenches.set(bench.name, arr); } arr.push({ commit: entry.commit, @@ -75,10 +66,29 @@ }); }); }); + if (suiteBenches.size === 0) return; + + var heading = document.createElement("h3"); + heading.textContent = suiteName; + container.appendChild(heading); + + suiteBenches.forEach(function (dataset, benchName) { + renderChart(container, benchName, dataset, theme, colorIdx); + colorIdx++; + }); }); - var colorIdx = 0; - allBenches.forEach(function (dataset, benchName) { + // Last update info + var info = document.createElement("p"); + info.className = "small"; + info.style.textAlign = "center"; + info.style.opacity = "0.7"; + info.textContent = "Last updated: " + new Date(data.lastUpdate).toLocaleString(); + container.appendChild(info); + } + + function renderChart(container, benchName, dataset, theme, colorIdx) { + { var wrapper = document.createElement("div"); wrapper.style.marginBottom = "1rem"; container.appendChild(wrapper); @@ -95,9 +105,6 @@ wrapper.appendChild(canvas); var color = COLORS[colorIdx % COLORS.length]; - colorIdx++; - - var theme = getThemeColors(); new Chart(canvas, { type: "line", @@ -109,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", @@ -182,15 +188,7 @@ }, }, }); - }); - - // Last update info - var info = document.createElement("p"); - info.className = "small"; - info.style.textAlign = "center"; - info.style.opacity = "0.7"; - info.textContent = "Last updated: " + new Date(data.lastUpdate).toLocaleString(); - container.appendChild(info); + } } // Re-render on Furo theme toggle diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 2d21f5997..5b2f8988b 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -1,9 +1,9 @@ # Benchmarks Performance improvements and regressions of `vmecpp` can be tracked below on a per-commit basis. A larger set of benchmarks against VMEC2000 can be found at [`proximafusion/vmecpp-benchmarks`](https://github.com/proximafusion/vmecpp-benchmarks). -A small but representative set of benchmarks runs automatically on every push to `main` and on pull requests, using [pytest-benchmark](https://pytest-benchmark.readthedocs.io/) and [github-action-benchmark](https://github.com/benchmark-action/github-action-benchmark/). +A small but representative set of benchmarks runs automatically on every push to `main` and on pull requests, using [pytest-benchmark](https://pytest-benchmark.readthedocs.io/) (Python, end-to-end) and [Google Benchmark](https://github.com/google/benchmark) (C++, function-level), both tracked via [github-action-benchmark](https://github.com/benchmark-action/github-action-benchmark/). -## Benchmark suite +## End-to-end benchmark suite (Python) | Benchmark | Description | |-----------|-------------| @@ -14,6 +14,22 @@ A small but representative set of benchmarks runs automatically on every push to | `test_bench_response_table_from_coils` | Magnetic field response table creation from coils file | | `test_bench_free_boundary` | Free-boundary solve with pre-computed response table | +## Microbenchmark suite (C++) + +These target individual hot functions so a regression can be attributed to a specific kernel rather than only showing up in the end-to-end timings above. Each is a Google Benchmark `cc_binary` under `src/vmecpp/cpp/`, swept over several `(mpol, ntor)` resolutions. + +| Benchmark | Description | +|-----------|-------------| +| `DftFourierToReal` / `FftFourierToReal` | Inverse transform (Fourier -> real-space geometry), DFT fallback vs FFTX path | +| `Dft` / `Fft` (real-space sweep) | Inverse transform cost as the real-space grid (`ntheta`, `nzeta`) is scaled | +| `BM_FourierToReal_Parallel_W7x_4t` | Inverse transform under VMEC's persistent multi-threaded OpenMP call pattern | +| `DeAliasConstraintForce` | Spectral-condensation constraint-force de-aliasing (forward+inverse transform pair) | +| `LaplaceSolve` / `LaplaceDecompose` | Free-boundary (NESTOR) dense Laplace solve: assemble + LU factorize + back-substitute | +| `TransformGreensFunctionDerivative` | Free-boundary Green's-function-derivative Fourier transform | +| `ComputeOutputQuantities` | Post-solve output computation (wout, jxbout, Mercier, ...) -- time spent after the solve has converged | + +FFT-path benchmarks are skipped at resolutions for which no vendored FFTX codelet exists (the solver falls back to the DFT path there), so only resolutions with an actual FFT kernel appear in the FFT charts. + ```{raw} html @@ -27,6 +43,8 @@ Click on any data point to open the corresponding commit. ## Running locally +Python end-to-end benchmarks: + ```bash pip install -e .[benchmark] pytest benchmarks/test_benchmarks.py -v @@ -37,3 +55,16 @@ To produce a JSON report: ```bash pytest benchmarks/test_benchmarks.py --benchmark-json=benchmark_results.json ``` + +C++ microbenchmarks are built and run via Bazel, e.g.: + +```bash +cd src/vmecpp/cpp +bazel run --config=perf -- //vmecpp/vmec/ideal_mhd_model:fft_toroidal_bench +``` + +The other three targets are `//vmecpp/vmec/ideal_mhd_model:dealias_constraint_force_bench`, +`//vmecpp/free_boundary/laplace_solver:laplace_solver_bench`, and +`//vmecpp/vmec/output_quantities:output_quantities_bench`. CI builds, runs, and merges +all four into a single JSON report for the regression tracker (see +`.github/actions/run-benchmarks/action.yml`). diff --git a/examples/data/README.md b/examples/data/README.md index 00efb23fd..543b37f78 100644 --- a/examples/data/README.md +++ b/examples/data/README.md @@ -1,13 +1,13 @@ # Example input and output files for VMEC++ -A few cases are available for testing VMCE++ and experimenting with it: +A few cases are available for testing VMEC++ and experimenting with it: 1. `cth_like_fixed_bdy.json` - Stellarator case, similar to the Compact Toroidal Hybrid ([CTH](https://www.auburn.edu/cosam/departments/physics/research/plasma_physics/compact_toroidal_hybrid/index.htm)) device 1. `input.cth_like_fixed_bdy` - Fortran namelist input file to be used with Fortran VMEC 1. `cth_like_fixed_bdy.json` - JSON input file for VMEC++, derived from `input.cth_like_fixed_bdy` using [`indata2json`](https://github.com/jonathanschilling/indata2json) 1. `wout_cth_like_fixed_bdy.nc` - NetCDF output file, produced using [`PARVMEC`](https://github.com/ORNL-Fusion/PARVMEC) from `input.cth_like_fixed_bdy`, for testing the loading of a `wout` file using VMEC++'s tooling -1. `input.nfp4_QH_warm_start` - quasi-helically example for use with SIMSOPT +1. `input.nfp4_QH_warm_start` - quasi-helical example for use with SIMSOPT 1. `solovev` - axisymmetric Tokamak case, similar to the Solov'ev equilibrium used in the [1983 Hirshman & Whitson article](https://doi.org/10.1063/1.864116) 1. `input.solovev` - Fortran namelist input file for use with Fortran VMEC diff --git a/examples/external_optimizers.py b/examples/external_optimizers.py new file mode 100644 index 000000000..b4da85f26 --- /dev/null +++ b/examples/external_optimizers.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Drive a VMEC++ equilibrium from outside with general-purpose optimizers. + +VMEC's equilibrium is the stationary point of its augmented functional (MHD +energy plus the spectral-condensation and lambda constraints). The gradient of +that functional in the decomposed internal basis is the raw, unpreconditioned +force exposed by ``VmecModel.evaluate(precondition=False)``. Finding the +equilibrium is therefore the root problem F(x) = 0, which gradient- and +Hessian-based solvers can attack while reusing VMEC++'s forward model, its +preconditioner (``apply_preconditioner``, VMEC's approximate inverse Hessian), +and its Hessian-vector product (``hessian_vector_product``, a directional +derivative of the analytic force computed inside VMEC++). + +This module wires that residual to several solvers and is shared by the +benchmark ``main`` below and by the tests: + +* preconditioned descent (VMEC's own update direction), +* Jacobian-free Newton-Krylov, plain and preconditioned, and +* a true Newton-Krylov driven by VMEC++'s own Hessian-vector product. + +All converge to the same equilibrium as the native solver. Force evaluations are +counted inside VMEC++ (``force_eval_count``) so the comparison is fair across +methods, including the evaluations hidden in Hessian-vector products. + +Interpretation: what separates the methods is how they couple DOFs of +neighbouring flux surfaces. Plain JFNK does not -- it treats every degree of +freedom as independent, so it cannot damp the radial stiffness that lets a +surface cross its neighbour between iterations (BAD_JACOBIAN restarts in the +native solver). It still converges, but slowly. VMEC's couples each surface to +its jF +/- 1 neighbours; applying it -- as the inner Krylov preconditioner, or +refreshed every step in the HVP Newton -- is what rescues conditioning and cuts +the cost by an order of magnitude. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from scipy.optimize import newton_krylov +from scipy.sparse.linalg import LinearOperator, gmres + +import vmecpp +from vmecpp.cpp import _vmecpp # type: ignore[import] + +DEFAULT_INPUT = ( + Path(__file__).resolve().parents[1] + / "examples" + / "data" + / "cth_like_fixed_bdy.json" +) + + +_BAD_JACOBIAN = 2 # RestartReason::BAD_JACOBIAN (flow_control.h) + + +def _ensure_valid_initial_jacobian(model, max_reguess: int = 2) -> None: + """Reguess the magnetic axis until the initial Jacobian is non-singular. + + Inputs that ship no axis (e.g. cma.json, with raxis/zaxis all zero) start from a + degenerate geometry, so the bare forward model returns at the BAD_JACOBIAN + checkpoint with zero force. The native solver reguesses the axis on the first + iterate (vmec.cc SolveEquilibriumLoop); mirror that here via reinitialize() so a + cold start from any valid INDATA has a defined gradient. + """ + for _ in range(max_reguess): + model.evaluate(2, 2, False) + if model.restart_reason != _BAD_JACOBIAN: + return + model.reinitialize() + model.evaluate(2, 2, False) + + +def make_model(input_path: Path = DEFAULT_INPUT, ns: int = 11): + indata = _vmecpp.VmecINDATA.from_file(str(input_path)) + model = _vmecpp.VmecModel.create(indata, ns) + _ensure_valid_initial_jacobian(model) + return model + + +def residual(model): + """Return F(x) = raw internal-basis force; F(x) = 0 at equilibrium.""" + + def F(x): + model.set_state(np.ascontiguousarray(x, dtype=float)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), dtype=float) + + return F + + +@dataclass +class Result: + name: str + force_evals: int + outer_iters: int + seconds: float + residual_norm: float + energy: float + + +def reference_equilibrium(input_path: Path = DEFAULT_INPUT, ns: int = 11): + model = make_model(input_path, ns) + model.solve() + model.evaluate(2, 2, False) + return np.asarray(model.get_state(), float), model.mhd_energy + + +def _finish(model, name, x, outer_iters, t0): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return x, Result( + name, + model.force_eval_count, + outer_iters, + time.perf_counter() - t0, + float(np.linalg.norm(np.asarray(model.get_forces(), float))), + model.mhd_energy, + ) + + +def solve_vmecpp(input_path=DEFAULT_INPUT, ns=11): + """Native VMEC++ solve through the public API; reports the iteration count. + + Unlike the other variants this drives the ordinary ``vmecpp.run`` path rather + than the low-level ``VmecModel`` primitives, so the reported residual/energy + are the public wout diagnostics (invariant force residuals and total MHD + energy) rather than the raw internal-basis force this module counts elsewhere. + """ + vmec_input = vmecpp.VmecInput.from_file(input_path) + vmec_input.ns_array = np.asarray([ns], dtype=vmec_input.ns_array.dtype) + t0 = time.perf_counter() + output = vmecpp.run(vmec_input, verbose=False) + wout = output.wout + return output, Result( + name="VMEC++ (native, vmecpp.run)", + force_evals=wout.niter, + outer_iters=wout.niter, + seconds=time.perf_counter() - t0, + residual_norm=float(np.sqrt(wout.fsqt[-1])), + energy=wout.wb + wout.wp, + ) + + +def solve_preconditioned_descent( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, delt=0.9, momentum=0.5, max_iter=20000 +): + model = make_model(input_path, ns) + F = residual(model) + x = np.asarray(model.get_state(), float).copy() + v = np.zeros_like(x) + model.reset_force_eval_count() + it = 0 + t0 = time.perf_counter() + for _ in range(max_iter): + if np.linalg.norm(F(x)) < tol: + break + it += 1 + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, True) # preconditioned search direction + fprec = np.asarray(model.get_forces(), float) + v = momentum * v + delt * fprec + x = x + delt * v + return _finish(model, "preconditioned descent", x, it, t0) + + +def solve_newton_krylov( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=2500, preconditioned=False +): + model = make_model(input_path, ns) + F = residual(model) + x0 = np.asarray(model.get_state(), float) + inner_m = None + model.reset_force_eval_count() + if preconditioned: + # Assemble VMEC's preconditioner at x0 and use it, frozen, as the inner + # Krylov preconditioner. M^-1 approximates the inverse Hessian; it is a + # radial tridiagonal (Thomas) solve per Fourier mode, so it couples each + # flux surface to its jF +/- 1 neighbours -- the coupling the plain + # branch lacks. + model.evaluate(2, 2, True) + n_dof = x0.size + inner_m = LinearOperator( # type: ignore[call-overload] + (n_dof, n_dof), + matvec=lambda b: np.asarray( # type: ignore[call-overload] + model.apply_preconditioner(np.ascontiguousarray(b)), float + ), + ) + t0 = time.perf_counter() + x = newton_krylov( + F, x0, f_tol=tol, maxiter=max_iter, method="lgmres", inner_M=inner_m + ) + name = ( + "Newton-Krylov (preconditioned)" if preconditioned else "Newton-Krylov (JFNK)" + ) + return _finish(model, name, x, 0, t0) + + +def solve_newton_krylov_preconditioned(input_path=DEFAULT_INPUT, ns=11, tol=1e-9): + return solve_newton_krylov(input_path, ns, tol, preconditioned=True) + + +def solve_newton_hvp( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_newton=80, inner_tol=1e-3 +): + """Globalized Newton-Krylov using VMEC++'s own Hessian-vector product. + + Each Newton step solves H dx = -F with GMRES, where H v is hessian_vector_product + (the analytic force's directional derivative computed inside VMEC++) and the inner + solve is preconditioned by M^-1. A backtracking line search on ||F|| globalizes the + step, which is required on stiff 3D cases where the full Newton step overshoots. + """ + model = make_model(input_path, ns) + F = residual(model) + x = np.asarray(model.get_state(), float).copy() + n_dof = x.size + model.reset_force_eval_count() + t0 = time.perf_counter() + it = 0 + for _ in range(max_newton): + fk = F(x) + norm0 = np.linalg.norm(fk) + if norm0 < tol: + break + it += 1 + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, True) # assemble M at the current iterate + h_op = LinearOperator( # type: ignore[call-overload] + (n_dof, n_dof), + matvec=lambda v: np.asarray( # type: ignore[call-overload] + model.hessian_vector_product(np.ascontiguousarray(v)), float + ), + ) + m_op = LinearOperator( # type: ignore[call-overload] + (n_dof, n_dof), + matvec=lambda b: np.asarray( # type: ignore[call-overload] + model.apply_preconditioner(np.ascontiguousarray(b)), float + ), + ) + dx, _ = gmres(h_op, -fk, M=m_op, rtol=inner_tol, maxiter=100) + # Backtracking line search: accept the largest step that reduces ||F||. + alpha = 1.0 + for _ in range(30): + if np.linalg.norm(F(x + alpha * dx)) < norm0: + break + alpha *= 0.5 + else: + break # no decrease found; stop + x = x + alpha * dx + return _finish(model, "Newton (VMEC++ HVP + M^-1)", x, it, t0) + + +ALL_SOLVERS = ( + solve_vmecpp, + solve_preconditioned_descent, + solve_newton_krylov_preconditioned, + solve_newton_krylov, + solve_newton_hvp, +) + + +def main(): + _, w_star = reference_equilibrium() + print(f"reference equilibrium (native solve): W = {w_star:.8e}\n") + print( + f"{'optimizer':32s} {'F-evals':>8s} {'iters':>6s} {'time[s]':>8s} " + f"{'||F||':>10s} {'dW vs ref':>10s}" + ) + for solver in ALL_SOLVERS: + r = solver()[1] + print( + f"{r.name:32s} {r.force_evals:8d} {r.outer_iters:6d} {r.seconds:8.2f} " + f"{r.residual_norm:10.1e} {abs(r.energy - w_star):10.1e}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/fourier_resolution_increase.py b/examples/fourier_resolution_increase.py index 69dda7c91..375452153 100644 --- a/examples/fourier_resolution_increase.py +++ b/examples/fourier_resolution_increase.py @@ -10,14 +10,13 @@ resolution by :func:`vmecpp.interpolate_solution` (radial interpolation in ``sqrt(s)`` plus Fourier zero-padding). -``vmecpp.run_continuation`` performs the whole schedule in one call: +Setting ``VmecInput.mpol``/``.ntor`` to a sequence (instead of a plain int) makes +``vmecpp.run`` perform the whole schedule in one call: - output = vmecpp.run_continuation( - vmec_input, - ns_array=[15, 31, 31], - mpol_array=[5, 9, 13], - ntor_array=[4, 4, 4], - ) + vmec_input.ns_array = np.array([15, 31, 31]) + vmec_input.mpol = [5, 9, 13] + vmec_input.ntor = [4, 4, 4] + output = vmecpp.run(vmec_input) This example runs the schedule step by step instead, so it can report the per-step iteration counts and compare the total against solving at the target @@ -64,8 +63,8 @@ def step_input(ns: int, mpol: int, ntor: int) -> vmecpp.VmecInput: # --- Fourier continuation --------------------------------------------------- -# vmecpp.run_continuation(vmec_input, ns_array=ns_array, mpol_array=mpol_array, -# ntor_array=ntor_array) runs exactly this schedule in a single call; it is +# Setting vmec_input.ns_array/.mpol/.ntor to these schedules and calling +# vmecpp.run(vmec_input) runs exactly this schedule in a single call; it is # unrolled here to report the per-step iteration counts. print("Fourier continuation:") schedule = list(zip(ns_array, mpol_array, ntor_array, strict=True)) diff --git a/examples/visualize_magnetic_field.py b/examples/visualize_magnetic_field.py index 6b02ee7b5..43da0ca49 100644 --- a/examples/visualize_magnetic_field.py +++ b/examples/visualize_magnetic_field.py @@ -63,6 +63,7 @@ def calculate_magnetic_field( cosk = np.cos(kernel) sink = np.sin(kernel) cosk_nyq = np.cos(kernel_nyq) + sink_nyq = np.sin(kernel_nyq) # Position in cylindrical coordinates (Sec. 9.1) r = np.dot(rmnc[:, j], cosk) @@ -87,12 +88,31 @@ def calculate_magnetic_field( bphi = bsupv * r bz = bsupu * dzdtheta + bsupv * dzdphi - # Repeat the same process for the current density field + # Repeat the same process for the current density field. VMEC++ stores + # sqrt(g) J^theta and sqrt(g) J^zeta in currumnc/currvmnc, so divide by + # the Jacobian before transforming contravariant components to Cartesian + # coordinates. + gmnc = vmec_output.wout.gmnc + gmns = getattr(vmec_output.wout, "gmns", None) currumnc = vmec_output.wout.currumnc currvmnc = vmec_output.wout.currvmnc - - curru = np.dot(currumnc[:, j], cosk_nyq) - currv = np.dot(currvmnc[:, j], cosk_nyq) + currumns = vmec_output.wout.currumns + assert currumns is not None + currvmns = vmec_output.wout.currvmns + assert currvmns is not None + + sqrtg = np.dot(gmnc[:, j], cosk_nyq) + curru_sqrtg = np.dot(currumnc[:, j], cosk_nyq) + currv_sqrtg = np.dot(currvmnc[:, j], cosk_nyq) + if gmns is not None and gmns.size: + sqrtg += np.dot(gmns[:, j], sink_nyq) + if currumns is not None and currumns.size: + curru_sqrtg += np.dot(currumns[:, j], sink_nyq) + if currvmns is not None and currvmns.size: + currv_sqrtg += np.dot(currvmns[:, j], sink_nyq) + + curru = curru_sqrtg / sqrtg + currv = currv_sqrtg / sqrtg currr = curru * drdtheta + currv * drdphi currphi = currv * r 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/enzyme/local_force_hessian_test.cc b/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc new file mode 100644 index 000000000..302f0eb93 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc @@ -0,0 +1,315 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT + +// Exact Hessian of VMEC's local force map via composed autodiff. +// +// The shared force-chain kernels (Jacobian, metric, B^contra, B_cov, magnetic +// pressure, MHD force density, and the hybrid lambda force) compose into the +// local map +// g : real-space geometry -> real-space force density, +// the nonlinear core of VMEC's force. The full force is Táµ€ . g . T with the +// linear spectral transforms T, Táµ€; the exact force Hessian-vector product is +// therefore Táµ€ . J_g . T . v, and J_g is what Enzyme provides here. The +// remaining augmented term, the spectral-condensation constraint force, also +// carries a linear Fourier bandpass and is validated end-to-end against the +// finite-difference HVP in the pybind exact-HVP path. +// +// This test composes the production kernels over flat buffers and takes the +// Jacobian of g by forward and reverse mode, checks both against central finite +// differences and against each other. + +#include +#include +#include +#include +#include +#include + +#include "vmecpp/common/enzyme/enzyme.h" +#include "vmecpp/vmec/ideal_mhd_model/bco_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/bcontra_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/jacobian_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/lambda_force_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/metric_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/mhdforce_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/pressure_kernel.h" + +// Problem dimensions and the constant (non-differentiated) context. +struct Ctx { + int nZnT, nsH; // half-grid surfaces; full-grid has nsH + 1 + int nFull, nHalf; // (nsH+1)*nZnT, nsH*nZnT + const double* sqrtSF; // [nsH+1] + const double* sqrtSH; // [nsH] + const double* chipH; // [nsH] + const double* presH; // [nsH] + const double* radialBlending; // [nsH+1] + double deltaS; + double lamscale; + bool lthreed; +}; + +// 16 geometry blocks (each nFull) packed into geom; force blocks (each nHalf): +// 12 MHD force-density blocks + 4 lambda-force blocks (blmn_e/o, clmn_e/o). +// Everything else is intermediate work sliced from work. +enum { kGeomBlocks = 16, kForceBlocks = 16 }; + +// g: geometry -> force density, composing the MHD and lambda-force kernels. +__attribute__((noinline)) void LocalForce(const double* geom, double* work, + double* force, const Ctx* c) { + const int nF = c->nFull; + const int nH = c->nHalf; + const int nZnT = c->nZnT, nsH = c->nsH; + // geometry blocks + const double* r1e = geom + 0 * nF; + const double* r1o = geom + 1 * nF; + const double* z1e = geom + 2 * nF; + const double* z1o = geom + 3 * nF; + const double* rue = geom + 4 * nF; + const double* ruo = geom + 5 * nF; + const double* zue = geom + 6 * nF; + const double* zuo = geom + 7 * nF; + const double* rve = geom + 8 * nF; + const double* rvo = geom + 9 * nF; + const double* zve = geom + 10 * nF; + const double* zvo = geom + 11 * nF; + const double* lue = geom + 12 * nF; + const double* luo = geom + 13 * nF; + const double* lve = geom + 14 * nF; + const double* lvo = geom + 15 * nF; + // half-grid intermediates from work + double* p = work; + double* r12 = p; + p += nH; + double* ru12 = p; + p += nH; + double* zu12 = p; + p += nH; + double* rs = p; + p += nH; + double* zs = p; + p += nH; + double* tau = p; + p += nH; + double* gsqrt = p; + p += nH; + double* guu = p; + p += nH; + double* guv = p; + p += nH; + double* gvv = p; + p += nH; + double* bsupu = p; + p += nH; + double* bsupv = p; + p += nH; + double* bsubu = p; + p += nH; + double* bsubv = p; + p += nH; + double* tp = p; + p += nH; + // per-nZnT scratch for the force kernel (26 blocks) + double* sc = p; // 26 * nZnT + + vmecpp::ComputeHalfGridJacobian( + r1e, r1o, z1e, z1o, rue, ruo, zue, zuo, c->sqrtSH, c->deltaS, + /*dSHalfDsInterp=*/0.25, nZnT, 0, 0, nsH, r12, ru12, zu12, rs, zs, tau); + vmecpp::ComputeMetricElements(r1e, r1o, rue, ruo, zue, zuo, rve, rvo, zve, + zvo, tau, r12, c->sqrtSF, c->sqrtSH, c->lthreed, + nZnT, 0, 0, nsH, gsqrt, guu, guv, gvv); + vmecpp::ComputeBsupContra(lue, luo, lve, lvo, gsqrt, c->sqrtSH, c->lthreed, + nZnT, 0, 0, nsH, bsupu, bsupv); + for (int jH = 0; jH < nsH; ++jH) { + for (int kl = 0; kl < nZnT; ++kl) { + const int ih = jH * nZnT + kl; + bsupu[ih] += c->chipH[jH] / gsqrt[ih]; + } + } + vmecpp::ComputeBCo(guu, guv, gvv, bsupu, bsupv, c->lthreed, nH, bsubu, bsubv); + vmecpp::ComputeMagneticPressure(bsupu, bsubu, bsupv, bsubv, nH, tp); + for (int jH = 0; jH < nsH; ++jH) { + for (int kl = 0; kl < nZnT; ++kl) tp[jH * nZnT + kl] += c->presH[jH]; + } + double* s = sc; + double* P_i = s; + s += nZnT; + double* rup_i = s; + s += nZnT; + double* zup_i = s; + s += nZnT; + double* rsp_i = s; + s += nZnT; + double* zsp_i = s; + s += nZnT; + double* taup_i = s; + s += nZnT; + double* gbubu_i = s; + s += nZnT; + double* gbubv_i = s; + s += nZnT; + double* gbvbv_i = s; + s += nZnT; + double* P_o = s; + s += nZnT; + double* rup_o = s; + s += nZnT; + double* zup_o = s; + s += nZnT; + double* rsp_o = s; + s += nZnT; + double* zsp_o = s; + s += nZnT; + double* taup_o = s; + s += nZnT; + double* gbubu_o = s; + s += nZnT; + double* gbubv_o = s; + s += nZnT; + double* gbvbv_o = s; + s += nZnT; + double* P_avg = s; + s += nZnT; + double* P_wavg = s; + s += nZnT; + double* gbubu_avg = s; + s += nZnT; + double* gbubu_wavg = s; + s += nZnT; + double* gbvbv_avg = s; + s += nZnT; + double* gbvbv_wavg = s; + s += nZnT; + double* gbubv_avg = s; + s += nZnT; + double* gbubv_wavg = s; + s += nZnT; + // lambda-force radial-sweep scratch (carried inside half-grid point) + double* bsubu_i = s; + s += nZnT; + double* bsubv_i = s; + s += nZnT; + double* gvv_gsqrt_i = s; + s += nZnT; + double* guv_bsupu_i = s; + s += nZnT; + double* armn_e = force + 0 * nH; + double* armn_o = force + 1 * nH; + double* azmn_e = force + 2 * nH; + double* azmn_o = force + 3 * nH; + double* brmn_e = force + 4 * nH; + double* brmn_o = force + 5 * nH; + double* bzmn_e = force + 6 * nH; + double* bzmn_o = force + 7 * nH; + double* crmn_e = force + 8 * nH; + double* crmn_o = force + 9 * nH; + double* czmn_e = force + 10 * nH; + double* czmn_o = force + 11 * nH; + vmecpp::ComputeMHDForceDensity( + r1e, r1o, rue, ruo, zue, zuo, z1o, rve, rvo, zve, zvo, r12, ru12, zu12, + rs, zs, tau, tp, gsqrt, bsupu, bsupv, c->sqrtSF, c->sqrtSH, P_i, rup_i, + zup_i, rsp_i, zsp_i, taup_i, gbubu_i, gbubv_i, gbvbv_i, P_o, rup_o, zup_o, + rsp_o, zsp_o, taup_o, gbubu_o, gbubv_o, gbvbv_o, P_avg, P_wavg, gbubu_avg, + gbubu_wavg, gbvbv_avg, gbvbv_wavg, gbubv_avg, gbubv_wavg, c->deltaS, nZnT, + /*nsMinF=*/0, 0, 0, nsH, /*jMaxRZ=*/nsH, c->lthreed, armn_e, armn_o, + azmn_e, azmn_o, brmn_e, brmn_o, bzmn_e, bzmn_o, crmn_e, crmn_o, czmn_e, + czmn_o); + // lambda force (blmn_e/o, clmn_e/o) from the shared kernel + double* blmn_e = force + 12 * nH; + double* blmn_o = force + 13 * nH; + double* clmn_e = force + 14 * nH; + double* clmn_o = force + 15 * nH; + vmecpp::ComputeHybridLambdaForce( + bsubu, bsubv, gvv, gsqrt, guv, bsupu, lue, luo, c->sqrtSH, c->sqrtSF, + c->radialBlending, c->lamscale, c->lthreed, nZnT, /*nsMinF=*/0, + /*nsMinF1=*/0, /*nsMinH=*/0, nsH, /*nsMaxFIncludingLcfs=*/nsH, bsubu_i, + bsubv_i, gvv_gsqrt_i, guv_bsupu_i, blmn_e, blmn_o, clmn_e, clmn_o); +} + +// Scalar objective L = 0.5 ||force||^2; work and force are caller-owned +// scratch. +__attribute__((noinline)) double Loss(const double* geom, double* work, + double* force, const Ctx* c) { + LocalForce(geom, work, force, c); + const int n = kForceBlocks * c->nHalf; + double s = 0.0; + for (int i = 0; i < n; ++i) s += 0.5 * force[i] * force[i]; + return s; +} + +int main() { + const int nZnT = 24, nsH = 8; + Ctx c; + c.nZnT = nZnT; + c.nsH = nsH; + c.nFull = (nsH + 1) * nZnT; + c.nHalf = nsH * nZnT; + c.deltaS = 0.1; + c.lamscale = 0.7; + c.lthreed = true; + std::mt19937 rng(3); + std::uniform_real_distribution d(0.5, 1.5), s(-1.0, 1.0); + std::vector sqrtSF(nsH + 1), sqrtSH(nsH), chipH(nsH), presH(nsH), + radialBlending(nsH + 1); + for (int j = 0; j <= nsH; ++j) { + sqrtSF[j] = std::sqrt(0.05 + 0.9 * j / nsH); + radialBlending[j] = 0.3 + 0.4 * j / nsH; + } + for (int j = 0; j < nsH; ++j) { + sqrtSH[j] = std::sqrt(0.05 + 0.9 * (j + 0.5) / nsH); + chipH[j] = 0.3 + 0.1 * j; + presH[j] = 0.2 + 0.05 * j; + } + c.sqrtSF = sqrtSF.data(); + c.sqrtSH = sqrtSH.data(); + c.chipH = chipH.data(); + c.presH = presH.data(); + c.radialBlending = radialBlending.data(); + + const int nG = kGeomBlocks * c.nFull; + const int nW = 15 * c.nHalf + 30 * nZnT; + const int nFc = kForceBlocks * c.nHalf; + std::vector geom(nG), v(nG); + for (double& x : geom) x = d(rng); + for (double& x : v) x = s(rng); + std::vector work(nW, 0.0), dwork(nW, 0.0), force(nFc, 0.0), + dforce(nFc, 0.0), grev(nG, 0.0); + + // Reverse mode: full gradient dL/dgeom. + __enzyme_autodiff((void*)Loss, enzyme_dup, geom.data(), grev.data(), + enzyme_dup, work.data(), dwork.data(), enzyme_dup, + force.data(), dforce.data(), enzyme_const, &c); + // Forward mode: directional derivative dL.v. + const double dfwd = __enzyme_fwddiff( + (void*)Loss, enzyme_dup, geom.data(), v.data(), enzyme_dup, work.data(), + dwork.data(), enzyme_dup, force.data(), dforce.data(), enzyme_const, &c); + + const double h = 1e-6; + std::vector gp(geom), gm(geom), w2(nW), f2(nFc); + for (int i = 0; i < nG; ++i) { + gp[i] = geom[i] + h * v[i]; + gm[i] = geom[i] - h * v[i]; + } + const double dfd = (Loss(gp.data(), w2.data(), f2.data(), &c) - + Loss(gm.data(), w2.data(), f2.data(), &c)) / + (2 * h); + const double drev = + std::inner_product(grev.begin(), grev.end(), v.begin(), 0.0); + const double scale = std::fabs(dfd) + 1e-300; + + printf("exact Hessian of VMEC local force map (MHD + lambda kernels)\n"); + printf(" geom dofs=%d force outputs=%d\n", nG, nFc); + printf(" reverse dL.v vs finite-diff : %.2e\n", + std::fabs(drev - dfd) / scale); + printf(" forward dL.v vs finite-diff : %.2e\n", + std::fabs(dfwd - dfd) / scale); + printf(" forward / reverse agreement : %.2e\n", + std::fabs(dfwd - drev) / (std::fabs(drev) + 1e-300)); + + const bool ok = std::fabs(drev - dfd) < 1e-5 * scale && + std::fabs(dfwd - dfd) < 1e-5 * scale && + std::fabs(dfwd - drev) < 1e-9 * (std::fabs(drev) + 1e-300); + printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} 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/laplace_solver/BUILD.bazel b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/BUILD.bazel index 1a05d184e..d45bfdad1 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/BUILD.bazel @@ -30,3 +30,16 @@ cc_test( ], size = "small", ) + +cc_binary( + name = "laplace_solver_bench", + srcs = ["laplace_solver_bench.cc"], + deps = [ + ":laplace_solver", + "//vmecpp/common/fourier_basis_fast_toroidal", + "//vmecpp/common/sizes:sizes", + "//vmecpp/free_boundary/tangential_partitioning:tangential_partitioning", + "@google_benchmark//:benchmark_main", + ], + linkopts = ["-llapack"], +) diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_bench.cc b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_bench.cc new file mode 100644 index 000000000..752970dd8 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_bench.cc @@ -0,0 +1,206 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT + +// Microbenchmarks for the free-boundary Laplace solve (NESTOR). +// +// The dense linear system assembled and factorized here is the dominant +// free-boundary cost when NESTOR performs a full update: an +// mnpd x mnpd (mnpd = (mf+1)*(2*nf+1)) matrix is built, LU-factorized via +// LAPACK dgetrf, and back-substituted via dgetrs. We benchmark: +// +// * TransformGreensFunctionDerivative -- the largest of the Fourier +// transforms feeding the source term (uses the manufactured single-mode +// input from laplace_solver_test.cc). +// * BuildMatrix + DecomposeMatrix + SolveForPotential -- the assemble + +// factorize + solve chain, which is the actual "Laplace solve". +// +// Sizes are chosen to bracket real free-boundary runs (cth_like / solovev +// free-boundary configurations use ntor ~ 4-6, mpol ~ 5-8). + +#include +#include +#include +#include +#include +#include +#include + +#include "Eigen/Dense" +#include "benchmark/benchmark.h" +#include "vmecpp/common/fourier_basis_fast_toroidal/fourier_basis_fast_toroidal.h" +#include "vmecpp/common/sizes/sizes.h" +#include "vmecpp/free_boundary/laplace_solver/laplace_solver.h" +#include "vmecpp/free_boundary/tangential_partitioning/tangential_partitioning.h" + +namespace vmecpp { +namespace { + +struct ResParams { + int nfp; + int mpol; + int ntor; + const char* label; +}; + +// (nfp, mpol, ntor) sizes bracketing real free-boundary runs. +constexpr ResParams kResolutions[] = { + {5, 5, 4, "5x4"}, + {5, 8, 6, "8x6"}, + {5, 12, 8, "12x8"}, +}; + +// ---------------------------------------------------------------------------- +// Fixture matching laplace_solver_test.cc's minimal single-threaded setup: +// a non-mocked LaplaceSolver with flat matrixShare / iPiv / bvecShare backing +// storage. +// ---------------------------------------------------------------------------- +struct BenchFixture { + Sizes s; + FourierBasisFastToroidal fb; + TangentialPartitioning tp; + + int nf; + int mf; + int mnpd; + + std::vector matrixShare; + std::vector iPiv; + std::vector bvecShare; + + std::unique_ptr ls; + + // Seeded, well-conditioned system state re-applied before each timed solve. + Eigen::VectorXd amat_seed; // diagonally-dominant [mnpd*mnpd] + Eigen::VectorXd bvec_singular_seed; // [mnpd] + + // Manufactured single-mode Green's-function-derivative input. + Eigen::VectorXd greenp; + + explicit BenchFixture(int nfp, int mpol, int ntor) + : s(/*lasym=*/false, nfp, mpol, ntor, /*ntheta=*/0, + /*nzeta=*/4 * (ntor + 1)), + fb(&s), + tp(s.nZnT), + nf(ntor), + mf(mpol + 1), + mnpd((mf + 1) * (2 * nf + 1)), + matrixShare(mnpd * mnpd, 0.0), + iPiv(mnpd, 0), + bvecShare(mnpd, 0.0) { + ls = std::make_unique( + &s, &fb, &tp, nf, mf, std::span(matrixShare), + std::span(iPiv), std::span(bvecShare)); + + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + + // Diagonally-dominant matrix so dgetrf never hits a singular pivot. + amat_seed.resize(mnpd * mnpd); + for (int col = 0; col < mnpd; ++col) { + for (int row = 0; row < mnpd; ++row) { + amat_seed[col * mnpd + row] = dist(rng); + } + // BuildMatrix() adds +0.5 on the diagonal; add a large diagonal boost + // here too to guarantee well-conditioned LU. + amat_seed[col * mnpd + col] += 2.0 * mnpd; + } + + bvec_singular_seed.resize(mnpd); + for (int i = 0; i < mnpd; ++i) bvec_singular_seed[i] = dist(rng); + + // Manufactured single-mode greenp, as in laplace_solver_test.cc. + const int numLocal = tp.ztMax - tp.ztMin; + const int nThetaEven = s.nThetaEven; + const int nZeta = s.nZeta; + greenp = Eigen::VectorXd::Zero(numLocal * nThetaEven * nZeta); + const int m0 = std::min(2, mpol - 1); + const int n0 = std::min(1, ntor); + for (int klpRel = 0; klpRel < numLocal; ++klpRel) { + for (int l = 0; l < nThetaEven; ++l) { + const double theta = 2.0 * std::numbers::pi * l / nThetaEven; + for (int k = 0; k < nZeta; ++k) { + const double phi = 2.0 * std::numbers::pi * k / nZeta; + const int idx = klpRel * nThetaEven * nZeta + l * nZeta + k; + greenp[idx] = std::sin(m0 * theta - n0 * phi); + } + } + } + } + + // Re-seed the LaplaceSolver's system state to the well-conditioned baseline. + // dgetrf destroys matrixShare in place, so this must run before each + // decompose+solve iteration. + void ResetSystem() { + ls->amat_sin_sin = amat_seed; + ls->bvec_sin = Eigen::VectorXd::Zero(mnpd); + } +}; + +// The assemble + factorize + solve chain: this is the actual Laplace solve. +template +void BM_LaplaceSolve(benchmark::State& state) { + static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, + kResolutions[kIdx].ntor); + for (auto _ : state) { + state.PauseTiming(); + fx.ResetSystem(); + state.ResumeTiming(); + + fx.ls->BuildMatrix(); + fx.ls->DecomposeMatrix(); + fx.ls->SolveForPotential(fx.bvec_singular_seed); + benchmark::ClobberMemory(); + } + state.SetLabel(kResolutions[kIdx].label); +} + +// Just the LU factorization (dgetrf), the O(mnpd^3) hotspot. +template +void BM_LaplaceDecompose(benchmark::State& state) { + static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, + kResolutions[kIdx].ntor); + for (auto _ : state) { + state.PauseTiming(); + fx.ResetSystem(); + fx.ls->BuildMatrix(); + state.ResumeTiming(); + + fx.ls->DecomposeMatrix(); + benchmark::ClobberMemory(); + } + state.SetLabel(kResolutions[kIdx].label); +} + +// The largest Green's-function-derivative Fourier transform. +template +void BM_TransformGreensFunctionDerivative(benchmark::State& state) { + static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, + kResolutions[kIdx].ntor); + for (auto _ : state) { + fx.ls->TransformGreensFunctionDerivative(fx.greenp); + benchmark::ClobberMemory(); + } + state.SetLabel(kResolutions[kIdx].label); +} + +BENCHMARK_TEMPLATE(BM_LaplaceSolve, 0)->Name("LaplaceSolve/5x4"); +BENCHMARK_TEMPLATE(BM_LaplaceSolve, 1)->Name("LaplaceSolve/8x6"); +BENCHMARK_TEMPLATE(BM_LaplaceSolve, 2)->Name("LaplaceSolve/12x8"); + +BENCHMARK_TEMPLATE(BM_LaplaceDecompose, 0)->Name("LaplaceDecompose/5x4"); +BENCHMARK_TEMPLATE(BM_LaplaceDecompose, 1)->Name("LaplaceDecompose/8x6"); +BENCHMARK_TEMPLATE(BM_LaplaceDecompose, 2)->Name("LaplaceDecompose/12x8"); + +BENCHMARK_TEMPLATE(BM_TransformGreensFunctionDerivative, 0) + ->Name("TransformGreensFunctionDerivative/5x4"); +BENCHMARK_TEMPLATE(BM_TransformGreensFunctionDerivative, 1) + ->Name("TransformGreensFunctionDerivative/8x6"); +BENCHMARK_TEMPLATE(BM_TransformGreensFunctionDerivative, 2) + ->Name("TransformGreensFunctionDerivative/12x8"); + +} // namespace +} // namespace vmecpp + +BENCHMARK_MAIN(); 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 e8df6160b..7f5c618d6 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel @@ -115,3 +115,42 @@ cc_library( "@abseil-cpp//absl/status", ], ) + +cc_binary( + name = "fft_toroidal_bench", + srcs = ["fft_toroidal_bench.cc"], + deps = [ + ":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", + ], +) + +cc_binary( + name = "dealias_constraint_force_bench", + srcs = ["dealias_constraint_force_bench.cc"], + deps = [ + ":ideal_mhd_model", + "//vmecpp/common/flow_control:flow_control", + "//vmecpp/common/fourier_basis_fast_poloidal", + "//vmecpp/common/sizes:sizes", + "//vmecpp/common/vmec_indata:vmec_indata", + "//vmecpp/vmec/handover_storage:handover_storage", + "//vmecpp/vmec/radial_partitioning:radial_partitioning", + "//vmecpp/vmec/radial_profiles:radial_profiles", + "@eigen", + "@google_benchmark//:benchmark_main", + ], +) diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h new file mode 100644 index 000000000..b5078a6f1 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#ifndef VMECPP_VMEC_IDEAL_MHD_MODEL_CONSTRAINT_FORCE_KERNEL_H_ +#define VMECPP_VMEC_IDEAL_MHD_MODEL_CONSTRAINT_FORCE_KERNEL_H_ + +#include + +namespace vmecpp { + +// The two local (non-transform) pieces of VMEC's spectral-condensation +// constraint force. The Fourier-space bandpass between them is the separate +// shared free function deAliasConstraintForce. All full-grid buffers are +// indexed (jF-nsMinF)*nZnT except sqrtSF, which is indexed (jF-nsMinF1). +// Allocation-free, shared between the solver and the Enzyme autodiff path. + +// Effective constraint force gConEff = (rCon - rCon0) ru + (zCon - zCon0) zu. +// No constraint on the axis (it has no poloidal angle), so the axis surface is +// skipped, matching deAliasConstraintForce. +inline void ComputeEffectiveConstraintForce( + const double* rCon, const double* rCon0, const double* zCon, + const double* zCon0, const double* ruFull, const double* zuFull, int nZnT, + int nsMinF, int nsMaxFIncludingLcfs, double* gConEff) { + int jMin = 0; + if (nsMinF == 0) { + jMin = 1; + } + for (int jF = std::max(jMin, nsMinF); jF < nsMaxFIncludingLcfs; ++jF) { + for (int kl = 0; kl < nZnT; ++kl) { + int idx_kl = (jF - nsMinF) * nZnT + kl; + gConEff[idx_kl] = (rCon[idx_kl] - rCon0[idx_kl]) * ruFull[idx_kl] + + (zCon[idx_kl] - zCon0[idx_kl]) * zuFull[idx_kl]; + } // kl + } // jF +} + +// Add the bandpass-filtered constraint force gCon back into the MHD R/Z forces +// (brmn, bzmn) and write the constraint-force outputs (frcon, fzcon). +inline void AddConstraintForces( + const double* rCon, const double* rCon0, const double* zCon, + const double* zCon0, const double* ruFull, const double* zuFull, + const double* gCon, const double* sqrtSF, int nZnT, int nsMinF, int nsMinF1, + int nsMaxF, double* brmn_e, double* brmn_o, double* bzmn_e, double* bzmn_o, + double* frcon_e, double* frcon_o, double* fzcon_e, double* fzcon_o) { + for (int jF = nsMinF; jF < nsMaxF; ++jF) { + for (int kl = 0; kl < nZnT; ++kl) { + int idx_kl = (jF - nsMinF) * nZnT + kl; + + double brcon = (rCon[idx_kl] - rCon0[idx_kl]) * gCon[idx_kl]; + double bzcon = (zCon[idx_kl] - zCon0[idx_kl]) * gCon[idx_kl]; + + brmn_e[idx_kl] += brcon; + bzmn_e[idx_kl] += bzcon; + brmn_o[idx_kl] += brcon * sqrtSF[jF - nsMinF1]; + bzmn_o[idx_kl] += bzcon * sqrtSF[jF - nsMinF1]; + + frcon_e[idx_kl] = ruFull[idx_kl] * gCon[idx_kl]; + fzcon_e[idx_kl] = zuFull[idx_kl] * gCon[idx_kl]; + frcon_o[idx_kl] = frcon_e[idx_kl] * sqrtSF[jF - nsMinF1]; + fzcon_o[idx_kl] = fzcon_e[idx_kl] * sqrtSF[jF - nsMinF1]; + } // kl + } // jF +} + +} // namespace vmecpp + +#endif // VMECPP_VMEC_IDEAL_MHD_MODEL_CONSTRAINT_FORCE_KERNEL_H_ diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dealias_constraint_force_bench.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dealias_constraint_force_bench.cc new file mode 100644 index 000000000..b8c553bdd --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/dealias_constraint_force_bench.cc @@ -0,0 +1,132 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT + +// Microbenchmark for deAliasConstraintForce, the spectral-condensation +// constraint-force de-aliasing step (a forward+inverse poloidal/toroidal +// transform pair) run every iteration inside IdealMhdModel::update(). +// +// Sweeps the same four representative (nfp, mpol, ntor) resolutions used by +// fft_toroidal_bench.cc so the constraint-force cost can be compared directly +// against the geometry/force transforms at each resolution. + +#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/vmec_indata/vmec_indata.h" +#include "vmecpp/vmec/handover_storage/handover_storage.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" + +namespace vmecpp { +namespace { + +// ns = 51 is representative of a medium-resolution VMEC run (matches the +// convention used in fft_toroidal_bench.cc). +constexpr int kNs = 51; + +// (nfp, mpol, ntor) resolutions, matching fft_toroidal_bench.cc for direct +// comparability. +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"}, +}; + +// ---------------------------------------------------------------------------- +// Shared fixture data, built once per (nfp, mpol, ntor) combination. +// +// Mirrors the buffer sizes established in IdealMhdModel::allocate(): +// faccon: s.mpol, with faccon[m-1] = -0.25 / xmpq[m]^2 for m > 1 +// tcon: nsMaxFIncludingLcfs - nsMinF +// gConEff: nZnT * (nsMaxFIncludingLcfs - nsMinF) (nrztIncludingBoundary) +// gsc/gcs: ntor + 1 +// gCon: nrztIncludingBoundary +// ---------------------------------------------------------------------------- +struct BenchFixture { + Sizes s; + RadialPartitioning rp; + FourierBasisFastPoloidal fb; + + Eigen::VectorXd faccon; + Eigen::VectorXd tcon; + Eigen::VectorXd gConEff; + Eigen::VectorXd gsc; + Eigen::VectorXd gcs; + Eigen::VectorXd gCon; + + explicit BenchFixture(int nfp, int mpol, int ntor) + : s(/*lasym=*/false, nfp, mpol, ntor, /*ntheta=*/0, /*nzeta=*/0), fb(&s) { + rp.adjustRadialPartitioning(/*num_threads=*/1, /*thread_id=*/0, kNs, + /*lfreeb=*/false, /*printout=*/false); + + const int nrztIncludingBoundary = + s.nZnT * (rp.nsMaxFIncludingLcfs - rp.nsMinF); + + // faccon[m] = -0.25 / (m*(m-1))^2 for m > 1 (signOfJacobian = -1 gives the + // same magnitude sign convention as the solver; the exact sign does not + // affect the timing). + faccon.setZero(s.mpol); + for (int m = 2; m < s.mpol; ++m) { + const double xmpq = m * (m - 1); + faccon[m - 1] = -0.25 * (-1) / (xmpq * xmpq); + } + + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + auto rfill = [&](Eigen::VectorXd& v, int n) { + v.resize(n); + for (int i = 0; i < n; ++i) v[i] = dist(rng); + }; + + rfill(tcon, rp.nsMaxFIncludingLcfs - rp.nsMinF); + rfill(gConEff, nrztIncludingBoundary); + + gsc.setZero(s.ntor + 1); + gcs.setZero(s.ntor + 1); + gCon.setZero(nrztIncludingBoundary); + } +}; + +template +void BM_DeAliasConstraintForce(benchmark::State& state) { + static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, + kResolutions[kIdx].ntor); + for (auto _ : state) { + deAliasConstraintForce(fx.rp, fx.fb, fx.s, fx.faccon, fx.tcon, fx.gConEff, + fx.gsc, fx.gcs, fx.gCon); + benchmark::ClobberMemory(); + } + state.SetLabel(kResolutions[kIdx].label); +} + +BENCHMARK_TEMPLATE(BM_DeAliasConstraintForce, 0) + ->Name("DeAliasConstraintForce/4x4"); +BENCHMARK_TEMPLATE(BM_DeAliasConstraintForce, 1) + ->Name("DeAliasConstraintForce/7x1"); +BENCHMARK_TEMPLATE(BM_DeAliasConstraintForce, 2) + ->Name("DeAliasConstraintForce/12x12"); +BENCHMARK_TEMPLATE(BM_DeAliasConstraintForce, 3) + ->Name("DeAliasConstraintForce/16x18"); + +} // namespace +} // namespace vmecpp + +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 04191d42b..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,410 +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) { +void BM_ToroidalTransform_FourierToReal(benchmark::State& state) { static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, kResolutions[kIdx].ntor); for (auto _ : state) { - FourierToReal3DSymmFastPoloidalFft(*fx.phys_x, fx.xmpq, fx.rp, fx.s, - *fx.rprof, fx.fb, fx.plans, fx.geom); + fx.model->dft_FourierToReal_3d_symm(*fx.phys_x); benchmark::ClobberMemory(); } state.SetLabel(kResolutions[kIdx].label); } template -void BM_ForcesToFourier(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) { - ForcesToFourier3DSymmFastPoloidalFft(fx.forces, fx.xmpq, fx.rp, *fx.fc, - fx.s, fx.fb, fx.plans, - VacuumPressureState::kOff, *fx.ff); + fx.model->dft_ForcesToFourier_3d_symm(*fx.phys_f); benchmark::ClobberMemory(); } state.SetLabel(kResolutions[kIdx].label); } -template -void BM_DftFourierToReal(benchmark::State& state) { - static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, - kResolutions[kIdx].ntor); - for (auto _ : state) { - FourierToReal3DSymmFastPoloidal(*fx.phys_x, fx.xmpq, fx.rp, fx.s, *fx.rprof, - fx.fb, fx.geom); - 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); - 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_RS(I, NAME) \ - BENCHMARK_TEMPLATE(BM_DftFourierToReal_RealSpace, I)->Name("Dft/" NAME); \ - BENCHMARK_TEMPLATE(BM_FftFourierToReal_RealSpace, I)->Name("Fft/" NAME) - -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; +#define REGISTER_RES(IDX, LABEL) \ + BENCHMARK_TEMPLATE(BM_ToroidalTransform_FourierToReal, IDX) \ + ->Name("ToroidalFourierToReal/" LABEL); \ + BENCHMARK_TEMPLATE(BM_ToroidalTransform_ForcesToFourier, IDX) \ + ->Name("ToroidalForcesToFourier/" LABEL) - // 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; - }; +REGISTER_RES(0, "4x4"); +REGISTER_RES(1, "6x8"); +REGISTER_RES(2, "12x12"); +REGISTER_RES(3, "12x13"); - 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); - - 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 20a005679..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 @@ -22,7 +22,9 @@ #include "vmecpp/vmec/handover_storage/handover_storage.h" #include "vmecpp/vmec/ideal_mhd_model/bco_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/bcontra_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/jacobian_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/lambda_force_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/metric_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/mhdforce_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/pressure_kernel.h" @@ -419,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) { @@ -811,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 @@ -1628,121 +1633,16 @@ void IdealMhdModel::hybridLambdaForce() { #pragma omp barrier #endif // _OPENMP - // obtain first inside point - int j0 = r_.nsMinF; - double sqrtSHi = 0.0; - if (j0 > 0) { - sqrtSHi = m_p_.sqrtSH[j0 - 1 - r_.nsMinH]; - } - for (int kl = 0; kl < s_.nZnT; ++kl) { - if (j0 == 0) { - // defaults to 0: no contribution from half-grid point inside the axis - m_ls_.bsubu_i[kl] = 0.0; - m_ls_.bsubv_i[kl] = 0.0; - m_ls_.gvv_gsqrt_i[kl] = 0.0; // gvv / gsqrt - m_ls_.guv_bsupu_i[kl] = 0.0; // guv * bsupu - } else { - // for the j-th forces full-grid point, the (j-1)-th half-grid point is - // inside - int iHalf = (j0 - 1 - r_.nsMinH) * s_.nZnT + kl; - m_ls_.bsubu_i[kl] = bsubu[iHalf]; - m_ls_.bsubv_i[kl] = bsubv[iHalf]; - m_ls_.gvv_gsqrt_i[kl] = gvv[iHalf] / gsqrt[iHalf]; - if (s_.lthreed) { - m_ls_.guv_bsupu_i[kl] = guv[iHalf] * bsupu[iHalf]; - } - } - } // kl - - for (int jF = r_.nsMinF; jF < r_.nsMaxFIncludingLcfs; ++jF) { - double sqrtSHo = 0.0; - if (jF < r_.nsMaxH) { - sqrtSHo = m_p_.sqrtSH[jF - r_.nsMinH]; - } - - for (int kl = 0; kl < s_.nZnT; ++kl) { - // obtain next outside point - // defaults to 0: no contribution from half-grid point outside LCFS - double bsubv_o = 0.0; - // gvv / gsqrt - double gvv_gsqrt_o = 0.0; - // guv * bsupu - double guv_bsupu_o = 0.0; - if (jF < r_.nsMaxH) { - // for the j-th forces full-grid point, the j-th half-grid point is - // outside - int iHalf = (jF - r_.nsMinH) * s_.nZnT + kl; - bsubv_o = bsubv[iHalf]; - gvv_gsqrt_o = gvv[iHalf] / gsqrt[iHalf]; - if (s_.lthreed) { - guv_bsupu_o = guv[iHalf] * bsupu[iHalf]; - } - } - - // alternative way to interpolate bsubv onto the full-grid - double gvv_gsqrt_lu_e = 0.5 * (m_ls_.gvv_gsqrt_i[kl] + gvv_gsqrt_o) * - lu_e[(jF - r_.nsMinF1) * s_.nZnT + kl]; - double gvv_gsqrt_lu_o = - 0.5 * (m_ls_.gvv_gsqrt_i[kl] * sqrtSHi + gvv_gsqrt_o * sqrtSHo) * - lu_o[(jF - r_.nsMinF1) * s_.nZnT + kl]; - - double gvv_gsqrt_lu = gvv_gsqrt_lu_e + gvv_gsqrt_lu_o; - double bsubv_alternative = gvv_gsqrt_lu; - if (s_.lthreed) { - double guv_bsupu = 0.5 * (m_ls_.guv_bsupu_i[kl] + guv_bsupu_o); - bsubv_alternative += guv_bsupu; - } - - const double bsubv_average = 0.5 * (bsubv_o + m_ls_.bsubv_i[kl]); - - // blend together two ways of interpolating bsubv - double _blmn = - bsubv_average * (1.0 - m_p_.radialBlending[jF - r_.nsMinF1]) + - bsubv_alternative * m_p_.radialBlending[jF - r_.nsMinF1]; - - if (jF > 0) { - // TODO(jons): no lamscale and (-1) factor for axis lambda force? - // MINUS SIGN => HESSIAN DIAGONALS ARE POSITIVE - _blmn *= -constants_.lamscale; - } - - blmn_e[(jF - r_.nsMinF) * s_.nZnT + kl] = _blmn; - blmn_o[(jF - r_.nsMinF) * s_.nZnT + kl] = - _blmn * m_p_.sqrtSF[jF - r_.nsMinF1]; - - if (s_.lthreed) { - // obtain next outside point - // defaults to 0 for half-grid point outside LCFS - double bsubu_o = 0.0; - if (jF < r_.nsMaxH) { - bsubu_o = bsubu[(jF - r_.nsMinH) * s_.nZnT + kl]; - } - - double _clmn = 0.5 * (bsubu_o + m_ls_.bsubu_i[kl]); - - if (jF > 0) { - // TODO(jons): no lamscale and (-1) factor for axis lambda force? - // MINUS SIGN => HESSIAN DIAGONALS ARE POSITIVE - _clmn *= -constants_.lamscale; - } - - clmn_e[(jF - r_.nsMinF) * s_.nZnT + kl] = _clmn; - clmn_o[(jF - r_.nsMinF) * s_.nZnT + kl] = - _clmn * m_p_.sqrtSF[jF - r_.nsMinF1]; - - // shift to next point - m_ls_.bsubu_i[kl] = bsubu_o; - } // lthreed - - // shift to next point - m_ls_.bsubv_i[kl] = bsubv_o; - m_ls_.gvv_gsqrt_i[kl] = gvv_gsqrt_o; - if (s_.lthreed) { - m_ls_.guv_bsupu_i[kl] = guv_bsupu_o; - } - } // kl - sqrtSHi = sqrtSHo; - } // jF + // Lambda force on the full grid via the shared kernel + // (lambda_force_kernel.h), also used by the Enzyme autodiff path. + ComputeHybridLambdaForce( + bsubu.data(), bsubv.data(), gvv.data(), gsqrt.data(), guv.data(), + bsupu.data(), lu_e.data(), lu_o.data(), m_p_.sqrtSH.data(), + m_p_.sqrtSF.data(), m_p_.radialBlending.data(), constants_.lamscale, + s_.lthreed, s_.nZnT, r_.nsMinF, r_.nsMinF1, r_.nsMinH, r_.nsMaxH, + r_.nsMaxFIncludingLcfs, m_ls_.bsubu_i.data(), m_ls_.bsubv_i.data(), + m_ls_.gvv_gsqrt_i.data(), m_ls_.guv_bsupu_i.data(), blmn_e.data(), + blmn_o.data(), clmn_e.data(), clmn_o.data()); // } #ifdef _OPENMP @@ -2169,21 +2069,12 @@ absl::Status IdealMhdModel::constraintForceMultiplier() { } void IdealMhdModel::effectiveConstraintForce() { - // gConEff - - // no constraint on axis --> has no poloidal angle - int jMin = 0; - if (r_.nsMinF == 0) { - jMin = 1; - } - - for (int jF = std::max(jMin, r_.nsMinF); jF < r_.nsMaxFIncludingLcfs; ++jF) { - for (int kl = 0; kl < s_.nZnT; ++kl) { - int idx_kl = (jF - r_.nsMinF) * s_.nZnT + kl; - gConEff[idx_kl] = (rCon[idx_kl] - rCon0[idx_kl]) * ruFull[idx_kl] + - (zCon[idx_kl] - zCon0[idx_kl]) * zuFull[idx_kl]; - } // kl - } // jF + // gConEff via the shared kernel (constraint_force_kernel.h), also used by the + // Enzyme autodiff path. + ComputeEffectiveConstraintForce(rCon.data(), rCon0.data(), zCon.data(), + zCon0.data(), ruFull.data(), zuFull.data(), + s_.nZnT, r_.nsMinF, r_.nsMaxFIncludingLcfs, + gConEff.data()); } // perform Fourier-space bandpass filtering of constraint force @@ -2214,24 +2105,15 @@ void IdealMhdModel::assembleTotalForces() { } } - for (int jF = r_.nsMinF; jF < r_.nsMaxF; ++jF) { - for (int kl = 0; kl < s_.nZnT; ++kl) { - int idx_kl = (jF - r_.nsMinF) * s_.nZnT + kl; - - double brcon = (rCon[idx_kl] - rCon0[idx_kl]) * gCon[idx_kl]; - double bzcon = (zCon[idx_kl] - zCon0[idx_kl]) * gCon[idx_kl]; - - brmn_e[idx_kl] += brcon; - bzmn_e[idx_kl] += bzcon; - brmn_o[idx_kl] += brcon * m_p_.sqrtSF[jF - r_.nsMinF1]; - bzmn_o[idx_kl] += bzcon * m_p_.sqrtSF[jF - r_.nsMinF1]; - - frcon_e[idx_kl] = ruFull[idx_kl] * gCon[idx_kl]; - fzcon_e[idx_kl] = zuFull[idx_kl] * gCon[idx_kl]; - frcon_o[idx_kl] = frcon_e[idx_kl] * m_p_.sqrtSF[jF - r_.nsMinF1]; - fzcon_o[idx_kl] = fzcon_e[idx_kl] * m_p_.sqrtSF[jF - r_.nsMinF1]; - } - } + // add the bandpass-filtered constraint force to the MHD R/Z forces and write + // the constraint outputs via the shared kernel (constraint_force_kernel.h), + // also used by the Enzyme autodiff path. + AddConstraintForces(rCon.data(), rCon0.data(), zCon.data(), zCon0.data(), + ruFull.data(), zuFull.data(), gCon.data(), + m_p_.sqrtSF.data(), s_.nZnT, r_.nsMinF, r_.nsMinF1, + r_.nsMaxF, brmn_e.data(), brmn_o.data(), bzmn_e.data(), + bzmn_o.data(), frcon_e.data(), frcon_o.data(), + fzcon_e.data(), fzcon_o.data()); } void IdealMhdModel::forcesToFourier(FourierForces& m_physical_f) { 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/ideal_mhd_model/lambda_force_kernel.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/lambda_force_kernel.h new file mode 100644 index 000000000..105df3730 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/lambda_force_kernel.h @@ -0,0 +1,132 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#ifndef VMECPP_VMEC_IDEAL_MHD_MODEL_LAMBDA_FORCE_KERNEL_H_ +#define VMECPP_VMEC_IDEAL_MHD_MODEL_LAMBDA_FORCE_KERNEL_H_ + +namespace vmecpp { + +// Hybrid lambda force on the full grid: blmn (and clmn in 3D). The covariant +// field bsubv is interpolated from the half grid two ways (a plain average and +// an alternative from gvv/gsqrt times lambda plus, in 3D, guv*bsupu) and +// blended with the radialBlending profile. The radial sweep carries the inside +// half-grid point in scratch (bsubu_i, bsubv_i, gvv_gsqrt_i, guv_bsupu_i, each +// nZnT) and shifts it outward each surface. Allocation-free over flat buffers, +// shared between IdealMhdModel::hybridLambdaForce and the Enzyme autodiff path. +// +// Half-grid inputs (index (jH-nsMinH)*nZnT): bsubu, bsubv, gvv, gsqrt, guv, +// bsupu. Full-grid lambda (index (jF-nsMinF1)*nZnT): lu_e, lu_o. Profiles: +// sqrtSH (half), sqrtSF and radialBlending (full, index jF-nsMinF1). Outputs +// (index (jF-nsMinF)*nZnT): blmn_e/o, clmn_e/o. +inline void ComputeHybridLambdaForce( + const double* bsubu, const double* bsubv, const double* gvv, + const double* gsqrt, const double* guv, const double* bsupu, + const double* lu_e, const double* lu_o, const double* sqrtSH, + const double* sqrtSF, const double* radialBlending, double lamscale, + bool lthreed, int nZnT, int nsMinF, int nsMinF1, int nsMinH, int nsMaxH, + int nsMaxFIncludingLcfs, double* bsubu_i, double* bsubv_i, + double* gvv_gsqrt_i, double* guv_bsupu_i, double* blmn_e, double* blmn_o, + double* clmn_e, double* clmn_o) { + // first inside point + int j0 = nsMinF; + double sqrtSHi = 0.0; + if (j0 > 0) { + sqrtSHi = sqrtSH[j0 - 1 - nsMinH]; + } + for (int kl = 0; kl < nZnT; ++kl) { + if (j0 == 0) { + // no contribution from the half-grid point inside the axis + bsubu_i[kl] = 0.0; + bsubv_i[kl] = 0.0; + gvv_gsqrt_i[kl] = 0.0; + guv_bsupu_i[kl] = 0.0; + } else { + int iHalf = (j0 - 1 - nsMinH) * nZnT + kl; + bsubu_i[kl] = bsubu[iHalf]; + bsubv_i[kl] = bsubv[iHalf]; + gvv_gsqrt_i[kl] = gvv[iHalf] / gsqrt[iHalf]; + if (lthreed) { + guv_bsupu_i[kl] = guv[iHalf] * bsupu[iHalf]; + } + } + } // kl + + for (int jF = nsMinF; jF < nsMaxFIncludingLcfs; ++jF) { + double sqrtSHo = 0.0; + if (jF < nsMaxH) { + sqrtSHo = sqrtSH[jF - nsMinH]; + } + + for (int kl = 0; kl < nZnT; ++kl) { + // next outside point (defaults to 0 outside the LCFS) + double bsubv_o = 0.0; + double gvv_gsqrt_o = 0.0; + double guv_bsupu_o = 0.0; + if (jF < nsMaxH) { + int iHalf = (jF - nsMinH) * nZnT + kl; + bsubv_o = bsubv[iHalf]; + gvv_gsqrt_o = gvv[iHalf] / gsqrt[iHalf]; + if (lthreed) { + guv_bsupu_o = guv[iHalf] * bsupu[iHalf]; + } + } + + double gvv_gsqrt_lu_e = 0.5 * (gvv_gsqrt_i[kl] + gvv_gsqrt_o) * + lu_e[(jF - nsMinF1) * nZnT + kl]; + double gvv_gsqrt_lu_o = + 0.5 * (gvv_gsqrt_i[kl] * sqrtSHi + gvv_gsqrt_o * sqrtSHo) * + lu_o[(jF - nsMinF1) * nZnT + kl]; + + double gvv_gsqrt_lu = gvv_gsqrt_lu_e + gvv_gsqrt_lu_o; + double bsubv_alternative = gvv_gsqrt_lu; + if (lthreed) { + double guv_bsupu = 0.5 * (guv_bsupu_i[kl] + guv_bsupu_o); + bsubv_alternative += guv_bsupu; + } + + const double bsubv_average = 0.5 * (bsubv_o + bsubv_i[kl]); + + // blend the two interpolations of bsubv + double _blmn = bsubv_average * (1.0 - radialBlending[jF - nsMinF1]) + + bsubv_alternative * radialBlending[jF - nsMinF1]; + + if (jF > 0) { + // MINUS SIGN => HESSIAN DIAGONALS ARE POSITIVE + _blmn *= -lamscale; + } + + blmn_e[(jF - nsMinF) * nZnT + kl] = _blmn; + blmn_o[(jF - nsMinF) * nZnT + kl] = _blmn * sqrtSF[jF - nsMinF1]; + + if (lthreed) { + double bsubu_o = 0.0; + if (jF < nsMaxH) { + bsubu_o = bsubu[(jF - nsMinH) * nZnT + kl]; + } + + double _clmn = 0.5 * (bsubu_o + bsubu_i[kl]); + + if (jF > 0) { + _clmn *= -lamscale; + } + + clmn_e[(jF - nsMinF) * nZnT + kl] = _clmn; + clmn_o[(jF - nsMinF) * nZnT + kl] = _clmn * sqrtSF[jF - nsMinF1]; + + bsubu_i[kl] = bsubu_o; + } // lthreed + + bsubv_i[kl] = bsubv_o; + gvv_gsqrt_i[kl] = gvv_gsqrt_o; + if (lthreed) { + guv_bsupu_i[kl] = guv_bsupu_o; + } + } // kl + sqrtSHi = sqrtSHo; + } // jF +} + +} // namespace vmecpp + +#endif // VMECPP_VMEC_IDEAL_MHD_MODEL_LAMBDA_FORCE_KERNEL_H_ diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel index 951ca9324..6a6f895e1 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel @@ -20,6 +20,24 @@ cc_library( ], ) +cc_binary( + name = "output_quantities_bench", + srcs = ["output_quantities_bench.cc"], + data = [ + "//vmecpp/test_data:cma", + ], + deps = [ + ":output_quantities", + "@abseil-cpp//absl/status:statusor", + "@abseil-cpp//absl/strings:str_format", + "@google_benchmark//:benchmark_main", + "//util/file_io:file_io", + "//vmecpp/common/util:util", + "//vmecpp/common/vmec_indata:vmec_indata", + "//vmecpp/vmec/vmec:vmec", + ], +) + cc_test( name = "output_quantities_io_test", srcs = ["output_quantities_io_test.cc", "test_helpers.h"], diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc index 23445edc4..c39f402e6 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc @@ -5168,7 +5168,12 @@ vmecpp::WOutFileContents vmecpp::ComputeWOutFileContents( void vmecpp::CompareWOut(const WOutFileContents& test_wout, const WOutFileContents& expected_wout, - double tolerance, bool check_equal_niter) { + double tolerance, bool check_equal_niter, + double current_density_tolerance) { + // jcuru, jcurv compare looser only if the caller opts in; otherwise + // tolerance. + const double current_tolerance = + current_density_tolerance > 0.0 ? current_density_tolerance : tolerance; CHECK_EQ(test_wout.signgs, expected_wout.signgs); CHECK_EQ(test_wout.gamma, expected_wout.gamma); CHECK_EQ(test_wout.pcurr_type, expected_wout.pcurr_type); @@ -5250,11 +5255,11 @@ void vmecpp::CompareWOut(const WOutFileContents& test_wout, IsCloseRelAbs(expected_wout.phipf[jF], test_wout.phipf[jF], tolerance)); CHECK( IsCloseRelAbs(expected_wout.chipf[jF], test_wout.chipf[jF], tolerance)); - CHECK( - IsCloseRelAbs(expected_wout.jcuru[jF], test_wout.jcuru[jF], tolerance)) + CHECK(IsCloseRelAbs(expected_wout.jcuru[jF], test_wout.jcuru[jF], + current_tolerance)) << "jF = " << jF; - CHECK( - IsCloseRelAbs(expected_wout.jcurv[jF], test_wout.jcurv[jF], tolerance)) + CHECK(IsCloseRelAbs(expected_wout.jcurv[jF], test_wout.jcurv[jF], + current_tolerance)) << "jF = " << jF; CHECK( IsCloseRelAbs(expected_wout.specw[jF], test_wout.specw[jF], tolerance)); diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.h b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.h index 456535bcc..bf3a0814b 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.h +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.h @@ -1521,10 +1521,15 @@ WOutFileContents ComputeWOutFileContents( // Compare the contents of a test wout object against a reference wout object, // exiting with an error in case of mismatches. // The comparison is performed using the specified tolerance in the "relabs" -// metric. +// metric. The current densities jcuru, jcurv are curl(B) quantities that +// amplify the rounding of the converged state; under vectorized/optimized +// builds they do not reproduce to the same tolerance as the profiles across +// machines. Pass current_density_tolerance > 0 to compare those two with a +// looser bound while keeping every other quantity at tolerance. void CompareWOut(const WOutFileContents& test_wout, const WOutFileContents& expected_wout, double tolerance, - bool check_equal_niter = true); + bool check_equal_niter = true, + double current_density_tolerance = 0.0); } // namespace vmecpp #endif // VMECPP_VMEC_OUTPUT_QUANTITIES_OUTPUT_QUANTITIES_H_ diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_bench.cc b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_bench.cc new file mode 100644 index 000000000..7b9ff8600 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_bench.cc @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT + +// Microbenchmark for ComputeOutputQuantities, the post-solve output +// computation. +// +// ComputeOutputQuantities() runs once after the equilibrium has converged and +// produces the entire wout contents, bsubs on both grids, jxbout, Mercier +// stability, and the other diagnostics. None of this feeds back into the +// force-balance loop, so it is pure post-processing overhead -- this benchmark +// answers "how much time do we spend after the actual solve has finished". +// +// Setup (untimed): load a small fixed-boundary case, run the solver to +// convergence via Vmec::run(). Timed loop: re-invoke ComputeOutputQuantities +// on the converged state. ComputeOutputQuantities takes all inputs by const +// reference and returns a fresh OutputQuantities by value, so repeated calls +// are idempotent. + +#include +#include +#include + +#include "absl/strings/str_format.h" +#include "benchmark/benchmark.h" +#include "util/file_io/file_io.h" +#include "vmecpp/common/util/util.h" +#include "vmecpp/common/vmec_indata/vmec_indata.h" +#include "vmecpp/vmec/output_quantities/output_quantities.h" +#include "vmecpp/vmec/vmec/vmec.h" + +namespace vmecpp { +namespace { + +// Small fixed-boundary stellarator case; converges quickly during setup. +constexpr char kCase[] = "cma"; + +void BM_ComputeOutputQuantities(benchmark::State& state) { + // ---- Untimed setup: drive the solver to convergence. ---- + const std::string filename = + absl::StrFormat("vmecpp/test_data/%s.json", kCase); + absl::StatusOr indata_json = file_io::ReadFile(filename); + if (!indata_json.ok()) { + state.SkipWithError("failed to read input JSON"); + return; + } + absl::StatusOr indata = VmecINDATA::FromJson(*indata_json); + if (!indata.ok()) { + state.SkipWithError("failed to parse INDATA"); + return; + } + + absl::StatusOr> maybe_vmec = Vmec::FromIndata(*indata); + if (!maybe_vmec.ok()) { + state.SkipWithError("Vmec::FromIndata failed"); + return; + } + Vmec& vmec = **maybe_vmec; + + absl::StatusOr ran = vmec.run(); + if (!ran.ok()) { + state.SkipWithError("vmec.run() failed"); + return; + } + + constexpr int kSignOfJacobian = -1; + + // ---- Timed loop: re-run only the post-solve output computation. ---- + for (auto _ : state) { + OutputQuantities output_quantities = ComputeOutputQuantities( + kSignOfJacobian, vmec.indata_, vmec.s_, vmec.fc_, vmec.constants_, + vmec.t_, vmec.h_, vmec.mgrid_.mgrid_mode, vmec.r_, vmec.decomposed_x_, + vmec.m_, vmec.p_, VmecCheckpoint::NONE, + static_cast(vmec.get_ivac()), vmec.get_status(), + vmec.get_iter2()); + benchmark::DoNotOptimize(output_quantities.wout.volume); + benchmark::ClobberMemory(); + } + state.SetLabel(kCase); +} + +BENCHMARK(BM_ComputeOutputQuantities)->Name("ComputeOutputQuantities/cma"); + +} // namespace +} // namespace vmecpp + +BENCHMARK_MAIN(); 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 73c089227..5e2510d7d 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc @@ -217,9 +217,27 @@ class VmecModel { // decision vector. This is the body of Vmec::UpdateForwardModel with // caller-supplied counters; for free-boundary runs the caller should react to // `need_restart` (a vacuum-activation restart). - void Evaluate(int iter1, int iter2) { + // When precondition is true (default) this runs the full forward model and + // leaves decomposed_f holding the preconditioned search direction, exactly as + // the native solver uses it. When false, the forward model returns at the + // INVARIANT_RESIDUALS checkpoint (vmec.cc line ~836), so decomposed_f holds + // the raw, unpreconditioned force: the gradient of VMEC's augmented + // Lagrangian with respect to the (decomposed) state, including the + // 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. + // 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 = + precondition ? vmecpp::VmecCheckpoint::NONE + : vmecpp::VmecCheckpoint::INVARIANT_RESIDUALS; + const int checkpoint_after = precondition ? INT_MAX : 0; // Clear the restart reason before evaluating the forward model, exactly as // Vmec::Evolve does at its start (vmec.cc): the forward model only *sets* a // reason (BAD_JACOBIAN when the Jacobian flips, HUGE_INITIAL_FORCES at @@ -240,8 +258,8 @@ class VmecModel { *vmec_->decomposed_x_[0], *vmec_->physical_x_[0], *vmec_->decomposed_f_[0], *vmec_->physical_f_[0], need_restart, last_preconditioner_update_, last_full_update_nestor_, vmec_->fc_, - iter1, iter2, vmecpp::VmecCheckpoint::NONE, INT_MAX, - /*verbose=*/false); + iter1, iter2, checkpoint, checkpoint_after, + /*verbose=*/false, always_fix_m1_gauge); if (!s.ok()) { error_message = std::string(s.status().message()); } @@ -253,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, @@ -381,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; } @@ -453,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 @@ -1189,7 +1263,9 @@ PYBIND11_MODULE(_vmecpp, m) { py::class_(m, "VmecModel") .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")) + .def("evaluate", &VmecModel::Evaluate, py::arg("iter1"), py::arg("iter2"), + 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"), @@ -1205,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/vmec/vmec/vmec_in_memory_mgrid_test.cc b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec_in_memory_mgrid_test.cc index 7bcd608de..4c3b5937c 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec_in_memory_mgrid_test.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec_in_memory_mgrid_test.cc @@ -76,9 +76,12 @@ TEST(TestVmec, CheckInMemoryMgrid) { vmecpp::run(indata, magnetic_response_table); ASSERT_TRUE(output_with_inmemory_mgrid.ok()); - // compare wout contents + // compare wout contents. jcuru/jcurv are curl(B) currents whose two solve + // paths diverge by ~1.03e-7 across optimized/vectorized builds; keep every + // other quantity at 1e-7 and compare those two at 2e-7. vmecpp::CompareWOut(output_with_inmemory_mgrid->wout, original_output->wout, - /*tolerance=*/1e-7); + /*tolerance=*/1e-7, /*check_equal_niter=*/true, + /*current_density_tolerance=*/2e-7); } // Axisymmetric (ntor = 0, nzeta = 1) free-boundary tokamak (solovev_free_bdy). @@ -121,6 +124,10 @@ TEST(TestVmec, SolovevFreeBoundaryAxisymmetric) { ASSERT_TRUE(inmemory_output.ok()); // The in-memory makegrid path must reproduce the committed-mgrid run. + // jcuru/jcurv are curl(B) currents whose two solve paths diverge by ~1.03e-7 + // across optimized/vectorized builds; keep every other quantity at 1e-7 and + // compare those two at 2e-7. vmecpp::CompareWOut(inmemory_output->wout, disk_output->wout, - /*tolerance=*/1e-7); + /*tolerance=*/1e-7, /*check_equal_niter=*/true, + /*current_density_tolerance=*/2e-7); } 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 new file mode 100644 index 000000000..5e20b7fa0 --- /dev/null +++ b/tests/test_external_optimizers.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""External optimizers reach force balance on axisymmetric and 3D cases. + +The Solov'ev equilibrium has a reproducible internal state, so the 2D test compares the +full state vector with the native solver. In 3D, the poloidal parameterization is not +unique and spectral condensation only regularizes that coordinate freedom. The 3D test +therefore compares the residual and energy instead of coordinate-dependent Fourier +coefficients. +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from external_optimizers import ( # type: ignore + make_model, + reference_equilibrium, + residual, + solve_newton_hvp, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_preconditioned_descent, +) + +ROOT = Path(__file__).resolve().parents[1] +SOLOVEV = ROOT / "examples" / "data" / "solovev.json" +CTH_LIKE = ROOT / "examples" / "data" / "cth_like_fixed_bdy.json" +CMA = ROOT / "src" / "vmecpp" / "cpp" / "vmecpp" / "test_data" / "cma.json" + +SOLVERS = ( + solve_preconditioned_descent, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + solve_newton_hvp, +) + + +def _solve_all(input_path): + return {solver.__name__: solver(input_path, ns=11) for solver in SOLVERS} + + +@pytest.fixture(scope="module") +def reference_2d(): + return reference_equilibrium(SOLOVEV, ns=11) + + +@pytest.fixture(scope="module") +def reference_3d(): + return reference_equilibrium(CTH_LIKE, ns=11) + + +@pytest.fixture(scope="module") +def solutions_2d(): + return _solve_all(SOLOVEV) + + +@pytest.fixture(scope="module") +def solutions_3d(): + return _solve_all(CTH_LIKE) + + +def test_optimizers_reach_2d_equilibrium(reference_2d, solutions_2d): + x_star, w_star = reference_2d + for name, (x, result) in solutions_2d.items(): + assert result.residual_norm < 1e-7, name + assert abs(result.energy - w_star) < 1e-8, name + assert np.linalg.norm(x - x_star) < 1e-5, name + + +def test_optimizers_reach_3d_force_balance(reference_3d, solutions_3d): + _, w_star = reference_3d + for name, (_, result) in solutions_3d.items(): + assert result.residual_norm < 1e-7, name + assert np.isclose(result.energy, w_star, rtol=1e-4, atol=0.0), name + + +def test_preconditioner_reduces_newton_krylov_force_evaluations(solutions_3d): + plain = solutions_3d[solve_newton_krylov.__name__][1] + preconditioned = solutions_3d[solve_newton_krylov_preconditioned.__name__][1] + assert preconditioned.force_evals < plain.force_evals + + +def test_hvp_newton_uses_fewer_outer_iterations_than_descent(solutions_3d): + newton = solutions_3d[solve_newton_hvp.__name__][1] + descent = solutions_3d[solve_preconditioned_descent.__name__][1] + assert newton.outer_iters < 20 + assert newton.outer_iters < descent.outer_iters + + +def test_cma_cold_start_exercises_non_axisymmetric_paths(): + model = make_model(CMA, ns=25) + x0 = np.asarray(model.get_state(), float) + f0 = residual(model)(x0) + assert np.all(np.isfinite(f0)) + assert np.linalg.norm(f0) > 0.0 + + rng = np.random.default_rng(0) + v = rng.standard_normal(x0.size) + hv = np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float) + assert np.all(np.isfinite(hv)) + assert np.linalg.norm(hv) > 0.0 diff --git a/tests/test_hessian.py b/tests/test_hessian.py new file mode 100644 index 000000000..8034f5ace --- /dev/null +++ b/tests/test_hessian.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VmecModel.hessian_vector_product gives the augmented functional's curvature. + +The Hessian-vector product is a central directional derivative of the analytic +force (the gradient of VMEC's augmented functional), computed inside VMEC++: +H v = (F(x + eps v) - F(x - eps v)) / (2 eps). It must be linear in v and agree +with an independent finite difference of the force, and it restores the state. +""" + +from pathlib import Path + +import numpy as np + +from vmecpp.cpp import _vmecpp # type: ignore + +SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" + + +def _model(ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(SOLOVEV)), ns) + + +def _raw_force(model, x): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), float) + + +def test_hvp_matches_finite_difference(): + m = _model() + m.evaluate(2, 2, False) + x = np.asarray(m.get_state(), float) + rng = np.random.default_rng(0) + v = rng.standard_normal(x.size) + v /= np.linalg.norm(v) + + hv = np.asarray(m.hessian_vector_product(np.ascontiguousarray(v)), float) + + eps = 1e-6 + fd = (_raw_force(m, x + eps * v) - _raw_force(m, x - eps * v)) / (2 * eps) + assert np.linalg.norm(hv - fd) < 1e-5 * np.linalg.norm(fd) + + +def test_hvp_is_linear(): + m = _model() + m.evaluate(2, 2, False) + rng = np.random.default_rng(1) + v = rng.standard_normal(np.asarray(m.get_state()).size) + hv = np.asarray(m.hessian_vector_product(np.ascontiguousarray(v)), float) + hv2 = np.asarray(m.hessian_vector_product(np.ascontiguousarray(2.0 * v)), float) + assert np.linalg.norm(hv2 - 2.0 * hv) < 1e-9 * np.linalg.norm(hv) + + +def test_hvp_restores_state(): + m = _model() + m.evaluate(2, 2, False) + x0 = np.asarray(m.get_state(), float).copy() + rng = np.random.default_rng(2) + v = rng.standard_normal(x0.size) + m.hessian_vector_product(np.ascontiguousarray(v)) + assert np.allclose(np.asarray(m.get_state(), float), x0) diff --git a/tests/test_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 new file mode 100644 index 000000000..16d26825e --- /dev/null +++ b/tests/test_internal_gradient.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Tests for the unpreconditioned internal-basis gradient. + +VmecModel.evaluate(precondition=False) returns at the INVARIANT_RESIDUALS checkpoint, so +get_forces() yields the raw, unpreconditioned force: the gradient of VMEC's augmented +functional (MHD energy plus the spectral-condensation and lambda constraints) with +respect to the decomposed (internal-basis) state. This is the gradient an external +optimizer working in the internal basis needs. + +The preconditioned force (precondition=True) is the native solver's search direction and +points in a different direction; both vanish at the converged equilibrium. +""" + +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): + indata = _vmecpp.VmecINDATA.from_file(str(SOLOVEV)) + return _vmecpp.VmecModel.create(indata, ns) + + +def test_raw_force_differs_from_preconditioned(): + m = _model() + m.evaluate(2, 2, True) + f_prec = np.asarray(m.get_forces(), float) + m.evaluate(2, 2, False) + f_raw = np.asarray(m.get_forces(), float) + + assert np.all(np.isfinite(f_raw)) + assert np.linalg.norm(f_raw) > 0.0 + # The preconditioner is a non-trivial metric: the two vectors are different + # in direction, not just scale. + cos = np.dot(f_prec, f_raw) / (np.linalg.norm(f_prec) * np.linalg.norm(f_raw)) + assert abs(cos) < 0.99 + + +def test_raw_force_vanishes_at_equilibrium(): + m = _model() + m.evaluate(2, 2, False) + f_initial = np.linalg.norm(np.asarray(m.get_forces(), float)) + + m.solve() + m.evaluate(2, 2, False) + f_converged = np.linalg.norm(np.asarray(m.get_forces(), float)) + + # The augmented-functional gradient is the equilibrium residual: it drops by + # many orders of magnitude once the native solver has converged. + assert f_converged < 1e-6 * f_initial + + +def test_cold_start_is_excluded(): + # evaluate(1, 2) is the cold-start special case (forces initialised to 1.0); + # the raw-gradient path uses iter1 >= 2, where the force is well defined. + 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)