diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml new file mode 100644 index 0000000..8091dda --- /dev/null +++ b/.github/workflows/benchmarks.yml @@ -0,0 +1,30 @@ +name: benchmarks + +on: + workflow_dispatch: + schedule: + - cron: "29 8 * * 0" + +jobs: + verify: + runs-on: ubuntu-latest + timeout-minutes: 30 + env: + JAX_ENABLE_X64: "true" + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install -e '.[dev]' + - run: pytest -m numerics + - run: python benchmarks/benchmark_core.py + env: + PYQSC_BENCHMARK_OUTPUT: benchmarks/results/ci.json + - name: Record environment + run: python -m pip freeze + - uses: actions/upload-artifact@v4 + with: + name: benchmark-report + path: benchmarks/results/ci.json diff --git a/.github/workflows/compatibility.yml b/.github/workflows/compatibility.yml new file mode 100644 index 0000000..7983651 --- /dev/null +++ b/.github/workflows/compatibility.yml @@ -0,0 +1,88 @@ +name: compatibility + +on: + pull_request: + workflow_dispatch: + schedule: + - cron: "43 6 * * 3" + +jobs: + frozen-contracts: + runs-on: ubuntu-latest + timeout-minutes: 15 + env: + JAX_ENABLE_X64: "true" + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install -e '.[dev]' + - run: pytest tests/compatibility tests/regression + + upstream-head-smoke: + if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 20 + continue-on-error: true + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install -e '.[dev]' + - run: python -m pip install 'git+https://github.com/landreman/pyQSC.git' + - run: >- + python -c "from qsc import Qsc; + from pyqsc_jax.near_axis import near_axis; + kwargs = dict(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3, etabar=-0.9, nphi=31); + assert abs(float(Qsc(**kwargs).iota) - float(near_axis(**kwargs).iota)) < 1e-10" + + essos-field-jet: + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + JAX_ENABLE_X64: "true" + MPLBACKEND: Agg + steps: + - uses: actions/checkout@v4 + - uses: actions/checkout@v4 + with: + repository: uwplasma/ESSOS + ref: refactor/pyqsc-external-field-jet + path: essos-integration + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install -e '.[dev,plot]' + - run: python -m pip install -e ./essos-integration + - run: python -m pip install -e . + - run: >- + pytest + essos-integration/tests/test_field_jet.py + essos-integration/tests/test_field_jet_examples.py + + vmex-implicit: + runs-on: ubuntu-latest + timeout-minutes: 20 + env: + JAX_ENABLE_X64: "true" + PYQSC_RUN_VMEX: "1" + steps: + - uses: actions/checkout@v4 + - uses: actions/checkout@v4 + with: + repository: uwplasma/vmex + ref: main + path: vmex-integration + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + - run: python -m pip install -e '.[dev]' + - run: python -m pip install -e ./vmex-integration + - run: python -m pip install -e . + - run: pytest tests/integration/test_vmex_live.py diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..5f82801 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,30 @@ +name: docs + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + schedule: + - cron: "17 7 * * 1" + +jobs: + build: + runs-on: ubuntu-latest + timeout-minutes: 15 + env: + JAX_ENABLE_X64: "true" + MPLBACKEND: Agg + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install -e '.[docs,plot]' + - run: python -m sphinx -W --keep-going -b html docs docs/_build/html + - run: python -m sphinx -W --keep-going -b doctest docs docs/_build/doctest + - run: python examples/01_first_order_qa.py + - name: Check external links + if: github.event_name == 'schedule' + run: python -m sphinx -W --keep-going -b linkcheck docs docs/_build/linkcheck diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..5a38d3f --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,87 @@ +name: publish + +on: + release: + types: [published] + workflow_dispatch: + +permissions: + contents: read + +jobs: + verify-and-build: + runs-on: ubuntu-latest + timeout-minutes: 45 + env: + JAX_ENABLE_X64: "true" + MPLBACKEND: Agg + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Require production release commit on main + if: github.event_name == 'release' + run: >- + git fetch origin main:refs/remotes/origin/main && + git merge-base --is-ancestor "$GITHUB_SHA" origin/main + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install -e '.[dev,docs,plot]' + - run: ruff check src tests examples benchmarks + - run: ruff format --check src tests examples benchmarks + - run: pytest --cov=pyqsc_jax --cov-branch + - run: python -m sphinx -W --keep-going -b html docs docs/_build/html + - run: python -m sphinx -W --keep-going -b doctest docs docs/_build/doctest + - run: python -m build + - run: python -m twine check dist/* + - run: python -m venv "${{ runner.temp }}/wheel-env" + - run: "${{ runner.temp }}/wheel-env/bin/python -m pip install dist/*.whl" + - run: >- + JAX_ENABLE_X64=true + ${{ runner.temp }}/wheel-env/bin/python -c + "import pyqsc_jax as q; s=q.solve_configuration('qa', nphi=31); + assert s.root_report.converged and s.linear_report.converged" + - run: python -m venv "${{ runner.temp }}/sdist-env" + - run: "${{ runner.temp }}/sdist-env/bin/python -m pip install dist/*.tar.gz" + - run: >- + JAX_ENABLE_X64=true + ${{ runner.temp }}/sdist-env/bin/python -c + "from pyqsc_jax.near_axis import near_axis; + s=near_axis(rc=[1.0,0.045],zs=[0.0,-0.045],nfp=3,etabar=-0.9); + assert abs(float(s.iota)-0.41830690943386617)<1e-10" + - uses: actions/upload-artifact@v4 + with: + name: distributions + path: dist/ + + publish-testpypi: + if: github.event_name == 'workflow_dispatch' + needs: verify-and-build + runs-on: ubuntu-latest + environment: testpypi + permissions: + id-token: write + steps: + - uses: actions/download-artifact@v4 + with: + name: distributions + path: dist/ + - uses: pypa/gh-action-pypi-publish@release/v1 + with: + repository-url: https://test.pypi.org/legacy/ + + publish-pypi: + if: github.event_name == 'release' + needs: verify-and-build + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write + steps: + - uses: actions/download-artifact@v4 + with: + name: distributions + path: dist/ + - uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..522afc4 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,81 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + lint: + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install ruff + - run: ruff check src tests examples benchmarks + - run: ruff format --check src tests examples benchmarks + + test: + permissions: + contents: read + id-token: write + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-latest + python: "3.12" + - os: ubuntu-latest + python: "3.14" + - os: macos-latest + python: "3.14" + runs-on: ${{ matrix.os }} + timeout-minutes: 30 + env: + JAX_ENABLE_X64: "true" + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + cache: pip + - run: python -m pip install -e '.[dev,plot]' + - run: python -c "import jax, pyqsc_jax; print(jax.__version__, pyqsc_jax.__version__, pyqsc_jax.__file__)" + - run: pytest --cov=pyqsc_jax --cov-branch --cov-report=term --cov-report=xml + - uses: codecov/codecov-action@v5 + if: matrix.os == 'ubuntu-latest' && matrix.python == '3.14' + with: + use_oidc: true + fail_ci_if_error: true + + package: + runs-on: ubuntu-latest + timeout-minutes: 15 + env: + JAX_ENABLE_X64: "true" + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + cache: pip + - run: python -m pip install build twine + - run: python -m build + - run: python -m twine check dist/* + - run: python -m pip install --force-reinstall dist/*.whl + - run: >- + python -c "from pyqsc_jax.near_axis import near_axis; + q = near_axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3, etabar=-0.9); + assert abs(float(q.iota) - 0.41830690943386617) < 1e-10" diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a34cc79 --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.coverage +coverage.xml +htmlcov/ +.audit/ +docs/_build/ +build/ +dist/ +*.egg-info/ +examples/output/ +benchmarks/results/ diff --git a/.readthedocs.yaml b/.readthedocs.yaml new file mode 100644 index 0000000..98bf07a --- /dev/null +++ b/.readthedocs.yaml @@ -0,0 +1,17 @@ +version: 2 + +build: + os: ubuntu-24.04 + tools: + python: "3.12" + +sphinx: + configuration: docs/conf.py + fail_on_warning: true + +python: + install: + - method: pip + path: . + extra_requirements: + - docs diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..529e616 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,136 @@ +# Changelog + +All notable changes to pyQSC_JAX will be documented here. + +The project follows Semantic Versioning once the new canonical API reaches its +first stable release. + +## Unreleased + +### Added + +- Optional, lazy pyQSC_JAX-to-VMEX bridge for differentiable fixed-boundary + radial \(\iota\), quasisymmetry, magnetic-well, energy, volume, and + aspect-ratio quantities in vacuum and finite beta, with a live current-VMEX + compatibility job. +- Full-torus 3D surface helpers and a four-stellarator README gallery covering + database QA ID 139524, database QH ID 3, the \(B_{20}\)-optimized case, and + the large-singular-radius database case 107579. +- `plasma_stellarator`, an angle-dependent finite-pressure, exactly + zero-current database case whose pressure-driven plasma-field norm varies + by 14.04% and whose formal radius remains inside the singular surface. +- Reproducible README figures for QA/QH topology branches, resolution + convergence, synchronized JIT/JVP/VMAP timings, and five-way \(B_{20}\) + optimizer screening followed by staged eight-mode refinement. +- Vectorized, JIT-compiled VMEC boundary export with all four coefficient + families, deterministic INDATA output, conversion diagnostics, + legacy-adapter support, a frozen VMEC 9.0 `wout` regression, and opt-in + vacuum and finite-pressure, zero-current local-VMEC reruns. +- Dense-grid `b20_optimized_qa` reference with a verified + \(1.59\times10^{-6}\) weighted residual and reproducible comparisons across + L-BFGS-B, least-squares, differential-evolution, and multistart + Levenberg--Marquardt searches. +- `plasma_dominant_channel` validation case, for which the plasma field is + 33.3% of the total field at the documented formal radius. +- README evidence, benchmark reports, publication figures, executable + examples, and validation documentation for B20, plasma-field, and VMEC + behavior. +- Audited compatibility and numerical baselines. +- Physics traceability and architecture decision records. +- Modern package, documentation, test, and CI scaffolding. +- Immutable general Fourier axes with normalized coefficient packing. +- JAX-native periodic differentiation, interpolation, and integration. +- Sampled Frenet geometry with explicit validity diagnostics and independent + symmetric/asymmetric regression references. +- Converged, damped periodic sigma solve with residual, iteration, + backtracking, stagnation, and Jacobian-conditioning reports. +- Implicit differentiation of the converged sigma equation through SOLVAX. +- Immutable first-order shape, field, field-gradient, elongation, and + `L_grad_B` results with JIT, VMAP, JVP, and VJP validation. +- Thin ESSOS compatibility adapter preserving legacy array orientations, + coordinate conversion, boundary generation, and plotting without depending + on ESSOS. +- Complete finite-pressure/current r2 coefficient system, including + `X20`/`Y20`, second harmonics, `beta_1s`, `G2`, `B20`, derivatives, + untwisted boundary data, and direct `B20` diagnostics. +- Implicit differentiation and residual/conditioning reports for the coupled + dense second-order solve. +- Independent four-equation r2 residual checks and upstream regression cases + spanning vacuum QA, finite pressure/current, and QH topology. +- Regular-coordinate total-field value, gradient, and Hessian on the axis, + including coordinate-map conditioning and Maxwell-identity diagnostics. +- Mercier magnetic-well and geodesic terms matching vacuum and finite-pressure + pyQSC references. +- JIT/JVP, resolution, upstream Hessian, full vacuum-symmetry, divergence, and + derivative-of-divergence validation for the total field jet. +- Singular-radius diagnostics derived directly from the regular-coordinate + Jacobian determinant, with vectorized global angular seeding, Newton + refinement, residual reports, and pyQSC QA/QH/finite-current parity. +- Differentiable r3 flux-constraint surface corrections, untwisted boundary + coefficients, and two independent consistency checks with upstream pyQSC + QA, finite-pressure/current, and QH parity. +- Standard-MHS magnetic shear with periodic and secular integrating-factor + branches, explicit `B31c`, immutable diagnostics, legacy adapter support, + and upstream QA/QH/asymmetric parity. +- Branch-local target-transform solves for sign-preserving `etabar` or `I2`, + including implicit derivatives, round-trip validation, local response and + fold diagnostics, and propagation through r2. +- Full-state fixed-sign pseudo-arclength continuation for `etabar`, including + fold-crossing detection, corrected solution records, and explicit partial + completion statuses. +- Dense nonconstant-`B20` diagnostics, exact affine elimination of `B2c`, + degenerate-response handling, and independent doubled/quadrupled-grid + verification. +- Scalable Curvo et al. (2025) Table 3 screening criteria with configurable + thresholds, strict/inclusive comparison semantics, units, measured values, + signed margins, aggregate status, and differentiable array results. +- Deterministic bounded axis exploration, damped JAX-Jacobian local + least-squares refinement, normalized basin clustering, hard criteria, + canonical secondary selectors, staged Fourier warm starts, independent + verification, and explicit global-certificate status semantics. +- Nonzero \(B_{20}\) Fourier-\(L^1\) certificates in dense and + resolution-verification diagnostics. +- Exact positive-volume plasma-current source through quadratic radial order, + including independent pressure/current pathways, regular evaluation through + \(I_2=0\), and covariant-current/enclosed-ampere conversions. +- Matched on-axis free-space plasma field with a full-torus periodic + finite-part integral, elliptical core constants, second-order local shape + correction, arbitrary matching-length cancellation, asymptotic error + metadata, and independent resolved-volume Biot–Savart validation. +- Local straight-ellipse plasma gradient with divergence and Ampère + diagnostics, external symmetric-trace-free subtraction, vacuum reduction, + and a reversible five-component Cartesian STF representation. +- Complete local plasma Hessian from the affine-current, curved-channel, and + second-order shape potentials, including Frenet-connection derivatives, + external vacuum subtraction, asymptotic metadata, and a reversible + seven-component Cartesian STF representation. +- Named immutable QA, QH, and finite-pressure/current configurations. +- Lazy plotting helpers that return Matplotlib figure and axes objects. +- Fourteen direct tutorial scripts and ten deterministic publication-figure + scripts with commit/parameter metadata and clean-directory CI execution. +- Downstream ESSOS external-field-jet integration with normalized 3+5+7 + objectives, finite-beta stage-two and single-stage examples, and + finite-difference gradient validation. +- Complete documentation hierarchy, migration guide, validation reports, and + release checklist. + +### Fixed + +- VMEC export now records the cylindrical-angle inversion tolerance and + convergence state, rejects nonfinite or nonpositive radii, and refuses to + write an invalid boundary when the inversion has not converged. The VMEX + bridge applies the same guard, while traced invalid candidates map to + nonfinite coefficients instead of silently solving a different surface. +- VMEX tutorial and publication scripts now require the explicit + `PYQSC_RUN_VMEX=1` opt-in used by live integration CI, so installing VMEX + does not unexpectedly turn the ordinary example matrix into a multi-minute + equilibrium run. +- ESSOS finite-beta demonstrations now use a nonplanar database stellarator + with finite pressure and exactly zero toroidal current; the single-stage + design vector can no longer vary \(I_2\). +- A no-op assignment to legacy `dofs` no longer exchanges normal and binormal + cylindrical components. +- The legacy sigma calculation no longer assumes exactly five Newton updates. +- Canonical field gradients now use explicit + `(sample, field component, derivative direction)` ordering; the ESSOS + adapter retains its historical layout. diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..d994ca8 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,19 @@ +cff-version: 1.2.0 +message: "If you use pyQSC_JAX, please cite this software and the theory papers relevant to your calculation." +title: "pyQSC_JAX" +type: software +version: 0.2.0-dev +date-released: 2026-07-29 +authors: + - family-names: Jorge + given-names: Rogério + affiliation: University of Wisconsin–Madison +repository-code: "https://github.com/uwplasma/pyQSC_JAX" +url: "https://github.com/uwplasma/pyQSC_JAX" +license: MIT +keywords: + - stellarator + - near-axis expansion + - quasisymmetry + - JAX + - automatic differentiation diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..5dbfc68 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,37 @@ +# Contributing + +pyQSC_JAX changes must preserve the ESSOS legacy contract and provide evidence +for every physics result. + +## Developer setup + +```bash +python -m venv .venv +. .venv/bin/activate +python -m pip install -e '.[dev,docs,plot]' +pytest +ruff check src tests +python -m sphinx -W -b html docs docs/_build/html +``` + +Do not pin dependency versions in package metadata or committed requirements +files. Release reports record the versions actually tested. + +## Physics changes + +Before merging a physics result: + +1. add its primary equation source and conventions to the traceability table; +2. add a focused unit test; +3. add an independent regression or physical-identity test; +4. document shapes, units, validity assumptions, and failure modes; +5. verify JIT and automatic differentiation where the API promises them. + +Reference data must be small, local, checksummed, license-compatible, and tied +to an upstream commit. Tests must never download data. + +## Pull requests + +Keep commits small and green. Explain the physical source, numerical method, +validation, compatibility impact, and any remaining limitation. Do not combine +unrelated formatting or generated output with physics changes. diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..f3341df --- /dev/null +++ b/NOTICE @@ -0,0 +1,37 @@ +pyQSC_JAX +Copyright (c) 2026 UW Plasma + +This project is distributed under the MIT License in LICENSE. + +The landreman/pyQSC project is used as an upstream reference implementation +and source of regression data: + + pyQSC + Copyright (c) 2020 Matt Landreman + SPDX-License-Identifier: BSD-2-Clause + https://github.com/landreman/pyQSC + +Any source or reference data adapted from pyQSC is identified in the relevant +file and retains the following notice: + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are met: + + 1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + + 2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md index d1585a3..7464446 100644 --- a/README.md +++ b/README.md @@ -1 +1,326 @@ -# pyQSC_JAX \ No newline at end of file +# pyQSC_JAX + +Differentiable near-axis stellarator construction and surface-free +plasma–coil field jets in JAX. + +[![Tests](https://github.com/uwplasma/pyQSC_JAX/actions/workflows/tests.yml/badge.svg)](https://github.com/uwplasma/pyQSC_JAX/actions/workflows/tests.yml) +[![Coverage](https://codecov.io/gh/uwplasma/pyQSC_JAX/branch/main/graph/badge.svg)](https://codecov.io/gh/uwplasma/pyQSC_JAX) +[![PyPI](https://img.shields.io/pypi/v/pyqsc-jax.svg)](https://pypi.org/project/pyqsc-jax/) +[![Docs](https://readthedocs.org/projects/pyqsc-jax/badge/?version=latest)](https://pyqsc-jax.readthedocs.io/) +[![DOI pending](https://img.shields.io/badge/DOI-pending-lightgrey.svg)](docs/release_checklist.md) +[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE) + +> **Status:** development branch for the 0.2 refactor. The local physics, +> compatibility, VMEC/VMEX, packaging, and external 3+5+7 field-jet gates are +> validated. Maintainer review, trusted publishing, and archival release +> metadata remain external release gates. + +![A visually diverse gallery of screened QA, QH, optimized, and large-clearance stellarators](docs/_static/stellarator_gallery.png) + +Four geometries with four distinct jobs: the one-period QA +[database ID 139524](https://stellarator.physics.wisc.edu/app/plot/139524), +the four-period QH [database ID 3](https://stellarator.physics.wisc.edu/app/plot/3), +an eight-mode flat-\(B_{20}\) refinement of +[database ID 57409](https://stellarator.physics.wisc.edu/app/plot/57409), and +the large-clearance [database ID 107579](https://stellarator.physics.wisc.edu/app/plot/107579). +Every panel is independently re-solved and screened. The lead QA stellarator +has helicity zero, \(\lvert\iota\rvert=0.355\), finite pressure, exactly +\(I_2=0\), and nonzero torsion; it passes the Curvo/Table-3 profile with the +requested \(\lvert\iota\rvert\geq0.3\) gate. The other three pass at +\(\lvert\iota\rvert\geq0.4\). Plot radii and provenance are saved beside the +runnable figure as JSON. + +| \(B_{20}\) optimization, geometry, and singular margin | Angle-dependent plasma/external field | +| --- | --- | +| ![B20 optimization](docs/_static/B20_optimization.png) | ![Plasma and external jet](docs/_static/plasma_external_jet.png) | + +| QA versus QH topology branches | Resolution audit catches under-resolved \(B_{20}\) | +| --- | --- | +| ![QA and QH topology branches](docs/_static/QA_QH_branches.png) | ![Spectral resolution convergence](docs/_static/convergence.png) | + +| Measured JIT, JVP, and VMAP speedups | Five optimizers plus staged refinement | +| --- | --- | +| ![Measured JAX core performance](docs/_static/core_performance.png) | ![B20 optimizer comparison](docs/_static/optimizer_comparison.png) | + +| Differentiable radial VMEX quantities | VMEC2000 export validation | +| --- | --- | +| ![VMEX radial quantities](docs/_static/vmex_radial_profiles.png) | ![VMEC validation](docs/_static/vmec_validation.png) | + +## Install + +```bash +python -m pip install pyqsc-jax +``` + +Install the JAX build appropriate for your CPU or accelerator. pyQSC_JAX does +not configure a backend or enable 64-bit mode at import time. + +For 3D plots and differentiable finite-radius equilibria: + +```bash +python -m pip install 'pyqsc-jax[plot,vmex]' +``` + +## Quickstart + +```python +import pyqsc_jax as qsc + +configuration = qsc.get_configuration("database_qa_139524") +solution = configuration.solve(nphi=81) +criteria = qsc.Criteria.from_curvo_2025(minimum_abs_iota=0.3) + +assert solution.root_report.converged +assert solution.linear_report.converged +assert criteria.evaluate(solution).passed +assert int(solution.helicity) == 0 +assert solution.inputs.I2 == 0 +print("source:", configuration.source_url) +print("|iota|:", abs(solution.iota)) +print("B20 residual:", solution.B20_residual) +print("Mercier D r^2:", solution.DMerc_times_r2) +print("singular radius:", solution.r_singularity) +print("minimum L_grad_grad_B:", solution.L_grad_grad_B.min()) +``` + +The result is an immutable JAX pytree. Canonical field arrays use +sample-first ordering: `(nphi, 3)`, `(nphi, 3, 3)`, and +`(nphi, 3, 3, 3)`. + +## Pick a workflow + +| Goal | Start here | Main result | +| --- | --- | --- | +| Construct QA/QH near the axis | `qsc.Qsc(...)` | immutable near-axis solution | +| Flatten \(B_{20}\) | `optimize_B2c`, `search_axis` | verified dense-grid diagnostics | +| Inspect a full 3D surface | `plot_surface_3d` | figure plus explicit plotting radius | +| Compute radial MHD quantities | `to_vmex_problem(...).solve()` | differentiable \(\iota(s)\), QS profile, magnetic well | +| Match finite-beta coils | `plasma_hessian_on_axis` | external vacuum 3+5+7 target | +| Run conventional VMEC2000 | `to_vmec` | diagnosed `input.*` file | + +## Capabilities + +| Capability | Status | +| --- | --- | +| General Fourier axes, QA/QH topology, spectral Frenet geometry | Validated | +| Converged periodic sigma solve with implicit JVP/VJP | Validated | +| Complete finite-pressure/current r2 and total field Hessian | Validated | +| r3 boundary correction and magnetic shear | Validated documented specialization | +| Mercier, singular radius, scale lengths, \(B_{20}\) spectrum | Validated | +| Target-\(\iota\) inverse solve and pseudo-arclength folds | Validated | +| Exact \(B_{2c}\), criteria, bounded multistart axis search | Validated | +| Surface-free plasma field, gradient, Hessian | Validated asymptotic model | +| External vacuum 3+5+7 field-jet target | Validated | +| Fast fixed-boundary VMEC export, including asymmetric boundaries | Validated against VMEC 9.0 | +| Differentiable VMEX radial iota, QS profile, magnetic well | Validated in vacuum and finite beta | +| ESSOS vacuum/finite-beta stage-two and single-stage objectives | Draft integration PR | +| Legacy `pyqsc_jax.near_axis.near_axis` contract | Preserved | + +## Headline regression cases + +These are checked results, not illustrative targets. The detailed optimizer, +resolution, timing, and radius-convergence tables are in the +[B20](docs/theory/b20-optimization.md), +[plasma-field](docs/validation/plasma_field.md), and +[VMEC](docs/validation/vmec.md) validation pages. + +| Check | Result | Regression gate | +| --- | ---: | ---: | +| Database-seeded QH \(B_{20}\), weighted \(L^2\), `nphi=121` | \(1.2743\times10^{-10}\) | \(<1.4\times10^{-10}\) | +| Improvement over ID-57409 exact-\(B_{2c}\) solve | \(2.4321\times10^8\)× | \(>2.0\times10^8\)× | +| \(B_{20}\)-optimized singular radius | 0.249 m | displayed \(r=0.075\) m gives 3.3× clearance | +| Largest showcased singular radius, database ID 107579 | 0.432 m | \(>0.4\) m | +| Showcase QA, database ID 139524 | helicity 0, \(\lvert\iota\rvert=0.355\), \(\tau_\mathrm{RMS}=1.199\ \mathrm{m}^{-1}\) | Curvo profile with \(\lvert\iota\rvert\ge0.3\) | +| Four-case showcase design screen | all pass | QA gate 0.3; other cases \(\lvert\iota\rvert\ge0.4\) | +| Pressure-only \(I_2=0\) plasma fraction, min/mean/max at \(a=0.15\) m | 0.1761% / 0.1890% / 0.2026% | nonzero and resolved | +| Pressure-driven plasma-field norm peak-to-peak / mean | 14.04% | \(>10\%\) | +| Finite-beta showcase axis torsion, RMS | \(0.979\ \mathrm{m}^{-1}\) | \(>0.5\ \mathrm{m}^{-1}\) | +| VMEC/near-axis on-axis \(\iota\), \(r=0.0025\) | 0.418543 / 0.418307 | relative error \(<0.1\%\) | +| VMEC force residuals | \(2.4\)–\(7.6\times10^{-11}\) | maximum \(<10^{-9}\) | +| Finite-pressure, \(I_2=0\) database QA VMEC/near-axis \(\iota\), \(r=0.0015\) | \(-0.354726\) / \(-0.354802\) | relative error \(0.0214\%\) | +| Finite-pressure database QA VMEC force residuals | \(2.2\)–\(10.0\times10^{-12}\) | maximum \(<10^{-9}\) | +| `to_vmec`, Apple M4 CPU, compile+execute / warm | 0.476 s / 5.20 ms | warm unit gate \(<0.5\) s | +| First-order solve, `nphi=61`, Apple CPU, compile+execute / warm | 198.2 ms / 0.144 ms | synchronized median; 1,373× amortized speedup | + +```python +optimized = qsc.solve_configuration("b20_optimized_good", nphi=121) +diagnostics = qsc.b20_diagnostics(optimized) +print(diagnostics.weighted_l2, diagnostics.grid_maximum) + +plasma_case = qsc.solve_configuration("plasma_stellarator", nphi=61) +assert plasma_case.inputs.I2 == 0 +plasma = qsc.plasma_field_on_axis(plasma_case, formal_radius=0.15) +``` + +## Target rotational transform + +```python +axis = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) +inverse = qsc.solve( + axis=axis, + iota=0.42, + etabar=-1.0, + solve_for="etabar", +) +print(inverse.inputs.etabar, inverse.response_derivative, inverse.branch_fold) +``` + +The `etabar` sign and seed select a local branch. Use +`continue_etabar_branch` when the response approaches a fold. + +## Axis optimization + +```python +indices = qsc.stellarator_symmetric_variable_indices(axis, modes=(1,)) +problem = qsc.AxisSearchProblem( + axis=axis, + variable_indices=indices, + lower_bounds=[0.02, -0.08], + upper_bounds=[0.08, -0.02], + etabar=-0.9, +) +result = qsc.search_axis(problem) +print(result.status, result.search_budget, result.distinct_basins) +``` + +Only `verified_zero` certifies the known zero lower bound of the nonnegative +primary \(B_{20}\) residual. A nonzero result is `best_found`, never a +mathematical global-minimum claim. + +## The \(B_{20}\)-optimized stellarator in 3D + +```python +from pyqsc_jax.plotting import plot_surface_3d + +optimized = qsc.solve_configuration("b20_optimized_good", nphi=121) +figure, axis = plot_surface_3d(optimized, radius=0.075) + +print(abs(optimized.iota)) # 2.9636 +print(optimized.B20_residual) # 1.2743e-10 T/m^2 +print(optimized.r_singularity) # 0.249 m +``` + +This configuration is a staged eight-mode refinement of public database ID +57409, not an unconstrained toy optimum. It passes every named design +criterion with \(\lvert\iota\rvert=2.964\); the displayed surface is 3.3 times +inside the computed singular radius. The optimizer comparison, +Fourier-resolution audit, full coefficient set, provenance, and the +distinction between `best_found` and a certified zero are in the +[B20 optimization notes](docs/theory/b20-optimization.md). + +## Finite-beta coil targets + +```python +solution = qsc.solve_configuration("plasma_stellarator", nphi=61) +assert solution.inputs.I2 == 0 +assert solution.inputs.p2 != 0 +target = qsc.plasma_hessian_on_axis(solution, formal_radius=0.15) + +print(target.field.external_field.shape) # (nphi, 3) +print(target.field.external_gradient_independent.shape) # (nphi, 5) +print(target.external_hessian_independent.shape) # (nphi, 7) +``` + +ESSOS consumes these external vacuum targets for normalized field, gradient, +and Hessian coil objectives. The dependency remains one-way: +`ESSOS -> pyQSC_JAX`. + +This is public stellarator-database configuration 52521, with exactly +\(I_2=0\), finite \(p_2=-2.8248\times10^4\ \mathrm{Pa/m^2}\), +\(\lvert\iota\rvert=2.809\), \(r_\mathrm{sing}=0.392\) m, and RMS axis +torsion \(0.979\ \mathrm{m}^{-1}\). It independently passes the Curvo profile. +The stated formal radius \(a=0.15\) m is mandatory model metadata and remains +inside the singular-radius estimate. The public plot uses Frenet components +and a percent scale, making the pressure-driven angle dependence and exact +cancellation of plasma and external transverse components visible; +\(|B_\mathrm{total}|=B_0\) is constant on axis by construction. + +The pressure-only contribution is approximately \(0.19\%\), not 30%. The +larger fraction previously shown came from finite \(I_2\) and produced an +essentially tokamak-like example; it is intentionally excluded from the +showcase. Current-driven cases remain supported and validated by the library, +but all user-facing finite-beta examples use \(I_2=0\). + +## Differentiable radial equilibria with VMEX + +```python +import jax + +near_axis = qsc.solve_configuration("plasma_stellarator", nphi=61) +assert near_axis.inputs.I2 == 0 +problem = qsc.to_vmex_problem( + near_axis, + r=0.02, + qs_surfaces=(0.25, 0.5, 0.75, 1.0), + adjoint_tol=1e-8, +) +equilibrium = problem.solve() + +print(equilibrium.quantities.iota) # full radial profile +print(equilibrium.quantities.quasisymmetry) # requested flux surfaces +print(equilibrium.quantities.magnetic_well) + +well, gradient = jax.value_and_grad( + lambda parameters: qsc.vmex_radial_quantities( + problem, parameters + ).magnetic_well +)(problem.parameters) +print(gradient.rbc, gradient.pres_scale) +``` + +`problem.parameters_for(new_near_axis_solution)` traceably rebuilds the +boundary, flux, pressure, and current leaves, allowing derivatives to +propagate from pyQSC_JAX variables through VMEX's converged fixed point. +Vacuum and scalar finite-beta fixed-boundary paths are tested against +[uwplasma/VMEX](https://github.com/uwplasma/vmex); the exact validated commit +and current limitations are recorded in the +[VMEX integration validation](docs/validation/vmex_interface.md). + +## VMEC export + +```python +solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=61, + order="r2", +) +export = qsc.to_vmec(solution, "input.qa", r=0.005) +print(export.conversion_seconds) +print(export.boundary.toroidal_angle_converged) +print(export.boundary.maximum_R_reconstruction_error) +``` + +This small-radius QA input is a conversion-validation reference, not one of +the screened finite-beta showcase devices above. The exporter uses a +vectorized Newton inversion and a two-dimensional FFT, +supports all four VMEC boundary coefficient families, and returns conversion +diagnostics. An unconverged toroidal-angle inversion raises before any file is +written. The committed `wout` regression checks the resulting equilibrium, +including the on-axis transform; opt-in tests rerun local VMEC for both the +vacuum reference and the finite-pressure, zero-current database QA case. + +## Learn and validate + +- [Installation and model selection](docs/getting_started/) +- [Theory and conventions](docs/theory/) +- [Executable tutorials](docs/tutorials/) +- [API reference](docs/api/) +- [pyQSC, literature, plasma, and ESSOS validation](docs/validation/) +- [VMEC conversion, equilibrium, and performance validation](docs/validation/vmec.md) +- [Differentiable VMEX radial-equilibrium validation](docs/validation/vmex_interface.md) +- [Migration from pyQSC and the legacy adapter](docs/migration.md) +- [Known limitations](docs/advanced/limitations.md) + +The [draft pyQSC_JAX PR](https://github.com/uwplasma/pyQSC_JAX/pull/2) and +[draft ESSOS integration PR](https://github.com/uwplasma/ESSOS/pull/46) +record the active review state. + +## Citation and license + +Citation metadata is in [CITATION.cff](CITATION.cff). A Zenodo DOI will be +minted as part of the first reviewed release; the badge remains explicitly +pending until then. pyQSC_JAX is MIT licensed. Adapted pyQSC source and +reference data retain BSD-2-Clause attribution in [NOTICE](NOTICE). diff --git a/benchmarks/benchmark_b20_optimizers.py b/benchmarks/benchmark_b20_optimizers.py new file mode 100644 index 0000000..51bc2eb --- /dev/null +++ b/benchmarks/benchmark_b20_optimizers.py @@ -0,0 +1,223 @@ +"""Compare optimizers from screened stellarator-database configuration 57409.""" + +from __future__ import annotations + +import json +import os +import platform +import subprocess +from pathlib import Path +from time import perf_counter + +import jax +import jax.numpy as jnp +import numpy as np +from scipy.optimize import differential_evolution, least_squares, minimize + +import pyqsc_jax as qsc + +NPHI = 121 +MODES = (1, 2, 3) +OUTPUT = Path( + os.environ.get("PYQSC_B20_BENCHMARK_OUTPUT", "benchmarks/results/b20_optimizers.json") +) +RC = ( + 1.0, + -0.51677144, + -0.009499784, + -0.005914526, +) +ZS = ( + 0.0, + -0.5420635, + -0.012225689, + -0.0059485724, +) +AXIS = qsc.Axis(rc=RC, zs=ZS, nfp=4) +VARIABLE_INDICES = qsc.stellarator_symmetric_variable_indices(AXIS, modes=MODES) +INITIAL = AXIS.dofs[jnp.asarray(VARIABLE_INDICES)] +HALF_WIDTH = jnp.asarray((0.14, 0.07, 0.03, 0.14, 0.07, 0.03)) +LOWER = INITIAL - HALF_WIDTH +UPPER = INITIAL + HALF_WIDTH +ETABAR = -1.3295174 +P2 = -23501.281 +SOURCE_DATABASE_ID = 57409 + + +def projected_residual(variables): + axis = AXIS.with_dofs(AXIS.dofs.at[jnp.asarray(VARIABLE_INDICES)].set(variables)) + solution = qsc.solve( + axis=axis, + etabar=ETABAR, + B0=1.0, + B2c=0.0, + p2=P2, + nphi=NPHI, + order="r2", + ) + optimized = qsc.optimize_B2c(solution) + weights = optimized.solution.geometry.d_l_d_phi + return ( + jnp.sqrt(weights / jnp.sum(weights)) + * optimized.diagnostics.anomaly + / optimized.solution.inputs.B0 + ) + + +compiled_residual = jax.jit(projected_residual) +compiled_jacobian = jax.jit(jax.jacrev(projected_residual)) +compiled_value_and_gradient = jax.jit( + jax.value_and_grad(lambda value: 0.5 * jnp.sum(projected_residual(value) ** 2)) +) + + +def residual_numpy(variables): + return np.asarray(compiled_residual(jnp.asarray(variables)), dtype=float) + + +def jacobian_numpy(variables): + return np.asarray(compiled_jacobian(jnp.asarray(variables)), dtype=float) + + +def value_and_gradient_numpy(variables): + value, gradient = compiled_value_and_gradient(jnp.asarray(variables)) + return float(value), np.asarray(gradient, dtype=float) + + +def weighted_l2(variables): + return float(np.linalg.norm(residual_numpy(variables))) + + +print("Compiling shared B20 residual and derivatives...") +compiled_residual(INITIAL).block_until_ready() +compiled_jacobian(INITIAL).block_until_ready() +compiled_value_and_gradient(INITIAL)[0].block_until_ready() +rows = [ + { + "method": "exact B2c only", + "seconds": 0.0, + "function_evaluations": 1, + "weighted_l2": weighted_l2(INITIAL), + } +] + +start = perf_counter() +lbfgsb = minimize( + value_and_gradient_numpy, + np.asarray(INITIAL), + method="L-BFGS-B", + jac=True, + bounds=list(zip(np.asarray(LOWER), np.asarray(UPPER), strict=True)), + options={"maxiter": 60, "ftol": 1.0e-18, "gtol": 1.0e-12}, +) +rows.append( + { + "method": "SciPy L-BFGS-B", + "seconds": perf_counter() - start, + "function_evaluations": int(lbfgsb.nfev), + "weighted_l2": weighted_l2(lbfgsb.x), + } +) + +start = perf_counter() +least_squares_result = least_squares( + residual_numpy, + np.asarray(INITIAL), + jac=jacobian_numpy, + bounds=(np.asarray(LOWER), np.asarray(UPPER)), + max_nfev=60, + xtol=1.0e-13, + ftol=1.0e-13, + gtol=1.0e-13, + x_scale="jac", +) +rows.append( + { + "method": "SciPy least_squares", + "seconds": perf_counter() - start, + "function_evaluations": int(least_squares_result.nfev), + "weighted_l2": weighted_l2(least_squares_result.x), + } +) + +start = perf_counter() +differential = differential_evolution( + lambda variables: weighted_l2(variables) ** 2, + list(zip(np.asarray(LOWER), np.asarray(UPPER), strict=True)), + seed=7, + popsize=4, + maxiter=4, + polish=False, + updating="immediate", +) +rows.append( + { + "method": "SciPy differential_evolution, low budget", + "seconds": perf_counter() - start, + "function_evaluations": int(differential.nfev), + "weighted_l2": weighted_l2(differential.x), + } +) + +problem = qsc.AxisSearchProblem( + axis=AXIS, + variable_indices=VARIABLE_INDICES, + lower_bounds=LOWER, + upper_bounds=UPPER, + etabar=ETABAR, + B0=1.0, + p2=P2, + nphi=NPHI, + criteria=qsc.Criteria.from_curvo_2025(minimum_abs_iota=0.4), +) +options = qsc.AxisSearchOptions( + coarse_samples=12, + local_starts=3, + maximum_iterations=20, + verification_multipliers=(1,), + verification_tail_tolerance=1.0, +) +start = perf_counter() +multistart = qsc.search_axis(problem, options=options) +rows.append( + { + "method": "pyQSC_JAX multistart Levenberg-Marquardt", + "seconds": perf_counter() - start, + "function_evaluations": multistart.search_budget, + "weighted_l2": None if multistart.best is None else multistart.best.primary_residual, + "status": multistart.status, + "distinct_basins": multistart.distinct_basins, + } +) + +repository = Path(__file__).resolve().parents[1] +try: + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +report = { + "git_commit": commit, + "python": platform.python_version(), + "platform": platform.platform(), + "jax": jax.__version__, + "jax_backend": jax.default_backend(), + "jax_enable_x64": bool(jax.config.jax_enable_x64), + "nphi": NPHI, + "modes": MODES, + "source_database_id": SOURCE_DATABASE_ID, + "source_url": (f"https://stellarator.physics.wisc.edu/app/plot/{SOURCE_DATABASE_ID}"), + "selected_configuration": "b20_optimized_good", + "timing_note": ( + "SciPy timings exclude shared residual/Jacobian compilation; the " + "pyQSC_JAX multistart timing includes compilation of its independent closure." + ), + "results": rows, +} +OUTPUT.parent.mkdir(parents=True, exist_ok=True) +OUTPUT.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") +for row in rows: + print(row) +print("saved:", OUTPUT) diff --git a/benchmarks/benchmark_core.py b/benchmarks/benchmark_core.py new file mode 100644 index 0000000..6af4345 --- /dev/null +++ b/benchmarks/benchmark_core.py @@ -0,0 +1,116 @@ +"""Record synchronized first-order compile, execution, JVP, and VMAP timings.""" + +from __future__ import annotations + +import json +import os +import platform +import statistics +import subprocess +import time +from pathlib import Path + +import jax +import jax.numpy as jnp + +import pyqsc_jax as qsc + +AXIS = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) +ETABAR = -0.9 +RESOLUTIONS = (15, 31, 61) +WARM_REPETITIONS = 7 +BATCH_SIZE = 8 +OUTPUT = Path(os.environ.get("PYQSC_BENCHMARK_OUTPUT", "benchmarks/results/local.json")) + + +def elapsed_seconds(callable_, *arguments): + """Synchronize a scalar JAX result and return value plus wall time.""" + + start = time.perf_counter() + value = callable_(*arguments) + jax.tree.map(lambda leaf: leaf.block_until_ready(), value) + return value, time.perf_counter() - start + + +def benchmark_resolution(nphi): + """Measure one static toroidal resolution.""" + + def transform(etabar): + return qsc.solve(axis=AXIS, etabar=etabar, nphi=nphi).iota + + compiled = jax.jit(transform) + jax.clear_caches() + value, cold_seconds = elapsed_seconds(compiled, jnp.asarray(ETABAR)) + warm_seconds = [ + elapsed_seconds(compiled, jnp.asarray(ETABAR))[1] for _ in range(WARM_REPETITIONS) + ] + + differentiated = jax.jit(lambda etabar: jax.jvp(transform, (etabar,), (jnp.ones_like(etabar),))) + _, jvp_cold_seconds = elapsed_seconds(differentiated, jnp.asarray(ETABAR)) + jvp_warm_seconds = [ + elapsed_seconds(differentiated, jnp.asarray(ETABAR))[1] for _ in range(WARM_REPETITIONS) + ] + + batched = jax.jit(jax.vmap(transform)) + batch_parameters = jnp.linspace(-1.0, -0.8, BATCH_SIZE) + _, vmap_cold_seconds = elapsed_seconds(batched, batch_parameters) + vmap_warm_seconds = [ + elapsed_seconds(batched, batch_parameters)[1] for _ in range(WARM_REPETITIONS) + ] + return { + "nphi": nphi, + "iota": float(value), + "cold_compile_and_execute_seconds": cold_seconds, + "warm_median_seconds": statistics.median(warm_seconds), + "warm_samples_seconds": warm_seconds, + "jvp_cold_compile_and_execute_seconds": jvp_cold_seconds, + "jvp_warm_median_seconds": statistics.median(jvp_warm_seconds), + "jvp_warm_samples_seconds": jvp_warm_seconds, + "vmap_batch_size": BATCH_SIZE, + "vmap_cold_compile_and_execute_seconds": vmap_cold_seconds, + "vmap_warm_median_seconds": statistics.median(vmap_warm_seconds), + "vmap_warm_samples_seconds": vmap_warm_seconds, + } + + +results = [] +print("Benchmarking first-order solve kernels...") +for nphi in RESOLUTIONS: + row = benchmark_resolution(nphi) + results.append(row) + print( + "nphi=", + nphi, + "cold=", + f"{row['cold_compile_and_execute_seconds']:.3f}s", + "warm=", + f"{row['warm_median_seconds']:.6f}s", + "JVP warm=", + f"{row['jvp_warm_median_seconds']:.6f}s", + "VMAP warm=", + f"{row['vmap_warm_median_seconds']:.6f}s", + ) + +repository = Path(__file__).resolve().parents[1] +try: + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +report = { + "git_commit": commit, + "python": platform.python_version(), + "platform": platform.platform(), + "processor": platform.processor(), + "jax": jax.__version__, + "jax_backend": jax.default_backend(), + "jax_devices": [str(device) for device in jax.devices()], + "jax_enable_x64": bool(jax.config.jax_enable_x64), + "warm_repetitions": WARM_REPETITIONS, + "results": results, +} +OUTPUT.parent.mkdir(parents=True, exist_ok=True) +OUTPUT.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") +print("saved:", OUTPUT) diff --git a/benchmarks/benchmark_vmec_export.py b/benchmarks/benchmark_vmec_export.py new file mode 100644 index 0000000..547a354 --- /dev/null +++ b/benchmarks/benchmark_vmec_export.py @@ -0,0 +1,90 @@ +"""Record synchronized cold and warm VMEC boundary-export timings.""" + +from __future__ import annotations + +import json +import os +import platform +import statistics +import subprocess +import tempfile +from pathlib import Path + +import jax + +import pyqsc_jax as qsc + +RADIUS = 0.03 +NPHI = 61 +NTHETA = 40 +MPOL = 12 +NTOR = 14 +WARM_REPETITIONS = 7 +OUTPUT = Path(os.environ.get("PYQSC_VMEC_BENCHMARK_OUTPUT", "benchmarks/results/vmec_export.json")) + +solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=NPHI, + order="r2", +) +with tempfile.TemporaryDirectory() as temporary_directory: + destination = Path(temporary_directory) / "input.benchmark" + jax.clear_caches() + cold = qsc.to_vmec( + solution, + destination, + r=RADIUS, + ntheta=NTHETA, + mpol=MPOL, + ntor=NTOR, + ) + warm = [ + qsc.to_vmec( + solution, + destination, + r=RADIUS, + ntheta=NTHETA, + mpol=MPOL, + ntor=NTOR, + ).conversion_seconds + for _ in range(WARM_REPETITIONS) + ] + +repository = Path(__file__).resolve().parents[1] +try: + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +report = { + "git_commit": commit, + "python": platform.python_version(), + "platform": platform.platform(), + "processor": platform.processor(), + "jax": jax.__version__, + "jax_backend": jax.default_backend(), + "jax_enable_x64": bool(jax.config.jax_enable_x64), + "case": { + "radius": RADIUS, + "nphi": NPHI, + "ntheta": NTHETA, + "mpol": MPOL, + "ntor": NTOR, + }, + "cold_compile_and_execute_seconds": cold.conversion_seconds, + "warm_median_seconds": statistics.median(warm), + "warm_samples_seconds": warm, + "maximum_toroidal_angle_residual": float(cold.boundary.maximum_toroidal_angle_residual), + "maximum_R_reconstruction_error": float(cold.boundary.maximum_R_reconstruction_error), + "maximum_Z_reconstruction_error": float(cold.boundary.maximum_Z_reconstruction_error), +} +OUTPUT.parent.mkdir(parents=True, exist_ok=True) +OUTPUT.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") +print("cold compile + execution [s]:", cold.conversion_seconds) +print("warm median [s]:", report["warm_median_seconds"]) +print("saved:", OUTPUT) diff --git a/benchmarks/benchmark_vmex_interface.py b/benchmarks/benchmark_vmex_interface.py new file mode 100644 index 0000000..1987797 --- /dev/null +++ b/benchmarks/benchmark_vmex_interface.py @@ -0,0 +1,97 @@ +"""Record synchronized VMEX forward and implicit-gradient timings.""" + +from __future__ import annotations + +import json +import os +import platform +import subprocess +from pathlib import Path +from time import perf_counter + +import jax +import numpy as np + +import pyqsc_jax as qsc + +CONFIGURATION = "plasma_stellarator" +RADIUS = 0.02 +OUTPUT = Path( + os.environ.get( + "PYQSC_VMEX_BENCHMARK_OUTPUT", + "benchmarks/results/vmex_interface.json", + ) +) + +near_axis = qsc.solve_configuration(CONFIGURATION, nphi=31) +problem = qsc.to_vmex_problem( + near_axis, + r=RADIUS, + qs_surfaces=(0.2, 0.4, 0.6, 0.8, 1.0), + ntheta=8, + mpol=3, + ntor=2, + ns_array=(7,), + ftol=1.0e-7, + max_iterations=1200, + adjoint_tol=1.0e-8, + multigrid=False, +) + +jax.clear_caches() +start = perf_counter() +equilibrium = problem.solve() +jax.block_until_ready(equilibrium.quantities) +forward_seconds = perf_counter() - start + + +def objective(parameters): + """Return the scalar used to benchmark VMEX's implicit gradient.""" + + return qsc.vmex_radial_quantities(problem, parameters).magnetic_well + + +start = perf_counter() +well, gradient = jax.value_and_grad(objective)(problem.parameters) +jax.block_until_ready((well, gradient)) +value_and_gradient_seconds = perf_counter() - start + +repository = Path(__file__).resolve().parents[1] +try: + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +report = { + "git_commit": commit, + "python": platform.python_version(), + "platform": platform.platform(), + "processor": platform.processor(), + "jax": jax.__version__, + "jax_backend": jax.default_backend(), + "jax_enable_x64": bool(jax.config.jax_enable_x64), + "vmex": problem.vmex_version, + "vmex_validated_commit": problem.validated_commit, + "case": { + "configuration": CONFIGURATION, + "radius": RADIUS, + "nphi": near_axis.inputs.nphi, + "ntheta": problem.ntheta, + "mpol_maximum": problem.mpol, + "ntor": problem.ntor, + "ns_array": np.asarray(problem.input.ns_array).tolist(), + "adjoint_tolerance": problem.adjoint_tol, + }, + "forward_seconds": forward_seconds, + "value_and_gradient_seconds": value_and_gradient_seconds, + "magnetic_well": float(well), + "gradient_pres_scale": float(gradient.pres_scale), + "gradient_rbc_norm": float(np.linalg.norm(np.asarray(gradient.rbc))), +} +OUTPUT.parent.mkdir(parents=True, exist_ok=True) +OUTPUT.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") +print("VMEX forward [s]:", forward_seconds) +print("VMEX magnetic-well value + gradient [s]:", value_and_gradient_seconds) +print("saved:", OUTPUT) diff --git a/benchmarks/reports/2026-07-29-apple-m4.json b/benchmarks/reports/2026-07-29-apple-m4.json new file mode 100644 index 0000000..7612b8b --- /dev/null +++ b/benchmarks/reports/2026-07-29-apple-m4.json @@ -0,0 +1,118 @@ +{ + "date": "2026-07-29", + "git_commit": "69b05da2dd8e7d6e734249e4d9f266f425dac256", + "hardware": "Apple M4, 10 CPU cores, 24 GB RAM", + "jax": "0.11.0", + "jax_backend": "cpu", + "jax_enable_x64": true, + "platform": "macOS-26.5.1-arm64-arm-64bit", + "python": "3.12.13", + "warm_repetitions": 7, + "results": [ + { + "nphi": 15, + "iota": 0.41818389451162313, + "cold_compile_and_execute_seconds": 0.2421206250146497, + "warm_samples_seconds": [ + 0.00009366602171212435, + 0.000059833022532982, + 0.00006033299723640084, + 0.0000551670091226697, + 0.000051125010941177607, + 0.000047290988732129335, + 0.000045833003241568804 + ], + "jvp_cold_compile_and_execute_seconds": 0.18876837499556132, + "jvp_warm_samples_seconds": [ + 0.00005804200191050768, + 0.00006270798621699214, + 0.00006150000263005495, + 0.000057124998420476913, + 0.00005129099008627236, + 0.00005458298255689442, + 0.00004983300459571183 + ], + "vmap_batch_size": 8, + "vmap_cold_compile_and_execute_seconds": 0.19951383300940506, + "vmap_warm_samples_seconds": [ + 0.0001232920039910823, + 0.00009733298793435097, + 0.0000914999982342124, + 0.0000757500238250941, + 0.00007770801312290132, + 0.00006995900184847414, + 0.00007645899313502014 + ] + }, + { + "nphi": 31, + "iota": 0.41830690943386595, + "cold_compile_and_execute_seconds": 0.1799230409960728, + "warm_samples_seconds": [ + 0.00007320797885768116, + 0.00006308397860266268, + 0.00006616700557619333, + 0.00006529100937768817, + 0.000054208998335525393, + 0.00006212500738911331, + 0.00005641698953695595 + ], + "jvp_cold_compile_and_execute_seconds": 0.20157658398966305, + "jvp_warm_samples_seconds": [ + 0.00009200000204145908, + 0.0000849170028232038, + 0.00008412500028498471, + 0.00006916699931025505, + 0.00006808299804106355, + 0.00007066698162816465, + 0.00006887500057928264 + ], + "vmap_batch_size": 8, + "vmap_cold_compile_and_execute_seconds": 0.2098307500127703, + "vmap_warm_samples_seconds": [ + 0.00025675000506453216, + 0.0002603330067358911, + 0.00024779202067293227, + 0.00023179201525636017, + 0.00022820799495093524, + 0.00023674999829381704, + 0.00023174998932518065 + ] + }, + { + "nphi": 61, + "iota": 0.41830691021517663, + "cold_compile_and_execute_seconds": 0.19822412499343045, + "warm_samples_seconds": [ + 0.00015645800158381462, + 0.00015158398309722543, + 0.0001516659976914525, + 0.000144416990224272, + 0.00014354099403135478, + 0.00013208299060352147, + 0.00013691699132323265 + ], + "jvp_cold_compile_and_execute_seconds": 0.21340283399331383, + "jvp_warm_samples_seconds": [ + 0.00022141702356748283, + 0.0001960420049726963, + 0.00019645900465548038, + 0.00018620799528434873, + 0.00018912501400336623, + 0.00019108300330117345, + 0.0001844169746618718 + ], + "vmap_batch_size": 8, + "vmap_cold_compile_and_execute_seconds": 0.21318525000242516, + "vmap_warm_samples_seconds": [ + 0.0004967089917045087, + 0.0004771670210175216, + 0.0004792909894604236, + 0.0004676670068874955, + 0.0004896250029560179, + 0.0004884999943897128, + 0.00046212499728426337 + ] + } + ] +} diff --git a/benchmarks/reports/2026-07-30-b20-optimizers-apple-m4.json b/benchmarks/reports/2026-07-30-b20-optimizers-apple-m4.json new file mode 100644 index 0000000..17704dd --- /dev/null +++ b/benchmarks/reports/2026-07-30-b20-optimizers-apple-m4.json @@ -0,0 +1,52 @@ +{ + "git_commit": "8fb457deeb8b71dca0bc3d06d2da5fd9cbc99c4f", + "jax": "0.11.0", + "jax_backend": "cpu", + "jax_enable_x64": true, + "modes": [ + 1, + 2, + 3 + ], + "nphi": 121, + "platform": "macOS-26.5.1-arm64-arm-64bit-Mach-O", + "python": "3.13.7", + "results": [ + { + "function_evaluations": 1, + "method": "exact B2c only", + "seconds": 0.0, + "weighted_l2": 0.030992772869247103 + }, + { + "function_evaluations": 70, + "method": "SciPy L-BFGS-B", + "seconds": 0.1407689999905415, + "weighted_l2": 0.00020139241185160374 + }, + { + "function_evaluations": 28, + "method": "SciPy least_squares", + "seconds": 1.3488249160000123, + "weighted_l2": 0.00010250101789370511 + }, + { + "function_evaluations": 120, + "method": "SciPy differential_evolution, low budget", + "seconds": 0.10409141599666327, + "weighted_l2": 0.2657663516222627 + }, + { + "distinct_basins": 3, + "function_evaluations": 130, + "method": "pyQSC_JAX multistart Levenberg-Marquardt", + "seconds": 21.77209062501788, + "status": "best_found", + "weighted_l2": 0.00010827903169153683 + } + ], + "selected_configuration": "b20_optimized_good", + "source_database_id": 57409, + "source_url": "https://stellarator.physics.wisc.edu/app/plot/57409", + "timing_note": "SciPy timings exclude shared residual/Jacobian compilation; the pyQSC_JAX multistart timing includes compilation of its independent closure." +} diff --git a/benchmarks/reports/2026-07-30-b20-vmec-apple-m4.json b/benchmarks/reports/2026-07-30-b20-vmec-apple-m4.json new file mode 100644 index 0000000..3a26ae8 --- /dev/null +++ b/benchmarks/reports/2026-07-30-b20-vmec-apple-m4.json @@ -0,0 +1,113 @@ +{ + "date": "2026-07-30", + "source_tree_base_commit": "f10549a3539b475cbfea2dba5cdf913a6516b75f", + "hardware": "Apple M4, 10 CPU cores, 24 GB RAM", + "jax": "0.11.0", + "jax_backend": "cpu", + "jax_enable_x64": true, + "platform": "macOS-26.5.1-arm64-arm-64bit-Mach-O", + "python": "3.13.7", + "b20_optimizer_comparison": { + "common_case": "Curvo stellarator-database configuration 57409; Fourier modes 1-3 varied; exact B2c eliminated at every evaluation; nphi=121 verification", + "source_database_id": 57409, + "source_url": "https://stellarator.physics.wisc.edu/app/plot/57409", + "timing_note": "SciPy timings exclude shared residual/Jacobian compilation; the pyQSC_JAX multistart timing includes compilation of its independent closure.", + "results": [ + { + "function_evaluations": 1, + "method": "exact B2c only", + "seconds": 0.0, + "weighted_l2": 0.030992772869247103 + }, + { + "function_evaluations": 70, + "method": "SciPy L-BFGS-B", + "seconds": 0.18148858298081905, + "weighted_l2": 0.00020139241185160374 + }, + { + "function_evaluations": 28, + "method": "SciPy least_squares", + "seconds": 1.309102333005285, + "weighted_l2": 0.00010250101789370511 + }, + { + "function_evaluations": 120, + "method": "SciPy differential_evolution, low budget", + "seconds": 0.11324654199415818, + "weighted_l2": 0.2657663516222627 + }, + { + "distinct_basins": 3, + "function_evaluations": 130, + "method": "pyQSC_JAX multistart Levenberg-Marquardt", + "seconds": 21.91377937499783, + "status": "best_found", + "weighted_l2": 0.00010827903169153683 + } + ], + "selected_refinement": { + "function_evaluations": 222, + "method": "staged bounded SciPy least_squares through Fourier mode 8", + "seconds": 24.129845416988246, + "curvo_profile_minimum_abs_iota": 0.4, + "curvo_profile_passed": true, + "abs_iota": 2.9635796242368357, + "singular_radius_m": 0.24925039358398768, + "verification_nphi": [121, 241, 481], + "verification_weighted_l2": [ + 1.27431807e-10, + 1.27645060e-10, + 1.28075313e-10 + ] + } + }, + "vmec_export": { + "case": "QA r2, radius=0.03, nphi=61, ntheta=40, mpol=12, ntor=14", + "cold_compile_and_execute_seconds": 0.4759971250023227, + "warm_median_seconds": 0.005203832988627255, + "warm_samples_seconds": [ + 0.005135457991855219, + 0.005203832988627255, + 0.005606082995655015, + 0.005212290998315439, + 0.005193958000745624, + 0.0052505840139929205, + 0.0050905829994007945 + ], + "maximum_R_reconstruction_error_m": 3.4204787844327456e-06, + "maximum_Z_reconstruction_error_m": 2.261107830789366e-06, + "maximum_toroidal_angle_residual": 2.22e-16 + }, + "vmec_radius_convergence": { + "near_axis_iota": 0.41830691021517735, + "vmec_version": "9.0", + "results": [ + { + "radius": 0.0025, + "relative_iota_error": 0.0005645603050142573, + "vmec_iota_axis": 0.418543069691998 + }, + { + "radius": 0.005, + "relative_iota_error": 0.002196, + "vmec_iota_axis": 0.41922558234541574 + }, + { + "radius": 0.01, + "relative_iota_error": 0.009403, + "vmec_iota_axis": 0.42224025 + }, + { + "radius": 0.02, + "relative_iota_error": 0.050798, + "vmec_iota_axis": 0.43955604 + }, + { + "radius": 0.03, + "relative_iota_error": 0.173066, + "vmec_iota_axis": 0.49070165 + } + ] + } +} diff --git a/docs/_static/B20_optimization.png b/docs/_static/B20_optimization.png new file mode 100644 index 0000000..b294f60 Binary files /dev/null and b/docs/_static/B20_optimization.png differ diff --git a/docs/_static/QA_QH_branches.png b/docs/_static/QA_QH_branches.png new file mode 100644 index 0000000..503006e Binary files /dev/null and b/docs/_static/QA_QH_branches.png differ diff --git a/docs/_static/axis_and_surfaces.png b/docs/_static/axis_and_surfaces.png new file mode 100644 index 0000000..0973363 Binary files /dev/null and b/docs/_static/axis_and_surfaces.png differ diff --git a/docs/_static/convergence.png b/docs/_static/convergence.png new file mode 100644 index 0000000..359b592 Binary files /dev/null and b/docs/_static/convergence.png differ diff --git a/docs/_static/core_performance.png b/docs/_static/core_performance.png new file mode 100644 index 0000000..32dcb5e Binary files /dev/null and b/docs/_static/core_performance.png differ diff --git a/docs/_static/optimizer_comparison.png b/docs/_static/optimizer_comparison.png new file mode 100644 index 0000000..fef9efd Binary files /dev/null and b/docs/_static/optimizer_comparison.png differ diff --git a/docs/_static/plasma_external_jet.png b/docs/_static/plasma_external_jet.png new file mode 100644 index 0000000..76df750 Binary files /dev/null and b/docs/_static/plasma_external_jet.png differ diff --git a/docs/_static/stellarator_gallery.png b/docs/_static/stellarator_gallery.png new file mode 100644 index 0000000..22ff018 Binary files /dev/null and b/docs/_static/stellarator_gallery.png differ diff --git a/docs/_static/vmec_validation.png b/docs/_static/vmec_validation.png new file mode 100644 index 0000000..f45eab9 Binary files /dev/null and b/docs/_static/vmec_validation.png differ diff --git a/docs/_static/vmex_radial_profiles.png b/docs/_static/vmex_radial_profiles.png new file mode 100644 index 0000000..f6ee05a Binary files /dev/null and b/docs/_static/vmex_radial_profiles.png differ diff --git a/docs/adr/ADR-core-architecture.md b/docs/adr/ADR-core-architecture.md new file mode 100644 index 0000000..79f58d5 --- /dev/null +++ b/docs/adr/ADR-core-architecture.md @@ -0,0 +1,38 @@ +# ADR: immutable functional core with compatibility adapters + +- Status: accepted +- Date: 2026-07-29 + +## Context + +The current mutable `near_axis` class mixes inputs, derived arrays, nonlinear +solves, coordinate conversion, field evaluation, plotting, and optimizer state. +Passing the mutable object as a static JIT argument makes cache behavior and +automatic differentiation fragile. ESSOS nevertheless requires the historical +class and mutable degree-of-freedom interface. + +## Decision + +Build a canonical immutable API from small frozen dataclasses registered as JAX +pytrees and pure functions. Keep array-valued physics data as pytree leaves and +small topology/resolution choices as static metadata. Ordinary dataclasses are +preferred initially; Equinox will be added only if it removes demonstrated +complexity. + +The public API will provide `Axis`, `Qsc`, and `solve`. The legacy +`pyqsc_jax.near_axis.near_axis` class will be a thin adapter that reconstructs +the immutable solution when its mutable degrees of freedom change. + +Modules will follow physics and ownership boundaries, beginning with models, +spectral operators, axis geometry, solvers, and first order. Empty speculative +modules and generic `utils.py` are prohibited. + +## Consequences + +- JIT, VMAP, JVP, and VJP operate on explicit array arguments. +- Solver reports and invalid-geometry diagnostics become part of results. +- Compatibility mutation remains isolated and testable. +- General asymmetric axis coefficients can be supported without duplicating + geometry formulas. +- The adapter may allocate a new solution after mutation; this cost is + acceptable for compatibility and must not constrain the core. diff --git a/docs/adr/ADR-solver-stack.md b/docs/adr/ADR-solver-stack.md new file mode 100644 index 0000000..b4a0d2c --- /dev/null +++ b/docs/adr/ADR-solver-stack.md @@ -0,0 +1,49 @@ +# ADR: solver stack and implicit differentiation + +- Status: accepted for first implementation +- Date: 2026-07-29 + +## Context + +First order requires a converged periodic nonlinear sigma solve. Second order +requires a dense coupled linear solve. Inverse transform modes add a scalar +unknown and continuation. Global/local axis optimization may later benefit from +Levenberg–Marquardt. + +The current code runs exactly five Newton iterations. Differentiating those +iterations does not represent the derivative of the converged equation and +provides no convergence evidence. + +## Options considered + +1. Raw JAX loops and `jax.numpy.linalg.solve`: minimal, but every implicit + derivative and transpose rule would be local maintenance. +2. Equinox/Optimistix: capable, but adds a broad solver dependency before a + missing capability is demonstrated. +3. SOLVAX `root_solve` and `linear_solve`: already provide JAX custom-root and + custom-linear-solve implicit differentiation around user-supplied primal + solvers. +4. A new SOLVAX solver API: justified only by a generic, independently tested + need beyond pyQSC_JAX. + +## Decision + +Use small pyQSC_JAX primal Newton and dense linear callbacks with explicit +residual, step, convergence, iteration, and conditioning reports. Wrap +converged roots and linear solutions with SOLVAX for implicit differentiation. + +Do not change SOLVAX initially. Revisit a general dense +Levenberg–Marquardt implementation only after the pyQSC_JAX optimization +requirements are concrete and the capability is demonstrably reusable. + +Pseudo-arclength continuation and fold detection belong in pyQSC_JAX because +they encode near-axis branch semantics. + +## Acceptance + +- primal residuals satisfy configured absolute and relative tolerances; +- iteration limits produce explicit failure reports; +- JVP/VJP agree with finite differences away from singular points; +- transpose solves are tested; +- second-order conditioning is reported, especially near `iota_N = 0`; +- inverse solves detect folds instead of silently changing branches. diff --git a/docs/advanced/autodiff.md b/docs/advanced/autodiff.md new file mode 100644 index 0000000..e7ed74e --- /dev/null +++ b/docs/advanced/autodiff.md @@ -0,0 +1,18 @@ +# Automatic differentiation + +The immutable solution records are registered JAX pytrees. Root and linear +solves differentiate converged equations implicitly instead of +backpropagating through iteration histories. Functions intended for +transformation accept physical arrays explicitly; static grid sizes and model +choices remain Python metadata. + +```{literalinclude} ../../examples/12_autodiff_check.py +:language: python +:linenos: +``` + +The test suite checks eager/JIT parity, VMAP batches, JVP/VJP consistency, and +finite differences for axis geometry, first and second order, inverse solves, +optimization, and plasma jets. Nonsmooth decisions such as basin selection +are outside a local derivative; differentiate a selected candidate's smooth +residual, not the discrete search procedure. diff --git a/docs/advanced/continuation.md b/docs/advanced/continuation.md new file mode 100644 index 0000000..fbe05c3 --- /dev/null +++ b/docs/advanced/continuation.md @@ -0,0 +1,17 @@ +# Continuation and folds + +Target-\(\iota\) inversion is local and can be multivalued when +\(\partial\iota/\partial\bar\eta=0\). `continue_etabar_branch` instead traces +the complete sigma collocation state using a secant predictor and augmented +pseudo-arclength corrector. It preserves the sign of `etabar`. + +```{literalinclude} ../../examples/11_continuation_scan.py +:language: python +:linenos: +``` + +`ContinuationResult` retains every corrected immutable solution, tangent, +response derivative, and fold flag. A partial branch reports +`solver_failure` or `branch_zero_crossing`; it is never labeled complete. +The [inverse/continuation derivation](../theory/inverse-solves.md) gives the +augmented equation and acceptance rules. diff --git a/docs/advanced/custom_criteria.md b/docs/advanced/custom_criteria.md new file mode 100644 index 0000000..167078a --- /dev/null +++ b/docs/advanced/custom_criteria.md @@ -0,0 +1,24 @@ +# Custom criteria + +`Criteria.from_curvo_2025` reproduces a named published profile with scalable +dimensionful thresholds. Every threshold can be overridden: + +```python +import pyqsc_jax as qsc + +criteria = qsc.Criteria.from_curvo_2025( + major_radius=1.7, + B0=2.5, + maximum_elongation=8.0, + minimum_beta=2.0e-4, +) +report = criteria.evaluate(solution) +for evaluation in report.evaluations: + print(evaluation.name, evaluation.value, evaluation.margin) +``` + +Margins are signed so positive is favorable for both lower and upper bounds. +The profile is a screening tool, not a universal definition of feasibility. +Coil engineering, finite-radius equilibrium, and particle confinement remain +separate analyses. Definitions and normalization are in +[named design criteria](../theory/criteria.md). diff --git a/docs/advanced/global_search.md b/docs/advanced/global_search.md new file mode 100644 index 0000000..74eda80 --- /dev/null +++ b/docs/advanced/global_search.md @@ -0,0 +1,19 @@ +# Global-search semantics + +The workflow is global-to-local, not a proof-producing global optimizer: + +1. deterministic branch-specific and Halton seeds; +2. batched coarse evaluation; +3. damped local least-squares refinement; +4. normalized basin clustering; +5. hard criteria and conditioning filters; +6. independent resolution verification. + +Only a verified zero reaches the known global lower bound of the nonnegative +mean-free \(B_{20}\) residual. Nonzero output is `best_found` within the +reported finite budget. The implementation never turns the best sampled +point into a global claim. + +See the [full search derivation](../theory/global-search.md) for status values, +bounded variables, high-mode safeguards, and continuation across Fourier +stages. diff --git a/docs/advanced/limitations.md b/docs/advanced/limitations.md new file mode 100644 index 0000000..f5f1323 --- /dev/null +++ b/docs/advanced/limitations.md @@ -0,0 +1,32 @@ +# Limitations + +- Near-axis results are asymptotic and do not prove a finite-radius + equilibrium or particle confinement. +- Frenet coordinates fail where axis curvature vanishes. +- Target-transform inversions are branch-local; folds require continuation. +- Pseudo-arclength continuation currently traces `etabar` with a fixed step. +- Magnetic shear currently supports the documented standard-MHS sign + specialization. +- A finite multistart budget cannot certify a nonzero global minimum. +- `r_singularity` diagnoses the truncated coordinate map, not equilibrium + existence. +- Plasma-field and Hessian remainder fields are asymptotic scales, not + rigorous bounds. +- The optional VMEX bridge covers differentiable fixed-boundary equilibria. + It does not differentiate a reconverged free-boundary NESTOR root. +- Traceable VMEX quasisymmetry profiles currently require stellarator + symmetry; use `qs_surfaces=()` for an asymmetric equilibrium without QS. +- VMEX magnetic well is an endpoint scalar, not a radial profile. +- VMEC/VMEX boundary conversion requires the near-axis surface to admit the + cylindrical-toroidal angle inversion at the requested radius. The + conversion reports this condition and disk-writing APIs reject failures; + a singular-radius estimate alone does not guarantee cylindrical + star-shapedness. +- “Surface-free” plasma/coil separation still requires a positive formal + radius or equivalent current/flux normalization. +- Coil feasibility and fast-particle confinement remain ESSOS or external + calculations. + +Nonconverged solvers return inspectable result/report objects. Production +callers must reject false `converged` or `finite` flags and must enforce their +own conditioning and resolution thresholds. diff --git a/docs/advanced/performance.md b/docs/advanced/performance.md new file mode 100644 index 0000000..9a27df0 --- /dev/null +++ b/docs/advanced/performance.md @@ -0,0 +1,94 @@ +# Performance + +Performance must be separated into compilation and execution: + +- cold compile plus first execution; +- warm repeated execution; +- batched VMAP throughput; +- reverse- or forward-mode derivative cost; +- scaling with `nphi` and Fourier mode count. + +The scheduled benchmark workflow records the exact dependency environment and +runs correctness checks before timing. Shared-runner timing is reported, not +used as a fragile pass/fail threshold. The main dense costs scale with the +periodic sigma Jacobian and r2 linear system; publication searches should +reuse static shapes to avoid recompilation. + +For reproducible local measurements, synchronize JAX results before stopping a +timer: + +```python +value = compiled_function(arguments) +value.iota.block_until_ready() +``` + +The release report records the benchmark command, hardware, versions, and +raw samples. No absolute performance claim is inferred from a single laptop +or CI runner. + +## Refactor benchmark + +The following synchronized CPU measurements were recorded on an Apple M4 +(10 CPU cores, 24 GB RAM), macOS arm64, Python 3.12.13, JAX 0.11.0 with +64-bit mode, at commit `69b05da`. “Cold” clears JAX caches and includes +compilation plus execution for that static shape; it does not include Python +process or XLA-runtime startup. Warm values are medians of seven samples. +VMAP evaluates a batch of eight `etabar` values. + +| `nphi` | solve cold [s] | solve warm [µs] | JVP cold [s] | JVP warm [µs] | VMAP cold [s] | VMAP warm [µs] | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 15 | 0.242 | 55.2 | 0.189 | 57.1 | 0.200 | 77.7 | +| 31 | 0.180 | 63.1 | 0.202 | 70.7 | 0.210 | 236.7 | +| 61 | 0.198 | 144.4 | 0.213 | 191.1 | 0.213 | 479.3 | + +Compiler overhead dominates these small cold cases, while warm execution +increases with collocation size. These values are a reproducible development +snapshot, not a cross-platform speed guarantee. The exact command was: + +```bash +JAX_ENABLE_X64=true PYTHONPATH=src python benchmarks/benchmark_core.py +``` + +All raw timing samples are stored in +`benchmarks/reports/2026-07-29-apple-m4.json`. + +## VMEC boundary export + +The VMEC converter is separately benchmarked because its old scalar +point-by-point root solves obscured the cost of the actual Fourier projection. +The replacement performs one JIT-compiled, vectorized Newton inversion and a +two-dimensional FFT. + +| grid and spectrum | compile + execute | warm median | max \(R\) error | max \(Z\) error | +| --- | ---: | ---: | ---: | ---: | +| `nphi=61`, `ntheta=40`, `mpol=12`, `ntor=14` | 0.476 s | 5.20 ms | 3.42 µm | 2.26 µm | + +The command is: + +```bash +JAX_ENABLE_X64=true PYTHONPATH=src python benchmarks/benchmark_vmec_export.py +``` + +The reported conversion timer excludes the near-axis solve and file-system +startup. Both cold and warm calls synchronize the boundary arrays before +stopping the timer. + +## VMEX implicit equilibrium + +The 5.20 ms number above is only the boundary conversion. A converged VMEX +equilibrium is a separate fixed-point solve. Its cold cost includes XLA +compilation and its gradient includes an implicit adjoint linear solve. + +The low-resolution live compatibility case (`ns=7`, `mpol=4`, `ntor=2`) took +23.1 s for a cold forward solve and 17.5 s for a subsequent magnetic-well +value-and-gradient on the development Apple CPU before persistent-cache +reuse. These are integration-smoke timings, not production performance +claims. Repeated optimization should reuse one `VmexProblem`, static shapes, +and VMEX's compilation cache. + +The synchronized benchmark that records the dependency versions, validated +VMEX commit, case parameters, objective value, and gradient norms is: + +```bash +JAX_ENABLE_X64=true PYTHONPATH=src python benchmarks/benchmark_vmex_interface.py +``` diff --git a/docs/api/axis-optimization.md b/docs/api/axis-optimization.md new file mode 100644 index 0000000..90a4092 --- /dev/null +++ b/docs/api/axis-optimization.md @@ -0,0 +1,43 @@ +# Axis-optimization API + +```{eval-rst} +.. autofunction:: pyqsc_jax.stellarator_symmetric_variable_indices +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.search_axis +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.continue_axis_search +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.AxisSearchProblem + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.AxisSearchOptions + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.AxisSearchCandidate + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.AxisSearchResult + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.AxisSearchContinuation + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.LocalLeastSquaresReport + :members: +``` diff --git a/docs/api/axis.md b/docs/api/axis.md new file mode 100644 index 0000000..0e980c1 --- /dev/null +++ b/docs/api/axis.md @@ -0,0 +1,29 @@ +# Axis and geometry API + +```{eval-rst} +.. autoclass:: pyqsc_jax.Axis + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.axis.evaluate_axis +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.axis.AxisSamples + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.geometry.compute_axis_geometry +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.geometry.AxisGeometry + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.models.GeometryDiagnostics + :members: +``` diff --git a/docs/api/configurations.md b/docs/api/configurations.md new file mode 100644 index 0000000..4f1ad70 --- /dev/null +++ b/docs/api/configurations.md @@ -0,0 +1,18 @@ +# Reference configurations + +```{eval-rst} +.. autoclass:: pyqsc_jax.ReferenceConfiguration + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.available_configurations +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.get_configuration +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.solve_configuration +``` diff --git a/docs/api/continuation.md b/docs/api/continuation.md new file mode 100644 index 0000000..a9de35f --- /dev/null +++ b/docs/api/continuation.md @@ -0,0 +1,10 @@ +# Continuation API + +```{eval-rst} +.. autofunction:: pyqsc_jax.continue_etabar_branch +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.ContinuationResult + :members: +``` diff --git a/docs/api/criteria.md b/docs/api/criteria.md new file mode 100644 index 0000000..fcbb7f7 --- /dev/null +++ b/docs/api/criteria.md @@ -0,0 +1,16 @@ +# Criteria API + +```{eval-rst} +.. autoclass:: pyqsc_jax.Criteria + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.CriteriaReport + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.CriterionEvaluation + :members: +``` diff --git a/docs/api/field-jet.md b/docs/api/field-jet.md new file mode 100644 index 0000000..55d4ed7 --- /dev/null +++ b/docs/api/field-jet.md @@ -0,0 +1,28 @@ +# Field-jet and equilibrium diagnostics API + +```{eval-rst} +.. autoclass:: pyqsc_jax.FieldJet + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.total_field_jet +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.MercierDiagnostics + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.mercier_diagnostics +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.SingularityDiagnostics + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.singularity_diagnostics +``` diff --git a/docs/api/first-order.md b/docs/api/first-order.md new file mode 100644 index 0000000..35272fb --- /dev/null +++ b/docs/api/first-order.md @@ -0,0 +1,41 @@ +# First-order API + +```{eval-rst} +.. autofunction:: pyqsc_jax.Qsc +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.solve +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.NearAxisInputs + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.NearAxisSolution + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.RootSolveOptions + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.RootSolveReport + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.InverseSolveDiagnostics + :members: +``` + +## ESSOS compatibility + +```{eval-rst} +.. autoclass:: pyqsc_jax.near_axis.near_axis + :members: +``` diff --git a/docs/api/optimization.md b/docs/api/optimization.md new file mode 100644 index 0000000..6851ae5 --- /dev/null +++ b/docs/api/optimization.md @@ -0,0 +1,32 @@ +# Optimization API + +```{eval-rst} +.. autofunction:: pyqsc_jax.b20_diagnostics +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.optimal_B2c_value +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.optimize_B2c +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.verify_B20_resolution +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.B20Diagnostics + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.B2cOptimizationResult + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.B20ResolutionVerification + :members: +``` diff --git a/docs/api/plasma.md b/docs/api/plasma.md new file mode 100644 index 0000000..5b38375 --- /dev/null +++ b/docs/api/plasma.md @@ -0,0 +1,85 @@ +# Plasma/current API + +```{eval-rst} +.. autofunction:: pyqsc_jax.plasma_hessian_on_axis +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plasma_gradient_on_axis +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.elliptical_channel_gradient +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.project_symmetric_trace_free_rank2 +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.pack_symmetric_trace_free_rank2 +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.unpack_symmetric_trace_free_rank2 +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.project_symmetric_trace_free_rank3 +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.pack_symmetric_trace_free_rank3 +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.unpack_symmetric_trace_free_rank3 +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plasma_field_on_axis +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.regularized_axis_integral +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.matched_plasma_field_kernel +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plasma_current_source +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.evaluate_weighted_current +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.enclosed_current_from_covariant +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.covariant_current_from_enclosed +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.PlasmaCurrentSource + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.PlasmaFieldData + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.PlasmaGradientData + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.PlasmaHessianData + :members: +``` diff --git a/docs/api/plotting.md b/docs/api/plotting.md new file mode 100644 index 0000000..08e8dca --- /dev/null +++ b/docs/api/plotting.md @@ -0,0 +1,33 @@ +# Plotting + +Matplotlib is imported lazily and is available through the `plot` extra. +Every helper returns its figure and axes; it never calls `show` or saves +implicitly. + +```{eval-rst} +.. autofunction:: pyqsc_jax.plotting.plot_axis +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plotting.surface_coordinates +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plotting.plot_surface_3d +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plotting.plot_b20 +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plotting.plot_field_jet_norms +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plotting.field_split_frenet_components +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.plotting.plot_field_split_components +``` diff --git a/docs/api/second-order.md b/docs/api/second-order.md new file mode 100644 index 0000000..00bf4d2 --- /dev/null +++ b/docs/api/second-order.md @@ -0,0 +1,20 @@ +# Second-order API + +```{eval-rst} +.. autoclass:: pyqsc_jax.SecondOrderData + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.LinearSolveReport + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.SecondOrderResiduals + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.second_order_residuals +``` diff --git a/docs/api/third-order.md b/docs/api/third-order.md new file mode 100644 index 0000000..3a748ee --- /dev/null +++ b/docs/api/third-order.md @@ -0,0 +1,19 @@ +# Third-order API + +```{eval-rst} +.. autoclass:: pyqsc_jax.ThirdOrderData + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.solve_third_order +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.ShearData + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.solve_magnetic_shear +``` diff --git a/docs/api/vmec.md b/docs/api/vmec.md new file mode 100644 index 0000000..10635d7 --- /dev/null +++ b/docs/api/vmec.md @@ -0,0 +1,16 @@ +# VMEC export + +```{eval-rst} +.. automodule:: pyqsc_jax.vmec + :members: + :undoc-members: + :show-inheritance: +``` + +`to_vmec` is also available as a method on the legacy +`pyqsc_jax.near_axis.near_axis` adapter. That method preserves the upstream +call form and populates `RBC`, `RBS`, `ZBC`, and `ZBS` compatibility arrays. + +`VmecBoundary.toroidal_angle_converged` and +`maximum_toroidal_angle_residual` diagnose the vectorized inversion. The +writer never emits an input when the configured tolerance is not reached. diff --git a/docs/api/vmex.md b/docs/api/vmex.md new file mode 100644 index 0000000..be46906 --- /dev/null +++ b/docs/api/vmex.md @@ -0,0 +1,35 @@ +# VMEX equilibrium bridge + +VMEX is an optional dependency and is imported only when the equilibrium +bridge is used. + +```{eval-rst} +.. autoclass:: pyqsc_jax.VmexProblem + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.VmexEquilibrium + :members: +``` + +```{eval-rst} +.. autoclass:: pyqsc_jax.VmexRadialQuantities + :members: +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.to_vmex_problem +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.vmex_parameters_from_solution +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.solve_vmex +``` + +```{eval-rst} +.. autofunction:: pyqsc_jax.vmex_radial_quantities +``` diff --git a/docs/changelog.md b/docs/changelog.md new file mode 100644 index 0000000..8169802 --- /dev/null +++ b/docs/changelog.md @@ -0,0 +1,5 @@ +# Changelog + +```{literalinclude} ../CHANGELOG.md +:language: markdown +``` diff --git a/docs/concepts/coordinates_and_conventions.md b/docs/concepts/coordinates_and_conventions.md new file mode 100644 index 0000000..9473e53 --- /dev/null +++ b/docs/concepts/coordinates_and_conventions.md @@ -0,0 +1,24 @@ +# Coordinates and conventions + +The magnetic axis uses cylindrical azimuth \(\phi\). Boozer toroidal angle +\(\varphi\) is normalized to advance by \(2\pi/n_{\mathrm{fp}}\) per field +period. The helical angle is + +\[ +\vartheta=\theta-Nn_{\mathrm{fp}}\varphi, +\qquad +\iota_N=\iota+Nn_{\mathrm{fp}}. +\] + +Canonical Cartesian vectors are ordered `(x, y, z)` and sampled arrays place +the toroidal index first. Cylindrical vectors are ordered `(R, phi, Z)`. +The Frenet basis is right-handed `(tangent, normal, binormal)`, while pyQSC's +historical Hessian representation uses derivative axes in +`(normal, binormal, tangent)` order. + +`sG` is the sign of \(G_0\), `spsi` the sign of toroidal flux, and +`chi = sG * spsi` in current conversions. Torsion follows Landreman and +Sengupta, opposite to the original Garren--Boozer sign convention. + +See the [full coordinate derivation](../theory/coordinates-and-conventions.md) +for Fourier definitions, frame topology, and validity diagnostics. diff --git a/docs/concepts/inputs_and_outputs.md b/docs/concepts/inputs_and_outputs.md new file mode 100644 index 0000000..78ff84e --- /dev/null +++ b/docs/concepts/inputs_and_outputs.md @@ -0,0 +1,54 @@ +# Inputs and outputs + +## Inputs + +`Axis` accepts `rc`, `rs`, `zc`, `zs`, and integer `nfp`. +`Qsc`/`solve` normalize the physical inputs into `NearAxisInputs`: + +| Name | Meaning | +| --- | --- | +| `axis` or `rc, rs, zc, zs, nfp` | magnetic-axis Fourier representation | +| `etabar` | signed first-order field-strength coefficient [m\(^{-1}\)] | +| `B0` | on-axis field strength [T] | +| `sigma0` | initial cross-section tilt | +| `I2` | quadratic covariant-current coefficient [T/m] | +| `p2` | quadratic pressure coefficient [Pa/m\(^2\)] | +| `B2c`, `B2s` | second-order field-strength harmonics [T/m\(^2\)] | +| `nphi` | samples per field period | +| `order` | `r1`, `r2`, or `r3` | +| `sG`, `spsi` | signs of \(G_0\) and toroidal flux | +| `iota`, `solve_for` | optional inverse target and solved parameter | + +All physical scalar inputs must be scalar arrays. `nphi`, order, sign flags, +and the solved-parameter choice are static pytree metadata. + +## Common outputs + +Every `NearAxisSolution` contains `inputs`, `geometry`, `root_report`, +`sigma`, `iota`, `iotaN`, `helicity`, `G0`, the first-order coefficients +`X1s`, `X1c`, `Y1s`, `Y1c` and untwisted forms, `elongation`, +`mean_elongation`, `B_axis`, `grad_B_axis`, and `L_grad_B`. + +An r2 result adds the complete `SecondOrderData`: `V1`, `V2`, `V3`, +`X20`, `X2s`, `X2c`, `Y20`, `Y2s`, `Y2c`, `Z20`, `Z2s`, `Z2c`, their +needed derivatives and untwisted forms, `beta_1s`, `G2`, `B20`, +`B20_mean`, `B20_anomaly`, `B20_residual`, and `B20_variation`. The r2 +solution also attaches singular-radius diagnostics, Mercier terms, and the +total `FieldJet`. + +`FieldJet` contains `field`, `gradient`, `hessian`, `hessian_frenet`, +coordinate-map derivatives, reconstruction errors, Maxwell residuals, and +\(L_{\nabla\nabla B}\). Canonical shapes are `(nphi, 3)`, +`(nphi, 3, 3)`, and `(nphi, 3, 3, 3)`, with field component before +derivative indices. + +An r3 result adds the first- and third-poloidal-harmonic surface coefficients, +their untwisted forms, `B0_order_a_squared_to_cancel`, +`flux_constraint_residual`, and `consistency_error`. Explicit shear adds +`B31c`, `iota2`, and the intermediate `ShearData`. + +Inverse, continuation, optimization, criteria, singularity, and plasma APIs +return their own immutable records with reports rather than overloading a +single mutable object. The API pages enumerate their public fields and the +[physics traceability table](../development/physics-traceability.md) maps +them to equations and tests. diff --git a/docs/concepts/magnetic_axis.md b/docs/concepts/magnetic_axis.md new file mode 100644 index 0000000..5c935b1 --- /dev/null +++ b/docs/concepts/magnetic_axis.md @@ -0,0 +1,26 @@ +# Magnetic axis + +`Axis` stores general Fourier series + +\[ +R(\phi)=\sum_n[R_{cn}\cos(nn_{\mathrm{fp}}\phi)+R_{sn}\sin(nn_{\mathrm{fp}}\phi)], +\] + +\[ +Z(\phi)=\sum_n[Z_{cn}\cos(nn_{\mathrm{fp}}\phi)+Z_{sn}\sin(nn_{\mathrm{fp}}\phi)]. +\] + +The public coefficient order is always `rc`, `rs`, `zc`, `zs`. Arrays are +zero-padded to a common Fourier length. Stellarator symmetry is the special +case `rs = zc = 0`; it is not assumed by the geometry core. + +`compute_axis_geometry` evaluates analytic Fourier derivatives, arc length, +curvature, torsion, Frenet frames, helicity, the \(\phi\leftrightarrow\varphi\) +map, and periodic spectral operators. Reject a result if +`geometry.diagnostics.frenet_valid` or +`cylindrical_coordinates_valid` is false. Curvature approaching zero is a +coordinate failure, not a benign numerical warning. + +Implementation: `pyqsc_jax.axis` and `pyqsc_jax.geometry`. Validation: +`tests/physics/test_geometry.py` and the audited upstream regression data in +`tests/regression/test_geometry_reference.py`. diff --git a/docs/concepts/precision_and_units.md b/docs/concepts/precision_and_units.md new file mode 100644 index 0000000..765f4ff --- /dev/null +++ b/docs/concepts/precision_and_units.md @@ -0,0 +1,28 @@ +# Precision and units + +The code uses SI units: + +- length in metres; +- magnetic field in tesla; +- pressure in pascals; +- current in amperes when explicitly converted; +- `I2` in the pyQSC covariant-current normalization. + +Angles and transform are dimensionless. Fourier axis coefficients have units +of length; `etabar` has inverse-length units; `B2c`, `B2s`, and `B20` have +T/m\(^2\). + +Enable JAX 64-bit arithmetic in the process environment before import: + +```bash +JAX_ENABLE_X64=true python calculation.py +``` + +The library does not make that global choice for callers. Reference parity, +small Maxwell residuals, and high-order spectral differentiation require +64-bit mode. In 32-bit mode, use problem-appropriate tolerances and do not +compare against the documented double-precision residuals. + +Resolution is part of the numerical model. Recompute accepted r2/r3 and +plasma-field candidates on doubled grids; inspect Fourier tails and solver +condition numbers rather than treating `nphi` as a cosmetic plotting choice. diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..9594177 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,36 @@ +"""Sphinx configuration for pyQSC_JAX.""" + +from pyqsc_jax import __version__ + +project = "pyQSC_JAX" +author = "UW Plasma" +copyright = "2026, UW Plasma" +version = __version__ +release = __version__ + +extensions = [ + "myst_parser", + "sphinx.ext.autodoc", + "sphinx.ext.doctest", + "sphinx.ext.intersphinx", + "sphinx.ext.mathjax", + "sphinx.ext.napoleon", + "sphinx_copybutton", + "sphinxcontrib.bibtex", +] + +source_suffix = { + ".md": "markdown", + ".rst": "restructuredtext", +} +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] +html_theme = "furo" +html_title = f"pyQSC_JAX {version}" +html_static_path = ["_static"] +intersphinx_mapping = { + "jax": ("https://docs.jax.dev/en/latest/", None), + "python": ("https://docs.python.org/3/", None), +} +bibtex_bibfiles = ["references.bib"] +bibtex_reference_style = "author_year" +myst_enable_extensions = ["amsmath", "colon_fence", "dollarmath"] diff --git a/docs/development/final_review.md b/docs/development/final_review.md new file mode 100644 index 0000000..aa262fd --- /dev/null +++ b/docs/development/final_review.md @@ -0,0 +1,126 @@ +# Final hostile review + +This review treats every sign, convention, compatibility promise, solver +status, dependency, and global-search claim as suspect until supported by an +equation, an independent comparison, or a failure-sensitive test. + +## Disposition + +No known pyQSC_JAX correctness defect remains after the local final gate. +The branch is suitable for draft-PR review, not for an immediate production +release. Production publication still requires remote CI, PR approval, +TestPyPI trusted-publishing verification, release-version metadata, and a +Zenodo archive/DOI. + +## Evidence + +| Gate | Result | +| --- | --- | +| Ruff lint and format | clean over `src`, `tests`, `examples`, `benchmarks` | +| Full test/coverage matrix | 296 passed, 4 opt-in integration skips; coverage run measured 98.75% lines and 95.898% branches | +| Documentation | HTML and doctest builders pass with warnings as errors | +| Examples | 14 tutorials + 10 publication scripts pass from clean cwd | +| Package metadata | wheel and sdist pass `twine check` | +| Clean artifacts | wheel and independently rebuilt sdist pass physics smokes | +| Dependencies | both environments pass `pip check`; no known advisories | +| High-resolution parity | QA/finite/QH r2 max array delta \(1.11\times10^{-11}\) at `nphi=121` | +| VMEC live integration | vacuum and finite-pressure, exactly zero-current equilibria pass; the latter has 0.02143% on-axis-iota error and \(9.97\times10^{-12}\) force residual | +| VMEX live integration | 2 vacuum/finite-beta current-main solves pass, including implicit pressure and boundary gradients | +| ESSOS clean integration | 13 field-jet/objective/example tests pass with pressure-only, exactly zero-current targets | +| Mutation: sigma sign | upstream-reference test fails with 200% relative transform error | +| Mutation: external Hessian `-`→`+` | STF symmetry test fails with order-unity error | + +The dependency auditor necessarily skips the unreleased `pyqsc-jax` +distribution itself because it has no PyPI record. It does audit the installed +JAX/SOLVAX dependency tree. + +## Convention audit + +- Axis Fourier packing is exactly `rc, rs, zc, zs`; constructor and legacy + `dofs` setter round-trip in the same order. +- Canonical sample axes lead. Gradient and Hessian are field-component-first: + `D[n,i,j] = dB_i/dx_j` and + `H[n,i,j,k] = d²B_i/(dx_j dx_k)`. +- Cylindrical vector differentiation includes basis-rotation connection + terms; fixed Cartesian arrays are not incorrectly treated as + field-periodic. +- Torsion follows Landreman--Sengupta and is documented as opposite the + original Garren--Boozer convention. +- Frame helicity, `sG`, and `spsi` enter + `iotaN = iota + N*nfp` consistently in forward, inverse, and continuation + systems. +- Current conversion uses `chi = sG*spsi`; a positive formal radius is + enforced before enclosed-current conversion or plasma-field matching. +- External jets are `total - plasma`. The 5- and 7-component packers project + and reconstruct the documented Cartesian STF bases. +- `verified_zero` is restricted to the known zero lower bound of the + nonnegative primary \(B_{20}\) residual. Nonzero searches remain + `best_found`. + +## Solver and AD audit + +- The sigma root is convergence-tested; the old fixed-five-step path is gone. +- The r2 coupled system reports residual and condition number. +- SOLVAX differentiates converged root/linear equations, not Newton or dense + solver iteration histories. +- Nonconverged candidates retain structured reports and are never silently + certified. +- VMEC and concrete VMEX boundary conversion refuse nonfinite, nonpositive, + or cylindrically unconverged surfaces; traced invalid VMEX candidates become + nonfinite instead of silently changing the requested boundary. +- JIT, VMAP, JVP, VJP, and finite-difference tests cover first order, r2, + inverse branches, optimization, and plasma jets. +- Discrete basin selection is not advertised as differentiable; smooth + candidate residuals are. + +## Architecture and dependency audit + +- Physics lives in small immutable pytree/dataclass records and pure + functions; the mutable class is only a compatibility adapter. +- The package has no mutable global physics state and performs no import-time + JAX configuration. +- Runtime dependencies are only unpinned `jax` and `solvax`. Equinox remains + an indirect SOLVAX implementation dependency rather than a public modeling + requirement. +- The unused `pyevtk` plotting extra found during review was removed. +- No `setup.py`, separate `jaxlib` declaration, generic `utils.py`, or ESSOS + dependency remains. +- The wheel contains only package code and required license/NOTICE metadata; + the sdist contains sources, tests, docs, examples, benchmark report, and + citation files. + +## Documentation/API findings fixed during review + +- Replaced an unsafe tag-only production publish path with reviewed-release, + full-verification, main-ancestry, OIDC TestPyPI/PyPI jobs. +- Added missing Codecov, docs-example, clean ESSOS, and benchmark artifact + workflow coverage. +- Removed stale documentation names `V1r` and `B2cQI`, which are not public r2 + fields. +- Corrected the publication surface plot so the magnetic-axis overlay spans + the full torus. +- Kept the DOI badge explicitly pending instead of fabricating a Zenodo + identifier. + +## Exact remaining limitations + +- Near-axis and plasma-channel results are asymptotic; field/Hessian remainder + values are scales, not rigorous bounds. +- The plasma Hessian is validated through its independently derived interior + contact terms and Maxwell identities; differentiating a singular filament + quadrature is not a valid independent reference. +- Frenet geometry fails at zero curvature. +- `etabar` continuation has a fixed step and does not yet provide a generic + multiparameter continuation engine. +- Magnetic shear supports only the documented standard-MHS sign + specialization. +- A finite multistart budget cannot certify a nonzero global minimum. +- VMEX supplies differentiable fixed-boundary finite-radius equilibrium + quantities, but fast-particle confinement, free-boundary equilibrium, and + detailed coil engineering remain external validations. +- The ESSOS base branch retains unrelated collection and documentation + failures described in ESSOS PR #46; the new field-jet slice is independently + green. +- TestPyPI upload, PyPI trusted-publisher configuration, PR approval, and + Zenodo DOI creation require maintainer/external service actions and are not + represented as complete. diff --git a/docs/development/physics-traceability.md b/docs/development/physics-traceability.md new file mode 100644 index 0000000..6126582 --- /dev/null +++ b/docs/development/physics-traceability.md @@ -0,0 +1,34 @@ +# Physics traceability + +Every implemented physics block must point to a primary derivation, a code +location, and at least one independent validation. Equation numbers below refer +to the cited source, not to a duplicated derivation in this repository. + +| Physics block | Primary source | Reference implementation | Required validation | Status | +| --- | --- | --- | --- | --- | +| General Fourier axis and Frenet frame | Landreman & Sengupta (2018), arXiv:1809.10233 | `Axis`, `evaluate_axis`, `compute_axis_geometry`; pyQSC `init_axis.py` | analytic circle; asymmetric Fortran curvature/torsion/varphi; frame identities; JIT/VMAP/grad | implemented and documented | +| First-order QS and sigma equation | Landreman & Sengupta (2019), arXiv:1908.10253 | `first_order.solve_sigma`; pyQSC `solve_sigma_equation.py` | residual, convergence report, QA/QH/finite-current pyQSC values, JIT/VMAP/JVP/VJP, finite differences | implemented and documented | +| Target-transform inverse solve and continuation | implicit-function reformulation of the first-order sigma equation | `inverse.solve_target_iota`; `continuation.continue_etabar_branch` | `etabar` and `I2` round trips, sign branches, distinct local basins, fold crossing, independent forward correction, r2 propagation, JIT/JVP/finite differences | branch-local inverse and fixed-step `etabar` pseudo-arclength implemented | +| Nonconstant \(B_{20}\) and affine \(B_{2c}\) elimination | Landreman & Sengupta (2019), complete r2 system | `optimize.b20_diagnostics`; `optimize.optimize_B2c` | direct norm definitions, affine reconstruction, stationary optimum, circular degeneracy, spectral tail, doubled-grid verification, JIT/JVP/finite differences | implemented and documented | +| Branch-aware axis search | nonnegative \(B_{20}\) residual plus the package's traced r2 equations | `axis_optimization.search_axis`; `axis_optimization.continue_axis_search` | deterministic batched exploration, bounded LM steps, accepted/rejected damping, criteria failure, basin clustering, certified zero, nonzero best-found semantics, staged modes | implemented and documented | +| Field gradient and `L_grad_B` | Landreman (2021), arXiv:2012.00865, Eq. 3.12 | `first_order.first_order_solution`; pyQSC `grad_B_tensor.py` | upstream arrays, vacuum symmetry/divergence, resolution convergence | implemented and documented | +| Complete second-order coefficients | Landreman & Sengupta (2019), arXiv:1908.10253 | `second_order.solve_second_order`; pyQSC `calculate_r2.py`; manuscript Eqs. 36–54 | four independent equation residuals; vacuum QA, finite-pressure/current, QH upstream arrays; JIT/JVP/VJP/finite differences | implemented and documented | +| Total field Hessian | manuscript Eqs. 55–83 | `field.total_field_jet`; pyQSC generated tensor as comparison only | regular-coordinate chain rule, JIT/JVP, QA and finite-current upstream arrays, full vacuum symmetry, divergence and derivative-of-divergence | implemented and documented | +| Magnetic well and Mercier terms | Landreman (2021), arXiv:2012.00865 | `diagnostics.mercier_diagnostics`; pyQSC `mercier.py` | vacuum and finite-pressure/current upstream values | implemented and documented | +| Third-order flux constraint | Landreman & Sengupta (2019) | `third_order.solve_third_order`; pyQSC `calculate_r3.py` | two independent flux-constraint forms; QA, finite-pressure/current, and QH upstream arrays; r3 boundary data; JIT/JVP/finite differences | implemented and documented | +| Magnetic shear | Rodríguez, Sengupta & Bhattacharjee (2022), Appendix F | `shear.solve_magnetic_shear`; pyQSC `calculate_r3.py` | three upstream paper cases, asymmetric/secular integration, resolution convergence, `B31c`, JIT/JVP/VJP and full-solve AD | implemented and documented for standard MHS with positive signs | +| Singular radius and scale lengths | Landreman (2021), arXiv:2012.00865 | `singularity.singularity_diagnostics`; pyQSC `r_singularity.py` as comparison | direct regular-map determinant, QA/QH/finite-current upstream radii and arrays, residual, angular convergence, JIT/JVP/finite differences | implemented and documented | +| Good-stellarator criteria | Curvo, Ferreira & Jorge (2025), Table 3 | `criteria.Criteria.from_curvo_2025` | exact normalized thresholds, strict/inclusive comparisons, signed margins, scaling, JIT/JVP/finite differences | implemented and documented | +| Plasma current source | manuscript Eqs. 88–97 | `plasma.plasma_current_source`; `plasma.evaluate_weighted_current` | exact \(I_2\)/ampere normalization, regular radial power, pressure-only and vacuum limits, batch/JIT evaluation | implemented and documented | +| Plasma on-axis field | manuscript Eqs. 118 and 133–138 | `plasma.regularized_axis_integral`; `plasma.plasma_field_on_axis` | circular finite part/core/logarithm, full-torus construction, matching-length cancellation, angular/toroidal convergence, vacuum/pressure limits, JIT/JVP, resolved volume Biot–Savart \(a^4|\log a|\) scaling | implemented and documented | +| Plasma gradient | manuscript Eqs. 142–160 | `plasma.elliptical_channel_gradient`; `plasma.plasma_gradient_on_axis` | straight circular and sheared elliptical channels, frame covariance, Ampère antisymmetry, divergence, external symmetry/trace, vacuum reduction, 5-component pack/unpack, JIT/JVP/finite differences | implemented and documented | +| Plasma Hessian | manuscript Eqs. 194 and 210–219 | `plasma.plasma_hessian_on_axis` | circular finite-conductor contact-term limit, QA/QH oriented ellipses, commuting plasma derivatives, spectral convergence, JIT/JVP/finite differences | implemented and documented | +| External vacuum jet | manuscript Eqs. 218–224 | nested `PlasmaHessianData`; rank-2/rank-3 STF packers | 3+5+7 representation, full Hessian symmetry, trace-free gradient/Hessian, vacuum reduction, asymptotic error metadata, ESSOS coil/axis finite-difference gradients | implemented and documented; normalized stage-two and single-stage ESSOS objectives are in draft ESSOS PR #46 | +| VMEC fixed-boundary conversion | near-axis surface expansion plus VMEC Fourier boundary convention | `vmec.uniform_cylindrical_surface`; `vmec.vmec_boundary`; upstream pyQSC `to_vmec.py` as an independent comparison | vectorized/legacy surface agreement, Fourier reconstruction, asymmetric coefficients, deterministic input, hard failure on unconverged cylindrical-angle inversion, warm timing, frozen and live vacuum/finite-pressure zero-current VMEC on-axis iota and force residuals, radius convergence | implemented and documented | +| VMEX radial equilibrium | VMEX fixed-boundary ideal-MHD fixed point and public implicit adjoint | `vmex.to_vmex_problem`; `vmex.vmex_radial_quantities`; VMEX 0.3.0 commit `2a40d756` | vacuum and finite-beta live solves, radial iota sign audit, QS profile, magnetic well, boundary/profile gradients, current-main CI | implemented and documented | + +## Source policy + +pyQSC is BSD-2-Clause licensed. Adapted code, formulas expressed as code, and +derived reference data must preserve attribution. The attached manuscript is +not committed. Its locally audited SHA-256 is recorded in a gitignored note. diff --git a/docs/development/refactor-baseline.md b/docs/development/refactor-baseline.md new file mode 100644 index 0000000..4f55ee2 --- /dev/null +++ b/docs/development/refactor-baseline.md @@ -0,0 +1,144 @@ +# Refactor baseline + +Date: 2026-07-29 + +This audit establishes the starting point for the complete pyQSC_JAX refactor. +It is descriptive, not an endorsement of the current architecture. + +## Repository and permission status + +| Repository | Audited revision | Access | Local status | +| --- | --- | --- | --- | +| `uwplasma/pyQSC_JAX` | `8b5d8ea218a624071765bfbf62d8d8e0c7acb90c` | admin/push | clean before audit; feature branch created | +| `landreman/pyQSC` | `cd75359ea47548d5db7ccb458c100085c04ba1bc` | read-only reference | clean | +| `uwplasma/ESSOS` | `932997f198964ccfffea898e25acd73600adbf3a` | admin/push | existing worktree has unrelated generated output | +| `uwplasma/SOLVAX` | `bf219c79a40a80b2c3115b749bf670b9f5de3cb5` | admin/push | existing worktree has unrelated changes | + +GitHub authentication is active for `rogeriojorge`. Any later ESSOS or SOLVAX +work must use an isolated worktree so the pre-existing local changes remain +untouched. + +## Current pyQSC_JAX inventory + +The package consists of an empty `pyqsc_jax/__init__.py`, a 505-line +`pyqsc_jax/near_axis.py`, `setup.py`, a title-only README, and no tests or CI. +The monolithic mutable `near_axis` class currently combines: + +- stellarator-symmetric axis Fourier evaluation using only `rc` and `zs`; +- Frenet geometry and topology; +- the first-order sigma equation; +- field and gradient evaluation; +- coordinate conversion and boundary generation; +- mutable optimization degrees of freedom. + +`order`, `B2c`, and `p2` are accepted but do not activate second- or third-order +physics. The sigma equation is advanced by exactly five Newton steps with no +convergence criterion or report. Methods are JIT-compiled with mutable `self` +as a static argument. Package metadata constrains both `jax` and `jaxlib`. + +The unmerged `origin/en/essos_bridge_fix` branch adds plotting and VTK export, +but also imports ESSOS from pyQSC_JAX and commits generated artifacts. Required +behavior will be reimplemented without reversing the dependency direction. + +## Numerical baseline + +For + +```python +near_axis( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + etabar=-0.9, + nfp=3, + nphi=31, +) +``` + +with JAX 64-bit arithmetic, current pyQSC_JAX and upstream pyQSC agree to +floating-point precision: + +| Quantity | Baseline | +| --- | ---: | +| `iota` | 0.41830690943386617 | +| `axis_length` | 6.340238817434161 | +| minimum curvature | 0.5956163516982903 | +| maximum curvature | 1.3060121594235536 | +| minimum torsion | -2.272197039914583 | +| maximum torsion | 0.6057564303422814 | +| maximum absolute sigma | 0.9920706836908618 | +| `B_axis.shape` | `(3, 31)` | +| `grad_B_axis.shape` | `(3, 3, 31)` | + +The audit environment used unpinned JAX/JAXLIB 0.11.0, NumPy 2.5.1, +SciPy 1.18.0, and Python 3.12. + +## Confirmed defect + +Assigning `field.dofs = field.dofs` unpacks the normal and binormal cylindrical +components in a different order than the constructor. The no-op round trip +changes the frame by up to: + +- normal: `0.5460213389358983`; +- binormal: `1.1249025277689044`. + +The strict expected-failure regression test must become a passing test when the +compatibility adapter is corrected. + +## Compatibility tiers + +### Tier 1: ESSOS, mandatory throughout + +ESSOS PR `uwplasma/ESSOS#34` established the import: + +```python +from pyqsc_jax.near_axis import near_axis +``` + +ESSOS constructs the class from `rc`, `zs`, `etabar`, `nfp`, and `nphi`. +Production objectives currently consume `R0`, `Z0`, `phi`, `B_axis`, +`grad_B_axis`, and `iota`; optimization additionally consumes `rc`, `zs`, +`etabar`, `x`, `dofs`, `B0`, `sigma0`, `I2`, `nphi`, `spsi`, `sG`, and `nfp`. +Examples also rely on boundary conversion and plotting. + +### Tier 2: upstream pyQSC semantic parity, final acceptance + +DESC, SIMSOPT, and broader pyQSC workflows use a larger object protocol, +including general `rc`, `rs`, `zc`, `zs`, `Bbar`, `lasym`, pressure/current +coefficients, interpolation splines, second- and third-order boundary data, and +named configurations. This tier is staged after the immutable first-order core +but remains part of the final definition of done. + +## Upstream test audit + +At pyQSC commit `cd75359`, 31 tests passed in 22.80 seconds during the initial +audit when upstream `test_to_vmec.py` was excluded because its optional +MPI/VMEC runtime was absent from that audit environment. The refactored +package now has its own frozen-wout equilibrium regression and an opt-in local +VMEC rerun; see {doc}`../validation/vmec`. The audited upstream tests cover: + +- Newton convergence, Fourier differentiation, interpolation, and utilities; +- axis geometry, helicity, sigma, iota, field gradient, and scale lengths; +- comparison to independent Fortran netCDF outputs; +- complete second-order coefficients, `B20`, Mercier terms, and singular radius; +- third-order coefficients, shear, and boundary Fourier conversion; +- named paper configurations and independent curvature/torsion checks. + +Reference data imported later must be reduced to small arrays, retain BSD-2-Clause +attribution, record this upstream commit and checksums, and never download during +tests. + +## Initial risks and gates + +1. Frenet coordinates require nonvanishing axis curvature; invalid axes need an + explicit diagnostic rather than silent NaNs. +2. The target-iota inverse can cross folds where + `d(iota)/d(etabar) = 0`; scalar inversion alone is not globally valid. +3. The second-order operator can be ill-conditioned near `iota_N = 0`. +4. A small collocation residual does not establish a small dense-grid `B20` + variation. +5. The manuscript plasma Hessian is gated on independent volume-current + Biot–Savart validation and Maxwell identities. +6. “Surface-free” still requires a formal radius or equivalent flux/current + normalization. +7. Nonzero optimization results are “best found,” never proofs of a global + minimum. diff --git a/docs/development/refactor_status.md b/docs/development/refactor_status.md new file mode 100644 index 0000000..64a8ae0 --- /dev/null +++ b/docs/development/refactor_status.md @@ -0,0 +1,273 @@ +# Refactor status + +This log records green phase boundaries on +`refactor/pyqsc-jax-complete`. “Not benchmarked” means no performance claim +was made for that phase. + +## Phase 1 — audit and compatibility baseline + +- Commit: `3a87ba6`. +- Files: baseline report, architecture/solver ADRs, local reference policy, + compatibility regression tests. +- Verification: legacy imports, shapes, mutability, boundary behavior, and + upstream reference extraction. +- Coverage/residuals: baseline only; inherited five-step sigma behavior + recorded rather than accepted. +- Performance/API: not benchmarked; historical API frozen. +- Risk/next: mutable monolith and inconsistent DOF unpacking; build immutable + axis/geometry core. + +## Phase 2 — packaging and quality scaffold + +- Commits: `5895008`, `969d616`. +- Files: `pyproject.toml`, `src/` layout, CI, docs configuration, license and + citation metadata. +- Verification: clean editable install, wheel/sdist smoke, Ruff, x64 CI. +- Coverage/residuals: 95% line/branch gates established. +- Performance/API: no import-time JAX policy; `jaxlib` removed from metadata. +- Risk/next: physics still incomplete; establish spectral geometry and a + converged first-order solve. + +## Phase 3 — axis, geometry, and first order + +- Commits: `28955f8`, `16249fe`. +- Files: immutable axis/geometry/spectral modules, damped sigma solver, + first-order models, thin compatibility adapter. +- Verification: analytic circle, asymmetric reference geometry, QA/QH and + finite-current parity, JIT/VMAP/JVP/VJP/finite differences. +- Coverage/residuals: QA sigma residual \(1.2\times10^{-14}\); coverage above + the 95% gate. +- Performance/API: explicit pytrees replace static mutable `self`; no + benchmark claim. +- Risk/next: r2 outputs unavailable; implement the complete coupled system. + +## Phase 4 — complete r2 and diagnostics + +- Commits: `3978a6b`, `5e67237`, `0086d77`. +- Files: second-order solve, total regular-coordinate Hessian, Mercier and + singular-radius diagnostics. +- Verification: four independent r2 residual equations; vacuum QA, + finite-pressure/current, and QH upstream arrays; Maxwell identities; + singular-map roots; JIT and AD checks. +- Coverage/residuals: finite-current linear residual + \(1.6\times10^{-13}\); field reconstruction at roundoff. +- Performance/API: one implicit dense linear solve; not benchmarked. +- Risk/next: third-order boundary/shear parity; implement as a separate green + milestone. + +## Phase 5 — third order and shear + +- Commits: `99e81d5`, `935e1f8`. +- Files: r3 flux constraint, untwisted boundary harmonics, standard-MHS + magnetic shear. +- Verification: two flux-constraint forms; QA/QH/finite-current upstream + coefficients and boundaries; three shear cases; resolution and AD. +- Coverage/residuals: consistency residuals meet documented test tolerances; + coverage above 95%. +- Performance/API: shear remains explicit because `B31c` is independent. +- Risk/next: shear sign specialization is documented; add inverse branches + and folds. + +## Phase 6 — inverse solves and optimization + +- Commits: `7e87d04`, `abcbd0a`, `5169173`, `cf3f6da`, `6fab74b`. +- Files: target-transform inverse, full-state continuation, \(B_{20}\) + diagnostics, analytic `B2c`, criteria, bounded multistart axis search. +- Verification: branch round trips, two inverse basins, fold traversal, + affine reconstruction \(3.3\times10^{-15}\), exact scalar stationarity, + criteria scaling, basin/status/resolution tests. +- Coverage/residuals: complete suite remained above 95% branch coverage. +- Performance/API: deterministic batched coarse search and JAX Jacobians; + absolute timing not benchmarked. +- Risk/next: nonzero results explicitly lack global certificates; implement + the surface-free current and field. + +## Phase 7 — plasma source and external 3+5+7 jet + +- Commits: `d46b2bc`, `f173de7`, `f7d15b3`, `c7277fc`. +- Files: positive-volume current, full-torus matched field, local gradient and + Hessian, STF pack/unpack, asymptotic metadata. +- Verification: current conversions, straight/circular/elliptical channels, + resolved-volume Biot--Savart field scaling, Ampère/divergence, QA/QH + external STF identities, vacuum reduction, convergence, JIT/JVP/FD. +- Coverage/residuals: 203 tests passed; 98.90% branch-aware coverage; + `plasma.py` 100%. +- Performance/API: surface-free calls accept explicit formal radius; not + benchmarked. +- Risk/next: Hessian contact terms rely on the independently derived interior + potential; integrate only the external target into ESSOS. + +## Phase 8 — ESSOS integration + +- ESSOS commits: `8a1ce7f`, `d20c23d`; draft PR + [#46](https://github.com/uwplasma/ESSOS/pull/46). +- Files: external target/residual models, coil Hessian interface, stage-two + and single-stage examples, cross-repository tests. +- Verification: 24 focused tests; actual coil and near-axis gradients versus + finite differences; vacuum reduction; example objectives 10.74→4.09 and + 10.75→4.52. +- Coverage/performance: focused integration only; example budgets are smoke + demonstrations, not device-quality searches. +- API: dependency remains `ESSOS -> pyQSC_JAX`. +- Risk/next: inherited ESSOS base-suite collection/docs failures are disclosed + in the PR; complete pyQSC_JAX examples and documentation. + +## Phase 9 — examples and publication figures + +- Commit: `a8eb51a`. +- Files: named configurations, lazy plotting API, 12 tutorials, 5 + deterministic publication scripts, clean-cwd execution tests. +- Verification: 23 focused tests; every script produced its figure; all five + publication plots visually inspected; Ruff clean. +- Coverage/residuals: tutorial AD relative error \(9.4\times10^{-10}\); + tutorial continuation crosses a detected fold. +- Performance/API: full clean-directory example suite took 152 s on the + development Mac; this includes separate-process JAX compilation and is not a + package benchmark. +- Risk/next: finish docs, release hardening, clean artifact validation, and + hostile review. + +## Phase 10 — documentation and release hardening + +- Commit: `69b05da`. +- Files: complete docs hierarchy, tutorials via `literalinclude`, migration, + validation reports, release checklist, README, and compact reproducible + figures. +- Verification: warnings-as-errors HTML and doctest builds; 23 focused + configuration/plot/example tests; figure visual inspection. +- Performance/API: named configurations and lazy plotting helpers added; no + new physics formula. +- Risk/next: harden coverage, packaging, publishing, compatibility, and + benchmark workflows before final review. + +## Phase 11 — release hardening + +- Commit: `686be1e`. +- Files: synchronized benchmark and raw report, CI coverage/Codecov, docs + example smoke, ESSOS compatibility job, TestPyPI/production OIDC workflow, + clean wheel/sdist verification. +- Verification: 226 tests pass with 98.64% combined line/branch coverage; + HTML/doctest warning-free; wheel and sdist install independently; `pip + check` and dependency audit clean; 12 ESSOS field-jet tests pass in a fresh + environment. +- Numerical checks: direct `nphi=121` QA/finite/QH comparison with audited + pyQSC has maximum r2 array difference \(1.11\times10^{-11}\). +- Performance: 61-point first-order solve 0.198 s cold and 144 µs warm on the + documented Apple M4 run; raw samples and caveats are committed. +- API: removed unused `pyevtk`; core dependencies remain only `jax` and + `solvax`. +- Risk/next: remote CI, TestPyPI OIDC, PR approval, and Zenodo are external + release gates; perform hostile final review before requesting merge. + +## Phase 12 — final review + +- Commit: `fb7ddac`. +- Evidence: convention/API/dependency audits complete; two representative + physics mutations are killed by targeted tests; exact limitations and + external release gates are recorded in `final_review.md`. +- Verification: final local lint, docs, clean-state, artifact, upstream, and + ESSOS checks pass. The latest remote CI result is reported in the PR rather + than hard-coded into this versioned status page. + +## Phase 13 — high-signal examples and VMEC equilibrium validation + +- Files: vectorized VMEC exporter, deterministic INDATA writer, frozen VMEC + 9.0 equilibrium, B20/plasma-dominant named cases, optimizer/export + benchmarks, publication figures, and expanded validation docs. +- Verification: 255 tests pass with 98.74% combined line/branch coverage; + `vmec.py` has 100% line/branch coverage; warnings-as-errors documentation + builds; both frozen and live local-VMEC equilibrium checks pass. +- Numerical checks: optimized QA weighted B20 residual + \(1.5901\times10^{-6}\) at `nphi=121` and stable through `nphi=481`; + plasma-field fraction 33.26%; VMEC on-axis-iota error 0.0565% at radius + 0.0025 with force residuals below \(7.6\times10^{-11}\). +- Performance: synchronized VMEC conversion is 0.476 s including compilation + and 5.20 ms warm at the documented high-resolution boundary setting. +- Risk/next: on-axis-iota agreement is asymptotic and degrades at larger + export radius as documented; remote CI and review remain external gates. + +## Phase 14 — VMEX radial equilibria and visual design audit + +- Files: optional VMEX bridge, live current-main compatibility job, radial + equilibrium tutorial and publication figure, full-torus 3D plotting + helpers, four-case gallery, and angle-dependent plasma configuration. +- Verification: vacuum and finite-beta VMEX solves; radial \(\iota\), QS, + magnetic-well, energy, and boundary-gradient checks; traceable + `parameters_for` differentiation; focused plotting and plasma regressions. +- Numerical checks: low-resolution vacuum VMEX/near-axis on-axis-\(\iota\) + difference 0.249%; finite-beta magnetic well + \(-4.2297\times10^{-4}\); angle-dependent minimum plasma-field fraction + 32.70%; optimized \(B_{20}\) surface shown 3.3 times inside its singular + radius. +- Performance: the VMEX smoke benchmark separates the converged forward solve + from the magnetic-well value-and-implicit-gradient cost; it is documented + as an integration timing, not a production throughput claim. +- Risk/next: VMEX quasisymmetry diagnostics currently require stellarator + symmetry and the bridge covers fixed-boundary implicit differentiation, not + a reconverged free-boundary NESTOR adjoint. Remote CI and review remain + external gates. + +## Phase 15 — screened database showcase and constrained \(B_{20}\) + +- Files: traceable configurations from Wisconsin database IDs 3, 57409, + 107579, and the one-period QA ID 139524; an eight-mode constrained + refinement of ID 57409; direct-Frenet 3D surface plotting; topology, + convergence, measured-performance, and optimizer-comparison README figures. +- Verification: every showcased design independently passes the Curvo profile. + The QA lead passes with helicity zero and + \(\lvert\iota\rvert\ge0.3\); the other displayed cases pass with + \(\lvert\iota\rvert\ge0.4\). Database source IDs and URLs are frozen in + configuration and figure metadata, and every displayed surface has a + smoothness regression. +- Numerical checks: the derived ID-57409 case has + \(\lvert\iota\rvert=2.964\), \(r_\mathrm{sing}=0.249\) m, and weighted + \(B_{20}\) residual \(1.2743\times10^{-10}\) at `nphi=121`, stable through + `nphi=481`. The finite-beta showcase is database ID 52521 with exactly + \(I_2=0\), finite \(p_2\), RMS axis torsion \(0.979\ \mathrm{m}^{-1}\), + and 14.04% relative angular variation of the pressure-driven plasma-field + norm while passing the same screen. +- QA showcase checks: ID 139524 has helicity zero, + \(\lvert\iota\rvert=0.355\), RMS axis torsion + \(1.199\ \mathrm{m}^{-1}\), and \(r_\mathrm{sing}=0.097\) m. The complete + 287-test run passes with 98.79% combined line/branch coverage; all 24 public + scripts and strict HTML/doctest documentation builds pass. +- Performance evidence: synchronized Apple M4 measurements show 0.144 ms warm + first-order latency at `nphi=61`, a 1,373-fold cold-to-warm ratio, 0.191 ms + for a JVP, and 0.479 ms for an eight-case VMAP batch. +- Limitation: at formal radius 0.15 m the pressure-only plasma contribution is + approximately 0.19% of the total field. The former 30% showcase relied on + finite \(I_2\) and was removed because its nearly planar geometry presented + as a tokamak rather than a representative stellarator. + +## Phase 16 — final VMEC, VMEX, and ESSOS integration + +- Commits: pyQSC_JAX `060aef3`; ESSOS `da66774`. +- Files: convergence-safe VMEC/VMEX boundary conversion, a live + finite-pressure and exactly zero-current VMEC regression, explicit opt-in + VMEX examples, and pressure-only stage-two and single-stage ESSOS examples. +- VMEC verification: the database QA ID 139524 case uses finite pressure, + \(I_2=0\), and RMS axis torsion \(1.189\ {\rm m}^{-1}\). At export radius + 0.0015 m, VMEC converges to force residual \(9.97\times10^{-12}\) and its + signed on-axis transform differs from pyQSC_JAX by 0.02143%. Boundary + export now refuses to write when the cylindrical-angle inversion is + nonfinite or unconverged. +- VMEX verification: two live current-main tests pass against VMEX 0.3.0 + commit `2a40d756`, covering vacuum and finite-beta/zero-current radial + transform, quasisymmetry, magnetic well, and implicit pressure/boundary + gradients. Both opt-in example scripts also complete against that commit. +- ESSOS verification: 13 focused tests pass, including the two clean-process + optimization examples. Both use the nonplanar, finite-pressure database + stellarator ID 52521 with exactly \(I_2=0\), and every field/gradient/Hessian + residual block decreases in both stage-two and single-stage runs. +- Publication verification: all ten scripts were regenerated from + `060aef3`. Physics figures record that exact source commit; performance + figures retain the exact commit of their synchronized raw benchmark report. + The committed README PNGs were byte-stable and visually re-inspected. +- Complete pyQSC_JAX evidence: the branch-coverage run completed 294 tests + before two VMEX subprocesses reached their former 180 s instrumentation + timeout; after making the documented VMEX opt-in explicit, those exact two + cases pass. The combined complete matrix is therefore 296 passes and four + intentional live-integration skips, with 98.75% line and 95.898% branch + coverage in the coverage run. +- Risk/next: TestPyPI/PyPI trusted publishing, Zenodo creation, maintainer + review, and merge remain external release gates. diff --git a/docs/getting_started/choosing_a_model.md b/docs/getting_started/choosing_a_model.md new file mode 100644 index 0000000..cd189da --- /dev/null +++ b/docs/getting_started/choosing_a_model.md @@ -0,0 +1,17 @@ +# Choosing a model + +| Goal | Order/API | Required checks | +| --- | --- | --- | +| Transform, elongation, \(L_{\nabla B}\) | `order="r1"` | root convergence, geometry validity | +| \(B_{20}\), Mercier, Hessian, singular radius | `order="r2"` | root and linear reports, resolution | +| Flux-consistent boundary and magnetic shear | `order="r3"` plus `solve_magnetic_shear` | r3 residuals and applicable sign restrictions | +| Prescribed transform | `solve_for="etabar"` or `"I2"` | local response and fold flag | +| Surface-free plasma/external target | `plasma_hessian_on_axis` | positive formal radius and asymptotic metadata | +| Radial iota, QS, and magnetic well | `to_vmex_problem` | VMEX convergence, implicit-gradient and resolution checks | +| Legacy ESSOS consumer | `near_axis` adapter | compatibility orientation and mutable-DOF contract | + +Use the lowest order that supplies the observable being optimized. Higher +order adds physical information but also tighter geometry, conditioning, and +resolution requirements. A near-axis model is asymptotic; it does not by +itself prove finite-radius equilibrium, coil feasibility, or particle +confinement. diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md new file mode 100644 index 0000000..813cdfb --- /dev/null +++ b/docs/getting_started/installation.md @@ -0,0 +1,36 @@ +# Installation + +pyQSC_JAX requires Python 3.12 or newer. + +```bash +python -m pip install pyqsc-jax +``` + +Install optional plotting, VMEX equilibrium, or development tools explicitly: + +```bash +python -m pip install 'pyqsc-jax[plot]' +python -m pip install 'pyqsc-jax[plot,vmex]' +git clone https://github.com/uwplasma/pyQSC_JAX.git +cd pyQSC_JAX +python -m pip install -e '.[dev,docs,plot,vmex]' +``` + +The `vmex` extra enables differentiable fixed-boundary radial quantities. To +test the current upstream development head directly: + +```bash +python -m pip install 'git+https://github.com/uwplasma/vmex.git' +``` + +The project declares `jax`, not `jaxlib`, and never changes JAX configuration +at import time. Choose the CPU, CUDA, or other backend using the +[official JAX installation guide](https://docs.jax.dev/en/latest/installation.html). +Production calculations should enable 64-bit arithmetic before importing JAX: + +```bash +export JAX_ENABLE_X64=true +``` + +Dependencies are intentionally unpinned in package metadata. Release reports +record the concrete environment used for validation. diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md new file mode 100644 index 0000000..c2e22fb --- /dev/null +++ b/docs/getting_started/quickstart.md @@ -0,0 +1,34 @@ +# Quickstart + +The canonical entry point returns an immutable JAX pytree: + +```python +import pyqsc_jax as qsc + +solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, + order="r1", +) + +assert solution.root_report.converged +print("iota:", solution.iota) +print("residual:", solution.root_report.residual_norm) +print("axis length:", solution.axis_length) +print("minimum L_grad_B:", solution.L_grad_B.min()) +``` + +`rc` and `zs` are Fourier coefficients of cylindrical axis radius and height, +`nfp` is the field-period count, and `etabar` sets the leading +field-strength variation. Start with [choosing a model](choosing_a_model.md) +before requesting second- or third-order outputs. + +The complete executable version is included directly from the repository: + +```{literalinclude} ../../examples/01_first_order_qa.py +:language: python +:linenos: +``` diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..21aaf5a --- /dev/null +++ b/docs/index.md @@ -0,0 +1,157 @@ +# pyQSC_JAX + +pyQSC_JAX constructs differentiable near-axis stellarators in JAX and +separates their on-axis total field jet into plasma-generated and external +vacuum targets without constructing a finite-radius surface. + +```{toctree} +:maxdepth: 2 +:caption: Getting started + +getting_started/installation +getting_started/quickstart +getting_started/choosing_a_model +``` + +```{toctree} +:maxdepth: 2 +:caption: Concepts + +concepts/coordinates_and_conventions +concepts/magnetic_axis +concepts/inputs_and_outputs +concepts/precision_and_units +``` + +```{toctree} +:maxdepth: 2 +:caption: Theory + +theory/first_order +theory/second_order +theory/third_order +theory/field_tensors +theory/diagnostics +theory/axis_optimization +theory/plasma_coil_separation +``` + +```{toctree} +:maxdepth: 2 +:caption: Tutorials + +tutorials/first_order_qa +tutorials/first_order_qh +tutorials/second_order_finite_beta +tutorials/target_iota +tutorials/optimize_axis +tutorials/vmec_export +tutorials/vmex_equilibrium +tutorials/vacuum_coils +tutorials/finite_beta_coils +``` + +```{toctree} +:maxdepth: 2 +:caption: Advanced + +advanced/continuation +advanced/global_search +advanced/custom_criteria +advanced/autodiff +advanced/performance +advanced/limitations +``` + +```{toctree} +:maxdepth: 2 +:caption: Validation and migration + +validation/pyqsc_parity +validation/literature_cases +validation/plasma_field +validation/vmec +validation/vmex_interface +validation/essos +migration +``` + +```{toctree} +:maxdepth: 2 +:caption: API + +api/axis +api/configurations +api/first-order +api/second-order +api/third-order +api/field-jet +api/plasma +api/continuation +api/optimization +api/axis-optimization +api/criteria +api/plotting +api/vmec +api/vmex +``` + +```{toctree} +:maxdepth: 2 +:caption: Derivation details + +theory/coordinates-and-conventions +theory/first-order +theory/second-order +theory/third-order +theory/field-jet +theory/b20-optimization +theory/criteria +theory/global-search +theory/inverse-solves +theory/plasma-current +theory/plasma-field +``` + +```{toctree} +:maxdepth: 2 +:caption: Development + +development/refactor-baseline +development/physics-traceability +development/refactor_status +development/final_review +adr/ADR-core-architecture +adr/ADR-solver-stack +``` + +```{toctree} +:maxdepth: 1 +:caption: Project + +changelog +release_checklist +``` + +## Status + +The canonical API is immutable and JAX-transformable. The legacy +`pyqsc_jax.near_axis.near_axis` import remains available as a thin ESSOS +adapter. Each nonlinear and linear solve returns convergence and conditioning +evidence; callers should never accept a result solely because an object was +returned. + +The physics traceability table connects each equation block to its primary +source, implementation symbol, and independent tests. The surface-free field +split always requires a formal radius or equivalent current/flux +normalization: surface-free does not mean radius-free. + +## References + +The implementation begins with the near-axis construction of Garren and +Boozer and the direct cylindrical-coordinate formulations of Landreman and +collaborators +{cite}`garren1991existence,garren1991magnetic,landreman2018direct,landreman2019highorder`. + +```{bibliography} +``` diff --git a/docs/migration.md b/docs/migration.md new file mode 100644 index 0000000..578457d --- /dev/null +++ b/docs/migration.md @@ -0,0 +1,57 @@ +# Migration from pyQSC and the legacy adapter + +## Canonical immutable API + +Replace: + +```python +from qsc import Qsc + +stel = Qsc(rc=rc, zs=zs, nfp=nfp, etabar=etabar) +``` + +with: + +```python +import pyqsc_jax as qsc + +solution = qsc.Qsc(rc=rc, zs=zs, nfp=nfp, etabar=etabar) +``` + +Canonical sampled arrays put `nphi` first: `B_axis` is `(nphi, 3)`, +`grad_B_axis` is `(nphi, 3, 3)`, and `grad_grad_B_axis` is +`(nphi, 3, 3, 3)`. Results are frozen. Rebuild with changed explicit inputs +instead of assigning fields. + +Constructor `order` accepts `"r1"`, `"r2"`, or `"r3"`. Second-order fields +raise an attribute error on an r1 result rather than returning stale or +partially initialized values. Check `root_report`, `linear_report`, and +resolution diagnostics explicitly. + +## Existing ESSOS code + +Existing imports remain valid: + +```python +from pyqsc_jax.near_axis import near_axis +``` + +The adapter preserves historical component-first arrays, mutable `x` and +`dofs`, `R0`, `Z0`, `phi`, `B_axis`, `grad_B_axis`, `iota`, boundary +conversion, and plotting methods. Its constructor and `dofs` setter now use +the same normal/binormal packing order. It delegates all physics to the +immutable core. + +New code should prefer `Qsc`/`solve`. Keep the adapter only at a downstream +compatibility boundary. The package does not emit an import-time or runtime +deprecation warning while ESSOS depends on this contract. + +## Behavioral differences + +- even `nphi` is honored rather than silently changed; +- general `rs` and `zc` axis coefficients are supported; +- failed solves return structured nonconvergence evidence; +- sigma and r2 derivatives are implicit derivatives of converged equations; +- inverse transform solving is explicit through `solve_for`; +- nonzero global-search results never claim global optimality; +- plasma/external separation requires a positive formal radius. diff --git a/docs/references.bib b/docs/references.bib new file mode 100644 index 0000000..03c5f9a --- /dev/null +++ b/docs/references.bib @@ -0,0 +1,120 @@ +@article{garren1991existence, + author = {Garren, D. A. and Boozer, Allen H.}, + title = {Existence of quasihelically symmetric stellarators}, + journal = {Physics of Fluids B}, + volume = {3}, + pages = {2822--2834}, + year = {1991}, + doi = {10.1063/1.859916}, +} + +@article{garren1991magnetic, + author = {Garren, D. A. and Boozer, Allen H.}, + title = {Magnetic field strength of toroidal plasma equilibria}, + journal = {Physics of Fluids B}, + volume = {3}, + pages = {2805--2821}, + year = {1991}, + doi = {10.1063/1.859915}, +} + +@article{landreman2018direct, + author = {Landreman, Matt and Sengupta, Wrick}, + title = {Direct construction of optimized stellarator shapes. Part 1. Theory in cylindrical coordinates}, + journal = {Journal of Plasma Physics}, + volume = {84}, + pages = {905840616}, + year = {2018}, + doi = {10.1017/S0022377818001289}, + eprint = {1809.10233}, + archivePrefix = {arXiv}, +} + +@article{landreman2019numerical, + author = {Landreman, Matt and Sengupta, Wrick and Plunk, Gabriel G.}, + title = {Direct construction of optimized stellarator shapes. Part 2. Numerical quasisymmetric solutions}, + journal = {Journal of Plasma Physics}, + volume = {85}, + pages = {905850103}, + year = {2019}, + doi = {10.1017/S0022377818001344}, +} + +@article{landreman2019highorder, + author = {Landreman, Matt and Sengupta, Wrick}, + title = {Constructing stellarators with quasisymmetry to high order}, + journal = {Journal of Plasma Physics}, + volume = {85}, + pages = {815850601}, + year = {2019}, + doi = {10.1017/S0022377819000783}, + eprint = {1908.10253}, + archivePrefix = {arXiv}, +} + +@article{jorge2020arbitrary, + author = {Jorge, Rogerio and Sengupta, Wrick and Landreman, Matt}, + title = {Near-axis expansion of stellarator equilibrium at arbitrary order in the distance to the axis}, + journal = {Journal of Plasma Physics}, + volume = {86}, + pages = {905860106}, + year = {2020}, + doi = {10.1017/S0022377820000033}, + eprint = {1911.02659}, + archivePrefix = {arXiv}, +} + +@article{landreman2021figures, + author = {Landreman, Matt}, + title = {Figures of merit for stellarators near the magnetic axis}, + journal = {Journal of Plasma Physics}, + volume = {87}, + pages = {905870112}, + year = {2021}, + doi = {10.1017/S0022377820001658}, + eprint = {2012.00865}, + archivePrefix = {arXiv}, +} + +@article{rodriguez2022weakly, + author = {Rodríguez, Eduardo and Sengupta, Wrick and Bhattacharjee, Amitava}, + title = {Weakly quasisymmetric near-axis solutions to all orders}, + journal = {Physics of Plasmas}, + volume = {29}, + pages = {012507}, + year = {2022}, + doi = {10.1063/5.0076583}, +} + +@article{giuliani2022single, + author = {Giuliani, Andrew and Wechsung, Florian and Cerfon, Antoine and Stadler, Georg and Landreman, Matt}, + title = {Single-stage gradient-based stellarator coil design: optimization for near-axis quasi-symmetry}, + journal = {Journal of Computational Physics}, + volume = {459}, + pages = {111147}, + year = {2022}, + doi = {10.1016/j.jcp.2022.111147}, + eprint = {2010.02033}, + archivePrefix = {arXiv}, +} + +@article{landreman2022mapping, + author = {Landreman, Matt}, + title = {Mapping the space of quasisymmetric stellarators using optimized near-axis expansion}, + journal = {Journal of Plasma Physics}, + volume = {88}, + pages = {905880616}, + year = {2022}, + doi = {10.1017/S0022377822001258}, + eprint = {2209.11849}, + archivePrefix = {arXiv}, +} + +@article{curvo2025deep, + author = {Curvo, Ines and Ferreira, Jorge and Jorge, Rogerio}, + title = {Using deep learning to design high aspect ratio fusion devices}, + journal = {Journal of Plasma Physics}, + volume = {91}, + year = {2025}, + doi = {10.1017/S002237782400165X}, +} diff --git a/docs/release_checklist.md b/docs/release_checklist.md new file mode 100644 index 0000000..d497a2d --- /dev/null +++ b/docs/release_checklist.md @@ -0,0 +1,20 @@ +# Release checklist + +No release is made from an unreviewed branch. + +- [x] Full line and branch coverage is at least 95%. +- [x] Ruff, all tests, all examples, docs with warnings as errors, doctests, + and compatibility checks pass. +- [x] Wheel and sdist metadata pass `twine check`. +- [x] Wheel and sdist each install in clean environments and pass a physics + smoke test. +- [x] Upstream pyQSC parity and high-resolution cases are rerun. +- [x] ESSOS vacuum and finite-beta integration passes in a clean environment. +- [x] Dependency and license audits have no unresolved issue. +- [x] Performance report records hardware, versions, command, and raw samples. +- [x] Changelog, migration guide, README, API docs, and limitations are current. +- [x] Publication figures are regenerated with commit/parameter metadata. +- [ ] TestPyPI trusted-publishing dry run succeeds. +- [ ] PyPI OIDC environment is configured and production publish is approved. +- [ ] `CITATION.cff` version/date and Zenodo metadata are updated. +- [x] No correctness limitation or inherited downstream failure is hidden. diff --git a/docs/theory/axis_optimization.md b/docs/theory/axis_optimization.md new file mode 100644 index 0000000..6b5a64f --- /dev/null +++ b/docs/theory/axis_optimization.md @@ -0,0 +1,15 @@ +# Axis optimization + +The primary search residual is the weighted, mean-free \(B_{20}\) array. +For fixed axis and other physical inputs, `B2c` is eliminated analytically. +Bounded deterministic exploration feeds damped local least-squares solves; +endpoints are clustered into distinct basins and rebuilt on independent grids. + +Only `verified_zero` carries a global statement: it reaches the known zero +lower bound of the nonnegative primary residual within requested tolerances. +A nonzero `best_found` result is not a global-minimum proof. Secondary +selectors rank already refined basins and carry no separate certificate. + +See [branch-aware global search](global-search.md) and +[\(B_{20}\) diagnostics](b20-optimization.md) for equations, status +semantics, high-mode controls, and verification. diff --git a/docs/theory/b20-optimization.md b/docs/theory/b20-optimization.md new file mode 100644 index 0000000..d9d7c2d --- /dev/null +++ b/docs/theory/b20-optimization.md @@ -0,0 +1,146 @@ +# \(B_{20}\) diagnostics and analytic \(B_{2c}\) + +## Primary residual + +The optimization target is the nonconstant part of \(B_{20}\), not its mean. +For positive quadrature weights \(w_j\), define + +\[ +(PB_{20})_j=B_{20,j} +-\frac{\sum_k w_kB_{20,k}}{\sum_k w_k}. +\] + +`weighted_l2` is + +\[ +\left[ +\frac{\sum_j w_j(PB_{20})_j^2}{\sum_jw_j} +\right]^{1/2}\frac{1}{B_0}. +\] + +`b20_diagnostics` also returns a smooth high-\(p\) norm, the sampled maximum, +peak-to-peak variation, all nonzero toroidal Fourier coefficients, their norm, +their \(L^1\) certificate, and a high-mode tail ratio. The Fourier coefficients +use direct weighted quadrature in Boozer toroidal angle \(\varphi\), so they do +not assume that the uniform cylindrical-\(\phi\) samples are uniform in +\(\varphi\). + +## Exact affine elimination + +For fixed axis and all other physical inputs, the complete second-order solve +is affine in \(B_{2c}\): + +\[ +B_{20}(\varphi;B_{2c})=u(\varphi)+B_{2c}v(\varphi). +\] + +The implementation obtains \(u\) and \(v\) from two complete implicit linear +solves at \(B_{2c}=0\) and \(B_{2c}=1\). It then evaluates + +\[ +B_{2c}^{\star} +=-\frac{\langle Pu,Pv\rangle_w}{\langle Pv,Pv\rangle_w}. +\] + +This is the global minimizer in the one-dimensional affine subproblem whenever +\(\langle Pv,Pv\rangle_w>0\). `optimize_B2c` fully recomputes the r2 solution +at the optimum and, when present, recomputes r3 and magnetic shear so no stale +derived arrays survive. `affine_reconstruction_error` independently compares +that solution to \(u+B_{2c}^{\star}v\). + +For a circular axis the projected response can vanish to roundoff because +changing \(B_{2c}\) changes only the constant mode. Such cases are marked +`degenerate`, retain the supplied \(B_{2c}\), and already have a vanishing +nonconstant residual. + +## Independent resolution verification + +`verify_B20_resolution` holds the physical candidate fixed and rebuilds it at + +\[ +n_\phi(m)=m(n_\phi-1)+1 +\] + +for each requested multiplier. The default `(1, 2, 4)` therefore doubles and +quadruples the number of grid intervals while retaining an odd grid. It reports +every dense diagnostic and successive relative changes. This check is +deliberately separate from the optimization; a small collocation objective is +not accepted without fine-grid verification. + +## Validation + +Tests establish: + +- direct agreement of every scalar diagnostic with its definition; +- stationarity and strict improvement of the analytic optimum; +- machine-precision affine reconstruction; +- JIT and JVP agreement with finite differences; +- correct degenerate behavior for a circular axis; +- r3 and shear recomputation; +- convergence of weighted and maximum residuals under grid refinement. + +The branch-aware multistart workflow built on these residuals is described in +{doc}`global-search`. + +## Optimizer comparison and selected case + +The public [Curvo stellarator database](https://stellarator.physics.wisc.edu/) +was screened before optimization. Configuration +[57409](https://stellarator.physics.wisc.edu/app/plot/57409) was selected as +the traceable low-\(B_{20}\) seed: pyQSC_JAX independently reproduces +\(\lvert\iota\rvert=2.833\), \(r_\mathrm{sing}=0.238\) m, and a complete +Curvo/Table-3 pass when the minimum-transform threshold is tightened to 0.4. + +Several optimizers were then tested from exactly that downloaded axis. For the +common comparison, Fourier modes 1--3 were varied and \(B_{2c}\) was +eliminated exactly at every objective evaluation. Timings are warm local +measurements and exclude the shared JAX compilation. + +| method | evaluations | time [s] | independently evaluated weighted \(L^2\) | +| --- | ---: | ---: | ---: | +| exact \(B_{2c}\) only | 1 | — | \(3.09928\times10^{-2}\) | +| SciPy L-BFGS-B | 70 | 0.141 | \(2.01392\times10^{-4}\) | +| SciPy `least_squares` | 28 | 1.349 | \(1.02501\times10^{-4}\) | +| low-budget differential evolution | 120 | 0.104 | \(2.65766\times10^{-1}\) | +| pyQSC_JAX multistart Levenberg--Marquardt | 130 | 21.772* | \(1.08279\times10^{-4}\) | + +The SciPy timings exclude the shared JAX residual/Jacobian compilation; the +asterisked multistart timing includes compilation of its independent closure. +The coarse differential-evolution budget is included to show why a method +label alone does not establish global quality. It did not locate the narrow +good basin. Bounded `least_squares` was then continued one Fourier mode at a +time, re-eliminating \(B_{2c}\) at every evaluation: + +| highest varied mode | independently evaluated weighted \(L^2\) | +| ---: | ---: | +| 3 | \(1.02501\times10^{-4}\) | +| 4 | \(1.13643\times10^{-6}\) | +| 5 | \(5.58886\times10^{-7}\) | +| 6 | \(2.75804\times10^{-7}\) | +| 7 | \(3.02110\times10^{-8}\) | +| 8 | \(1.27645\times10^{-10}\) | + +The resulting `b20_optimized_good` configuration was held fixed for the +independent resolution checks: + +| `nphi` | weighted \(L^2\) | +| ---: | ---: | +| 121 | \(1.27431807\times10^{-10}\) | +| 241 | \(1.27645060\times10^{-10}\) | +| 481 | \(1.28075313\times10^{-10}\) | + +At `nphi=121`, its dense maximum is \(2.62\times10^{-10}\) and peak-to-peak +variation is \(5.23\times10^{-10}\). This is \(2.4321\times10^8\) times +smaller in weighted residual than exact \(B_{2c}\) optimization of its public +database seed. The independently recomputed design has +\(\lvert\iota\rvert=2.964\), \(r_\mathrm{sing}=0.2493\) m, maximum elongation +2.95, minimum \(L_{\nabla B}=0.437\) m, minimum +\(L_{\nabla\nabla B}=0.340\) m, positive Mercier margin, and a complete +Curvo-profile pass. The README surface is drawn at 0.075 m, giving 3.32-fold +radial clearance relative to the truncated-map singularity diagnostic. + +The result is a verified numerical basin, not a proof that no other basin is +better. The runnable comparison is +`benchmarks/benchmark_b20_optimizers.py`; its frozen report records the full +environment and raw values in +`benchmarks/reports/2026-07-30-b20-optimizers-apple-m4.json`. diff --git a/docs/theory/coordinates-and-conventions.md b/docs/theory/coordinates-and-conventions.md new file mode 100644 index 0000000..3a240cf --- /dev/null +++ b/docs/theory/coordinates-and-conventions.md @@ -0,0 +1,106 @@ +# Axis coordinates and conventions + +## Fourier axis + +`Axis` represents a closed magnetic axis in cylindrical coordinates over one +field period: + +\[ +R(\phi) = +\sum_{n=0}^{N_F-1} +\left[ +R_{cn}\cos(n n_{\mathrm{fp}}\phi) ++R_{sn}\sin(n n_{\mathrm{fp}}\phi) +\right], +\] + +\[ +Z(\phi) = +\sum_{n=0}^{N_F-1} +\left[ +Z_{cn}\cos(n n_{\mathrm{fp}}\phi) ++Z_{sn}\sin(n n_{\mathrm{fp}}\phi) +\right]. +\] + +The arrays `rc`, `rs`, `zc`, and `zs` are padded with zeros to a common +one-dimensional shape `(nfourier,)`. The packed degree-of-freedom order is +exactly `rc, rs, zc, zs`. A stellarator-symmetric axis has `rs = zc = 0`. + +```python +import pyqsc_jax as qsc + +axis = qsc.Axis( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, +) +``` + +The representation follows the direct-construction convention +{cite}`landreman2018direct`. Unlike the historical pyQSC constructor, even +`nphi` is not silently changed; grid resolution is an explicit caller choice. + +## Sampled geometry + +`compute_axis_geometry(axis, nphi=...)` samples one field period on the +endpoint-free grid + +\[ +\phi_j = \frac{2\pi j}{n_{\mathrm{fp}}n_\phi}, +\qquad j=0,\ldots,n_\phi-1. +\] + +Analytic Fourier derivatives provide the first three derivatives of \(R\) and +\(Z\). Vector components in the cylindrical frame are ordered `(R, phi, Z)`; +Cartesian components are ordered `(x, y, z)`. Arrays of sampled vectors have +shape `(nphi, 3)`. + +The tangent, normal, and binormal are + +\[ +\boldsymbol{t} = \frac{\mathrm{d}\boldsymbol{r}_0}{\mathrm{d}\ell}, +\qquad +\boldsymbol{n} = \frac{1}{\kappa} +\frac{\mathrm{d}\boldsymbol{t}}{\mathrm{d}\ell}, +\qquad +\boldsymbol{b} = \boldsymbol{t}\times\boldsymbol{n}. +\] + +They form a right-handed orthonormal frame. Torsion uses the convention of +Landreman and Sengupta and the usual differential-geometry formula; this is +the opposite sign to the original Garren–Boozer convention +{cite}`garren1991magnetic,landreman2019highorder`. + +The Boozer toroidal grid is normalized so that + +\[ +\frac{\mathrm{d}\varphi}{\mathrm{d}\phi} += \frac{B_0}{|G_0|} +\frac{\mathrm{d}\ell}{\mathrm{d}\phi}, +\qquad +\varphi(0)=0, +\] + +and advances by \(2\pi/n_{\mathrm{fp}}\) per field period. Periodic +differentiation uses Fourier collocation matrices on the same endpoint-free +grid. + +## Topology and validity + +`frame_helicity` counts the signed winding of the Frenet normal around the +axis over one field period. It is intrinsic to the axis. The first-order solve +later combines it with the flux and covariant-field signs. + +Frenet coordinates require nonzero speed and curvature everywhere. +`GeometryDiagnostics` reports: + +- minimum speed and curvature; +- minimum cylindrical radius; +- maximum frame orthogonality error; +- minimum frame determinant; +- Boolean Frenet and cylindrical-coordinate validity. + +Invalid samples are represented by non-finite frame/torsion values and a false +validity flag. High-level solvers reject them with a structured report rather +than allowing them to propagate silently. diff --git a/docs/theory/criteria.md b/docs/theory/criteria.md new file mode 100644 index 0000000..42cfe4f --- /dev/null +++ b/docs/theory/criteria.md @@ -0,0 +1,52 @@ +# Named design criteria + +`Criteria.from_curvo_2025` implements the screening profile in Table 3 of +Curvo, Ferreira, and Jorge {cite}`curvo2025deep`. It is a named, configurable +profile rather than a package-wide definition of a “good stellarator.” + +For the paper's normalization \(R_{c0}=1\,\mathrm{m}\) and +\(B_0=1\,\mathrm{T}\), the profile requires: + +- positive magnetic-axis length; +- \(\lvert\iota\rvert\geq0.2\); +- maximum elongation no greater than \(10\); +- \(\min L_{\nabla B}\geq0.1\,\mathrm{m}\); +- minimum cylindrical axis radius at least \(0.3\,\mathrm{m}\); +- singular radius at least \(0.05\,\mathrm{m}\); +- \(\min L_{\nabla\nabla B}\geq0.1\,\mathrm{m}\); +- \(B_{20}\) peak-to-peak variation no greater than + \(5\,\mathrm{T/m^2}\); +- \(\beta\geq10^{-4}\); and +- \(D_{\mathrm{Merc}}r^2>0\). + +The profile uses the paper's near-axis pressure proxy + +\[ +\beta=-\frac{\mu_0p_2r_{\mathrm{singularity}}^2}{B_0^2}. +\] + +When `major_radius` or `B0` differs from the reference normalization, length +thresholds scale with `major_radius` and the \(B_{20}\)-variation threshold +scales as \(B_0/R^2\). Dimensionless thresholds do not change. Callers can +override any threshold by name: + +```python +criteria = qsc.Criteria.from_curvo_2025( + major_radius=1.7, + B0=2.5, + maximum_elongation=8.0, + minimum_beta=2.0e-4, +) +report = criteria.evaluate(solution) +``` + +Every `CriterionEvaluation` contains the measured value, threshold, comparison +sense, units, pass flag, and a signed raw margin. A positive margin is +favorable for both minimum and maximum criteria. `report.passed` is true only +if all ten checks pass. + +The singular radius is the first coordinate-map singularity of the truncated +near-axis construction. It is useful as a screening proxy but does not prove +that a finite-radius equilibrium has nested, nonintersecting flux surfaces. +Likewise, this profile does not replace coil feasibility, fast-particle +confinement, MHD equilibrium, or free-boundary validation. diff --git a/docs/theory/diagnostics.md b/docs/theory/diagnostics.md new file mode 100644 index 0000000..8a412b4 --- /dev/null +++ b/docs/theory/diagnostics.md @@ -0,0 +1,20 @@ +# Diagnostics + +The package keeps diagnostics distinct from solver success: + +- `RootSolveReport` and `LinearSolveReport` record residuals, iterations, + conditioning, finiteness, and convergence; +- `GeometryDiagnostics` records speed, curvature, cylindrical radius, and + frame validity; +- `MercierDiagnostics` separates magnetic-well and geodesic contributions; +- `SingularityDiagnostics` finds the first zero of the quadratic regular-map + Jacobian and reports its residual; +- `B20Diagnostics` reports weighted \(L^2\), smooth and grid maxima, + peak-to-peak variation, Fourier coefficients, and tail decay; +- `CriteriaReport` returns every measurement, threshold, comparison, margin, + and pass flag. + +The detailed definitions are in [field tensors](field-jet.md), +[\(B_{20}\) optimization](b20-optimization.md), and +[named criteria](criteria.md). A small residual is meaningful only with a +finite, well-conditioned solve and independent resolution verification. diff --git a/docs/theory/field-jet.md b/docs/theory/field-jet.md new file mode 100644 index 0000000..582fa73 --- /dev/null +++ b/docs/theory/field-jet.md @@ -0,0 +1,126 @@ +# Total on-axis field jet + +## Regular coordinates + +The total magnetic field, gradient, and Hessian are computed without +constructing a finite-radius surface. Define the regular transverse +coordinates + +\[ +q_1=r\cos\vartheta,\qquad q_2=r\sin\vartheta. +\] + +Through quadratic order the position is + +\[ +\boldsymbol{x} +=\boldsymbol{r}_0+\boldsymbol{d}_1q_1+\boldsymbol{d}_2q_2 ++\frac12\boldsymbol{h}_{11}q_1^2 ++\boldsymbol{h}_{12}q_1q_2 ++\frac12\boldsymbol{h}_{22}q_2^2. +\] + +The vectors \(\boldsymbol d_a\) come from the first-order coefficients and +\(\boldsymbol h_{ab}\) from the complete second-order solution. Consequently, +the coordinate map and its inverse are nonsingular on axis even though polar +coordinates \((r,\vartheta)\) are not. + +The implementation follows equations 55–83 of the audited surface-free +plasma/coil derivation. It expands + +\[ +\boldsymbol B_{\mathrm{tot}} +=P\left[ +\boldsymbol x_{,\varphi} ++\iota_N(-q_2\boldsymbol x_{,q_1}+q_1\boldsymbol x_{,q_2}) +\right] +\] + +through second order in \(q_1,q_2\), then applies the ordinary inverse-map +chain rule. This avoids the long generated component expressions in pyQSC; +those expressions are retained only as an independent regression reference. + +## A subtle periodicity rule + +Scalar coefficients and cylindrical components are periodic over one field +period. Fixed Cartesian components of a vector generally are not: the +cylindrical basis rotates by \(2\pi/n_{\mathrm{fp}}\). A cylindrical vector +\(\boldsymbol v=(v_R,v_\phi,v_Z)\) is therefore differentiated using + +\[ +\frac{\mathrm d\boldsymbol v}{\mathrm d\varphi} +=\left(v_R'-\phi'v_\phi\right)\boldsymbol e_R ++\left(v_\phi'+\phi'v_R\right)\boldsymbol e_\phi ++v_Z'\boldsymbol e_Z. +\] + +Applying a periodic Fourier derivative directly to Cartesian frame samples +is incorrect for \(n_{\mathrm{fp}}>1\). The connection terms above are part of +the implementation and have dedicated Maxwell-identity and pyQSC regression +tests. + +## Array ordering and diagnostics + +Canonical arrays place the sample axis first: + +- `field[n, i]` is \(B_i\); +- `gradient[n, i, j]` is \(\partial B_i/\partial x_j\); +- `hessian[n, i, j, k]` is + \(\partial^2B_i/(\partial x_j\partial x_k)\). + +`hessian_frenet` follows pyQSC's historical +`(sample, derivative, derivative, field)` ordering in the +`(normal, binormal, tangent)` frame. The compatibility adapter exposes this +array as `grad_grad_B`. + +`FieldJet` includes the regular-coordinate Jacobian, inverse Jacobian, and +coordinate Hessian together with: + +- field and gradient agreement with the lower-order construction; +- \(\nabla\cdot\boldsymbol B\); +- symmetry in the two derivative indices; +- the gradient of \(\nabla\cdot\boldsymbol B\); +- \(L_{\nabla\nabla B}\) and its inverse. + +Vacuum tests additionally require full index symmetry. Finite-current tests +retain only the mixed-spatial-derivative symmetry that Maxwell's equations +require. + +## Mercier quantities + +`MercierDiagnostics` provides the leading near-axis quantities +`d2_volume_d_psi2`, `DGeod_times_r2`, `DWell_times_r2`, and +`DMerc_times_r2`. Their normalization and signs match pyQSC at the audited +upstream commit. The geodesic and well contributions are kept separate so a +caller can inspect cancellations rather than accepting only their sum. + +## Singular radius + +The quadratic regular map also gives a concise singular-radius calculation. +Its determinant through second order in \(r\) is + +\[ +\widehat g = +g_0+r(g_{1c}\cos\vartheta+g_{1s}\sin\vartheta) ++r^2(g_{20}+g_{2s}\sin2\vartheta+g_{2c}\cos2\vartheta). +\] + +The coefficients are obtained by collecting powers of \(q_1,q_2\) using the +multilinearity of the determinant of the regular coordinate map. This replaces +the long generated coefficient expressions in pyQSC while producing the same +quadratic Jacobian. As an internal quasisymmetry check, `g1s` must vanish. + +For each toroidal sample, all positive quadratic roots are enumerated on a +uniform angular grid. The best root seeds a vectorized Newton solve for + +\[ +\widehat g=0,\qquad +\frac{\partial\widehat g}{\partial\vartheta}=0. +\] + +The default 256 angular seeds and eight Newton steps reproduce the audited +pyQSC radii while retaining JIT and forward-mode differentiation. The result +contains the radius, angle, and residual at every toroidal sample, together +with the global minimum. It diagnoses loss of regularity of the truncated +near-axis map; it is not a guarantee that a finite-radius equilibrium exists +up to that radius. diff --git a/docs/theory/field_tensors.md b/docs/theory/field_tensors.md new file mode 100644 index 0000000..a0c82b6 --- /dev/null +++ b/docs/theory/field_tensors.md @@ -0,0 +1,21 @@ +# Field tensors + +Regular transverse coordinates \(q_1=r\cos\vartheta\) and +\(q_2=r\sin\vartheta\) remove the polar-coordinate singularity on axis. +Applying the inverse-coordinate chain rule to the complete r2 expansion gives + +\[ +B_i,\qquad D_{ij}=\partial_jB_i,\qquad +H_{ijk}=\partial_j\partial_kB_i. +\] + +Vacuum fields satisfy symmetric trace-free identities; finite-current total +fields retain derivative-index symmetry and Ampère-law antisymmetry in the +gradient. The result reports divergence, derivative asymmetry, and the +gradient of divergence. + +See [total on-axis field jet](field-jet.md) for the full regular-map +derivation, periodic cylindrical connection terms, tensor order, scale +lengths, and singular-radius construction. Implementation: +`pyqsc_jax.field.total_field_jet`. Validation: +`tests/physics/test_field_jet.py`. diff --git a/docs/theory/first-order.md b/docs/theory/first-order.md new file mode 100644 index 0000000..9aa9bfd --- /dev/null +++ b/docs/theory/first-order.md @@ -0,0 +1,133 @@ +# First-order quasisymmetric construction + +## Inputs and signs + +The first-order solve uses the direct-construction conventions of +{cite}`landreman2019highorder`. The leading field strength is `B0`; `sG` is +the sign of \(G_0\), and `spsi` is the sign of toroidal flux. The axis geometry +provides \(\kappa\), \(\tau\), \(\varphi\), and + +\[ +\frac{|G_0|}{B_0} = \frac{L}{2\pi}. +\] + +`frame_helicity` is the signed winding of the Frenet normal over one field +period. The sign-adjusted helicity and transform entering the sigma equation +are + +\[ +N = h\,s_\psi s_G, +\qquad +\iota_N = \iota + N n_{\mathrm{fp}}. +\] + +This convention gives \(N=0\) for the standard QA example and \(N=-1\) for +the documented four-field-period QH example when `spsi = sG = 1`. + +## Periodic sigma equation + +For prescribed \(\bar\eta\), the unknown state contains the rotational +transform in its first slot and samples of the periodic function +\(\sigma(\varphi)\) in the remaining slots. The first sigma sample is replaced +by the prescribed `sigma0`, which removes the otherwise redundant degree of +freedom. The collocation residual is + +\[ +\frac{\mathrm d\sigma}{\mathrm d\varphi} ++ \iota_N +\left[ +1 + \sigma^2 ++ \left(\frac{\bar\eta^2}{\kappa^2}\right)^2 +\right] +-2\frac{\bar\eta^2}{\kappa^2} +\left(-s_\psi\tau + \frac{I_2}{B_0}\right) +\frac{s_G|G_0|}{B_0} +=0. +\] + +Fourier collocation differentiates the periodic array. A damped dense Newton +solve uses the exact JAX Jacobian, residual-based backtracking, a step +stagnation test, and configurable absolute/relative tolerances. It returns a +`RootSolveReport` containing: + +- initial and final infinity-norm residuals; +- the effective residual tolerance and final step norm; +- Newton and accumulated backtracking counts; +- final Jacobian condition number; +- convergence, finiteness, and stagnation flags. + +Reaching an iteration or stagnation limit does not silently certify a +solution. The numerical candidate and a false convergence flag are returned +so batched workflows can apply their own rejection policy. + +## Implicit differentiation + +The converged candidate is supplied to SOLVAX's custom root wrapper after +`stop_gradient`. JVPs and VJPs therefore differentiate + +\[ +F(x,p)=0, +\qquad +\frac{\mathrm dx}{\mathrm dp} += +-\left(\frac{\partial F}{\partial x}\right)^{-1} +\frac{\partial F}{\partial p}, +\] + +not the sequence of Newton iterates. Tests compare JVP and VJP, finite +differences in \(\bar\eta\) and an axis Fourier coefficient, eager/JIT +evaluation, and VMAP batches. + +## First-order surface and field jet + +The Frenet-plane coefficients are + +\[ +X_{1s}=0,\qquad +X_{1c}=\frac{\bar\eta}{\kappa},\qquad +Y_{1s}=\frac{s_Gs_\psi\kappa}{\bar\eta},\qquad +Y_{1c}=\sigma Y_{1s}. +\] + +They are also provided in an untwisted frame for boundary conversion. The +solution contains the on-axis field in cylindrical and Cartesian bases, the +Cartesian and cylindrical field-gradient tensors, elongation, and + +\[ +L_{\nabla B} += B_0\sqrt{\frac{2}{\nabla\boldsymbol B:\nabla\boldsymbol B}}. +\] + +Canonical arrays use a leading sample axis. Thus `B_axis.shape == +(nphi, 3)` and `grad_B_axis.shape == (nphi, 3, 3)`, with +`grad_B_axis[n, i, j] = d B_i / d x_j`. The ESSOS adapter preserves the +historical component-axis layout. In the vacuum limit, resolution studies +verify the trace-free and symmetric gradient identities spectrally. + +## Reference configuration + +For + +```python +import pyqsc_jax as qsc + +solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, +) +``` + +the validated result is + +\[ +\iota = 0.41830690943386617, +\qquad +L = 6.340238817434161. +\] + +The sigma array, QA/QH topology, finite-current result, on-axis field, and +field gradient agree with pyQSC at the audited upstream revision recorded in +the refactor baseline. diff --git a/docs/theory/first_order.md b/docs/theory/first_order.md new file mode 100644 index 0000000..e7aa411 --- /dev/null +++ b/docs/theory/first_order.md @@ -0,0 +1,13 @@ +# First order + +First order solves the periodic sigma equation for cross-section orientation +and rotational transform, then constructs the elliptical surface coefficients +and on-axis field gradient. The nonlinear solve is damped, convergence-tested, +and differentiated at the converged root through the implicit-function +theorem. + +The complete equations, sign conventions, failure modes, implementation +symbols, and reference value are in the +[first-order derivation](first-order.md). Direct tests live in +`tests/physics/test_first_order.py`; AD checks are in +`tests/numerics/test_first_order_autodiff.py`. diff --git a/docs/theory/global-search.md b/docs/theory/global-search.md new file mode 100644 index 0000000..752d5fb --- /dev/null +++ b/docs/theory/global-search.md @@ -0,0 +1,83 @@ +# Branch-aware axis search + +## Objective and exact scalar elimination + +`search_axis` minimizes the full weighted, mean-free \(B_{20}\) residual on +the toroidal grid: + +\[ +r_j=\sqrt{\frac{w_j}{\sum_kw_k}}\, +\frac{B_{20,j}-\langle B_{20}\rangle_w}{B_0}. +\] + +It does not use a few hand-selected collocation points. At every axis +evaluation, \(B_{2c}\) is eliminated analytically before the residual is +formed. Target-transform problems use `solve_for="etabar"` or +`solve_for="I2"`, so the requested transform and the full sigma collocation +state are solved together on the selected branch. + +## Reproducible global-to-local workflow + +The implemented workflow is: + +1. include the supplied axis and any explicit branch-specific seeds; +2. fill the bounded coefficient box with a deterministic Halton sequence; +3. evaluate the coarse population with `jax.vmap`, replacing invalid + geometry or unconverged solves by an infinite score without contaminating + the batch with NaNs; +4. refine the best starts with a damped Gauss–Newton/Levenberg–Marquardt + iteration using the full JAX residual Jacobian; +5. cluster endpoints in normalized bounded-coordinate space; +6. apply hard geometry, solver, conditioning, and optional `Criteria` checks; +7. select lexicographically by primary tolerance and then by the requested + secondary selector; and +8. independently rebuild the selected physical candidate at the requested + toroidal resolutions. + +Every optimized variable is represented internally by a hyperbolic-tangent +map. Bounds are therefore never imposed by clipping a trial step. The width of +each coefficient interval supplies mode-aware variable scaling; callers should +tighten high-mode bounds as the Fourier order increases. + +`continue_axis_search` provides staged Fourier continuation. Each stage may +increase the retained mode count and `nphi`; all common Fourier coefficients +from the prior best basin warm-start the next stage. This is preferable to +activating many weakly resolved high modes at once. + +The built-in secondary selectors rank feasible basins by singular radius, +minimum first- or second-derivative scale length, maximum elongation, a +fourth-order axis Sobolev norm, distance from a target axis length, or minimum +normalized criteria margin. They are not silently mixed into the primary +quasisymmetry residual with arbitrary weights. + +## Status and global claims + +The result uses one of: + +- `verified_zero`; +- `best_found`; +- `no_feasible_candidate`; +- `branch_fold`; +- `ill_conditioned`; +- `solver_failure`; or +- `verification_failure`. + +`verified_zero` has one deliberately narrow meaning. The nonconstant +\(B_{20}\) residual is nonnegative, so a candidate whose weighted and maximum +residuals reach zero within the requested absolute, resolution, and spectral +tolerances attains the known global lower bound of that primary residual. The +flag does not certify that the secondary selector is globally optimal. + +Any nonzero result is `best_found` and reports its search budget and number of +distinct basins. It carries no global-minimum claim. A small coarse-grid value +that fails doubled-grid, maximum-error, or spectral-tail checks is +`verification_failure`. + +## High-frequency safeguards + +The search exposes coefficient bounds, bounded-coordinate scaling, staged +mode activation, a Sobolev selector, the complete nonzero Fourier spectrum, +its \(L^1\) certificate, tail ratio, and doubled/quadrupled-grid changes. +These controls are important because high-order axis optimization is +spectrally ill-conditioned: unresolved coefficients can move error between +collocation points without improving the continuous configuration. diff --git a/docs/theory/inverse-solves.md b/docs/theory/inverse-solves.md new file mode 100644 index 0000000..b1e7064 --- /dev/null +++ b/docs/theory/inverse-solves.md @@ -0,0 +1,92 @@ +# Target-transform inverse solves + +## Coupled formulation + +The forward first-order problem prescribes \(\bar\eta\) and solves the +periodic sigma equation for \(\sigma(\varphi)\) and \(\iota\). The inverse +problem instead fixes \(\iota_\star\) and promotes either \(\bar\eta\) or +\(I_2\) to the scalar unknown. The collocation state still has `nphi` +unknowns: one scalar parameter and `nphi - 1` sigma samples, with +\(\sigma(0)=\sigma_0\) imposed by substitution. + +For `solve_for="etabar"`, the scalar state is \(u\) and + +\[ +\bar\eta=s_\eta\exp(u), +\qquad s_\eta\in\{-1,1\}. +\] + +This prevents a Newton step from crossing the singular +\(\bar\eta=0\) surface and changing the orientation branch. The sign comes +from the seed. If no seed is supplied, `-1.0` selects the default negative +branch. A zero seed selects the positive branch from the smallest +representable positive magnitude. + +For `solve_for="I2"`, the scalar state is \(I_2\) directly. In both cases the +same converged dense Newton policy and SOLVAX implicit-root rule used by the +forward sigma solve apply to the coupled inverse system. + +## Local folds + +After convergence, the forward sigma residual is differentiated with respect +to its ordinary state \(y=(\iota,\sigma_1,\ldots)\) and the solved parameter +\(p\). The response is obtained from + +\[ +\frac{\partial F}{\partial y}\frac{\mathrm dy}{\mathrm dp} +=-\frac{\partial F}{\partial p}. +\] + +`response_derivative` is \(\mathrm d\iota/\mathrm dp\). Its reciprocal is the +implicit derivative of the branch-local inverse away from a fold. +`branch_fold` is true when this response is nonfinite or its magnitude is no +larger than `fold_tolerance`. + +A fold does not imply that the full solution branch is singular. It means the +projection \(p(\iota)\) is locally multivalued. The direct tests contain two +different negative-\(\bar\eta\) solutions with the same transform and +opposite response slopes. Therefore a nonconverged target solve must not be +interpreted as proof that no branch exists. + +## Pseudo-arclength continuation + +`continue_etabar_branch` begins from two converged forward points. A secant in +physical \((\bar\eta,\iota)\) coordinates supplies the oriented tangent and +predictor. The corrector solves the complete sigma collocation system together +with + +\[ +\boldsymbol t\mathbin{\cdot} +\left[ +(\bar\eta,\iota)-(\bar\eta,\iota)_{\mathrm{predict}} +\right]=0. +\] + +The nonlinear state contains the log-magnitude of \(\bar\eta\), \(\iota\), and +all free sigma samples. Thus the corrector can pass through +\(\mathrm d\iota/\mathrm d\bar\eta=0\) without changing the selected +\(\bar\eta\) sign. A response-slope sign change marks a fold even when no +discrete point lands exactly at zero slope. + +`ContinuationResult` retains every corrected immutable solution, the branch +tangents, response derivatives, and fold flags. It stops with +`status="solver_failure"` after a failed corrector or +`status="branch_zero_crossing"` before leaving the fixed-sign branch. + +## Validation and acceptance + +Tests cover: + +- negative and positive \(\bar\eta\) sign branches; +- forward/inverse round trips for \(\bar\eta\) and finite \(I_2\); +- two distinct local \(\bar\eta\) basins at one target transform; +- full-state pseudo-arclength traversal across their intervening fold; +- independent forward-solve checks of corrected continuation points; +- r2 coefficient propagation; +- JIT and implicit JVPs against the reciprocal forward response and centered + finite differences; +- invalid and non-scalar inputs. + +Every caller must reject a result unless `root_report.converged` and +`root_report.finite` are true. Crossing a detected fold requires +pseudo-arclength continuation rather than repeated fixed-transform solves. diff --git a/docs/theory/plasma-current.md b/docs/theory/plasma-current.md new file mode 100644 index 0000000..5339e30 --- /dev/null +++ b/docs/theory/plasma-current.md @@ -0,0 +1,52 @@ +# Positive-volume plasma-current source + +The free-space Biot–Savart calculation must integrate the physical volume +measure. Expanding current density and the polar-coordinate Jacobian +separately makes both expressions look singular at \(r=0\) and makes it easy +to lose a factor of \(r\). The attached derivation instead defines the regular, +positive-volume source + +\[ +\mathbf W=\chi\mu_0\mathcal J_r\mathbf J, +\qquad \chi=s_Gs_\psi, +\] + +and obtains the exact identity + +\[ +\mathbf W=\chi\left[ +(I_r-r\bar B\beta_\vartheta)\mathbf x_\varphi +-(G_r+NI_r-r\bar B\beta_\varphi)\mathbf x_\vartheta +\right]. +\] + +`plasma_current_source` expands this identity directly: + +\[ +\frac{\mathbf W}{L} +=r\mathbf w_1+r^2\mathbf w_2+O(r^3), +\qquad +\mathbf w_1=j\mathbf t, +\qquad +j=2\chi I_2=\mu_0J_\parallel(0). +\] + +The quadratic term retains the independently pressure-driven +\(\beta_{1s}\) and \(C_2=G_2+NI_2\) paths. It is represented by cosine and +sine vector coefficients without division by \(I_2\), so it remains regular +as the on-axis parallel current passes through zero. + +For a mandatory formal current-channel radius \(a>0\), + +\[ +\mu_0 I_p(a)=\pi a^2j=2\pi\chi I_2a^2. +\] + +`enclosed_current_from_covariant` and `covariant_current_from_enclosed` +implement this conversion. Holding \(I_2\) fixed while varying \(a\) holds the +on-axis current density fixed and makes the enclosed current scale as \(a^2\). +Holding enclosed amperes fixed instead requires \(I_2\propto a^{-2}\), which +is a different asymptotic ordering. + +`evaluate_weighted_current` evaluates the regular truncated source at any +radial/helical-angle batch without constructing a finite-radius surface. diff --git a/docs/theory/plasma-field.md b/docs/theory/plasma-field.md new file mode 100644 index 0000000..ba713a8 --- /dev/null +++ b/docs/theory/plasma-field.md @@ -0,0 +1,148 @@ +# Matched free-space plasma field + +The on-axis plasma field is not determined by local Ampère-law data alone. +Remote parts of the current distribution add a harmonic field. The decay +condition at spatial infinity selects that field uniquely and leads to a +matched volume-current Biot–Savart calculation. + +`regularized_axis_integral` evaluates the full-torus periodic finite part + +\[ +\mathcal R(\varphi)= +\int_0^{2\pi} +\left[ +L\frac{\mathbf t(\varphi')\times +[\mathbf r_0(\varphi)-\mathbf r_0(\varphi')]} +{|\mathbf r_0(\varphi)-\mathbf r_0(\varphi')|^3} +-\frac{\kappa(\varphi)\mathbf b(\varphi)} +{4|\sin[(\varphi'-\varphi)/2]|} +\right]d\varphi'. +\] + +All field-period copies are included. At the coincident grid point the +symmetric finite-part value is used; the two one-sided limits are equal and +opposite. + +The inner elliptical channel supplies the finite constants + +\[ +C_b=-\frac12-\frac12\log\frac{\mathcal T+2}{4} ++\frac{x^2+1}{\mathcal T+2}, +\qquad +C_n=-\frac{\chi\sigma}{\mathcal T+2}, +\] + +and a smooth angular correction \(\mathbf S\) containing the first-harmonic +weighted current and the second-order cross-sectional displacement. The +assembled leading field is + +\[ +\mathbf B_p(\mathbf r_0)= +\frac{ja^2}{4} +\left[ +\mathcal R+\kappa\mathbf b +\left(\log\frac{8L}{a}+C_b\right) ++\kappa C_n\mathbf n +\right] ++\mathbf S ++O(a^4\log a). +\] + +`matched_plasma_field_kernel` also exposes the same result at an arbitrary +matching reference length. Changing that length shifts the finite part and +the local logarithm separately but leaves their sum invariant. + +Validation includes: + +- exact cancellation of the arbitrary matching length; +- a circular-axis finite part below \(10^{-13}\); +- vanishing circular-ellipse core constants; +- the circular local-induction logarithm; +- vacuum and pressure-only limits; +- angular and toroidal resolution convergence; +- JIT and radius JVP checks; and +- an independent resolved volume-current Biot–Savart integral whose error + follows the predicted \(a^4|\log a|\) scaling. + +The returned `estimated_field_remainder` is an order-of-magnitude asymptotic +scale, not a rigorous error bound. `formal_radius_to_curvature_radius` should +be inspected before treating the slender-channel expansion as accurate. + +## Local elliptical gradient + +At leading gradient order the current channel is locally straight. In the +ordered Frenet basis \((\mathbf t,\mathbf n,\mathbf b)\), with the first tensor +index denoting derivative direction, the manuscript gives + +\[ +D^p=\frac{j}{\mathcal T+2} +\begin{pmatrix} +0&0&0\\ +0&\chi\sigma&1+(1+\sigma^2)/x^2\\ +0&-(1+x^2)&-\chi\sigma +\end{pmatrix}. +\] + +The canonical package tensor is field-component-first, so +`gradient_frenet` is the transpose of this matrix. +`elliptical_channel_gradient` implements the formula independently of a +near-axis solution and optionally rotates it into Cartesian coordinates. + +The local tensor satisfies + +\[ +\nabla\cdot\mathbf B_p=0,\qquad +D^p_{nb}-D^p_{bn}=j=\mu_0J_\parallel(0). +\] + +`plasma_gradient_on_axis` subtracts it from the total near-axis gradient. +The resulting external gradient is symmetric and trace-free within the +spectral resolution of the total solve. The result also supplies its five +independent Cartesian STF components in the order +`(xx, yy, xy, xz, yz)`. + +## Curved-channel Hessian + +The plasma Hessian cannot be obtained by naively differentiating the singular +Biot–Savart kernel through the current-carrying region. Those derivatives +produce local contact terms. `plasma_hessian_on_axis` instead evaluates the +cubic interior logarithmic potential of the resolved elliptical channel, +including three separately determined contributions: + +- the affine first-harmonic weighted current; +- the universal curved-channel metric term; and +- the second-order deformation of the current cross-section. + +The transverse derivatives are evaluated algebraically in the oriented +principal axes of the ellipse. Tangential derivatives are obtained by +spectrally differentiating the regular plasma gradient and applying the +Frenet connection. The public tensors use field-component-first ordering, + +\[ +H_{ijk}=\partial_j\partial_k B_i. +\] + +Subtracting the plasma Hessian from the regular-coordinate total Hessian gives +the external target. In a vacuum neighborhood this tensor is fully symmetric +and trace-free: + +\[ +H^c_{ijk}=H^c_{(ijk)},\qquad H^c_{iik}=0. +\] + +`external_hessian_independent` packs its seven Cartesian STF components in the +order `(xxx, xxy, xxz, xyy, xyz, yyy, yyz)`. The corresponding +`pack_symmetric_trace_free_rank3` and +`unpack_symmetric_trace_free_rank3` functions are reversible. + +Validation covers the circular finite-conductor curvature limit, QA and QH +oriented ellipses, vacuum reduction, spectral convergence of the Maxwell +identities, JIT/JVP agreement, and all rank-three tensor permutations. The +resolved-volume field test above validates the free-space source and matching +normalization. It is intentionally not replaced by differentiation of the +singular quadrature, which would omit the contact terms that the interior +potential supplies. + +As for the lower-order jet, `estimated_hessian_remainder` is a conservative +asymptotic scale rather than a rigorous bound. The predicted remainder is +\(O(a^2|\log a|)\) for fixed near-axis inputs. diff --git a/docs/theory/plasma_coil_separation.md b/docs/theory/plasma_coil_separation.md new file mode 100644 index 0000000..5aca6b6 --- /dev/null +++ b/docs/theory/plasma_coil_separation.md @@ -0,0 +1,26 @@ +# Plasma--coil separation + +The total on-axis jet is decomposed as + +\[ +(\boldsymbol B,D,H)_{\mathrm{tot}} += +(\boldsymbol B,D,H)_p ++ +(\boldsymbol B,D,H)_c. +\] + +The plasma part is the matched free-space field generated by the +positive-volume near-axis current source. The external part is obtained by +subtraction and represented by 3 field components, 5 independent +symmetric-trace-free gradient components, and 7 independent fully symmetric +trace-free Hessian components. + +A positive `formal_radius` is mandatory: it sets enclosed current and the +slender-channel matching scale. Surface-free means no finite-radius surface +is constructed; it does not mean radius-free. Field and Hessian remainder +metadata are asymptotic scales, not rigorous error bounds. + +The [plasma-current source](plasma-current.md) and +[matched field derivation](plasma-field.md) give the equations and validation +limits. ESSOS consumes the external target; pyQSC_JAX never imports ESSOS. diff --git a/docs/theory/second-order.md b/docs/theory/second-order.md new file mode 100644 index 0000000..78def4d --- /dev/null +++ b/docs/theory/second-order.md @@ -0,0 +1,131 @@ +# Finite-pressure/current second order + +## Scope + +Setting `order="r2"` computes the complete second-order surface coefficients +and \(B_{20}\) system of {cite}`landreman2019highorder`. The implementation +follows the staged dependencies in manuscript equations 36–54 and the +independently audited pyQSC implementation, while replacing mutable row-wise +assembly with pure JAX array operations. + +The pressure and current inputs use + +\[ +p = p_0 + r^2 p_2 + O(r^4), +\qquad +I = r^2 I_2 + O(r^4). +\] + +Both may be nonzero. `B2c` and `B2s` prescribe the constant second-harmonic +field-strength coefficients. + +## Algebraic stages + +Define + +\[ +V_1=X_{1c}^2+Y_{1c}^2+Y_{1s}^2,\qquad +V_2=2Y_{1s}Y_{1c},\qquad +V_3=X_{1c}^2+Y_{1c}^2-Y_{1s}^2. +\] + +The tangent-direction coefficients are obtained directly: + +\[ +Z_{20}=-\frac{B_0}{8|G_0|}V_1', +\] + +\[ +Z_{2s}=-\frac{B_0}{8|G_0|} +\left(V_2'-2\iota_N V_3\right), +\qquad +Z_{2c}=-\frac{B_0}{8|G_0|} +\left(V_3'+2\iota_N V_2\right), +\] + +where a prime means \(\mathrm d/\mathrm d\varphi\). These quantities determine +\(X_{2s}\) and \(X_{2c}\) algebraically from the second-harmonic +field-strength equations. + +The pressure response is + +\[ +\beta_{1s} +=-\frac{4s_\psi s_G\mu_0 p_2\bar\eta |G_0|} +{\iota_N B_0^3}. +\] + +The two remaining periodic unknowns are stacked as + +\[ +u=(X_{20},Y_{20}) +\] + +and satisfy a dense collocation system \(A u=b\). `A` contains the Fourier +derivative blocks and pointwise couplings; it is assembled from four +vectorized blocks, with no Python loop over grid rows. After the solve, +\(Y_{2s}\) and \(Y_{2c}\) follow from the two area constraints. + +## Implicit linear differentiation + +The primal callback uses `jax.numpy.linalg.solve`. SOLVAX wraps the matrix +action, primal solve, transpose action, and transpose solve with a custom +linear solve. JVPs and VJPs therefore differentiate the equation \(Au=b\), +not the internals of the dense factorization. + +`linear_report` records: + +- absolute and relative infinity-norm residuals; +- matrix condition number; +- finite, converged, and well-conditioned flags. + +The default report marks condition numbers above \(10^{12}\) as poorly +conditioned. Small \(|\iota_N|\) is a known physical conditioning risk and +must be assessed through this report rather than a warning emitted inside +JIT-compiled code. + +## \(B_{20}\) and direct diagnostics + +The nonconstant second-order field strength is + +\[ +B_{20}=B_0\left[ +\kappa X_{20} +-\frac{B_0}{|G_0|}Z_{20}' ++\frac{\bar\eta^2}{2} +-\frac{\mu_0p_2}{B_0^2} +-\frac{B_0^2}{4|G_0|^2} +\left(q_c^2+q_s^2+r_c^2+r_s^2\right) +\right], +\] + +with the standard first-order combinations \(q_c,q_s,r_c,r_s\). The result +includes the axis-length-weighted mean, anomaly, normalized weighted +\(L^2\) residual, and peak-to-peak variation. It also provides + +\[ +G_2=-\frac{\mu_0p_2G_0}{B_0^2}-\iota I_2. +\] + +All second-order coefficient derivatives needed by the field tensor and +diagnostics are stored. Untwisted zero- and second-harmonic +coefficients are used by the ESSOS-compatible boundary conversion. + +## Validation + +Four equations are evaluated independently of matrix assembly: two coupled +force-balance equations and two area constraints. Their maximum residual is +tested directly. Upstream regression cases cover: + +- the vacuum QA configuration of section 5.1; +- the finite-pressure/current configuration of section 5.3; +- the QH configuration of section 5.4; +- an asymmetric-axis development comparison. + +JIT, VMAP, JVP/VJP consistency, finite differences, and legacy r2 boundary +availability are also tested. Reference values are tied to the audited pyQSC +commit recorded in the refactor baseline. + +An r2 solve also computes the regular-coordinate total-field Hessian and the +leading Mercier quantities described in [](field-jet.md). Singular-radius +diagnostics remain a separate validation gate. diff --git a/docs/theory/second_order.md b/docs/theory/second_order.md new file mode 100644 index 0000000..c0c186a --- /dev/null +++ b/docs/theory/second_order.md @@ -0,0 +1,12 @@ +# Second order + +Second order supplies the complete finite-pressure/current coefficient system, +including the coupled periodic `(X20, Y20)` solve, `G2`, `beta_1s`, and +`B20`. The dense linear system is solved once for the primal result and +differentiated implicitly through SOLVAX. + +See the [complete second-order derivation](second-order.md) for the ordered +solution steps, residual equations, condition report, and pyQSC references. +The implementation is `pyqsc_jax.second_order`; independent residual and AD +tests are in `tests/physics/test_second_order.py` and +`tests/numerics/test_second_order_autodiff.py`. diff --git a/docs/theory/third-order.md b/docs/theory/third-order.md new file mode 100644 index 0000000..c055bf4 --- /dev/null +++ b/docs/theory/third-order.md @@ -0,0 +1,132 @@ +# Third-order flux constraint + +## Scope + +Setting `order="r3"` first constructs the complete r2 solution, then adds the +third-order surface displacement required for consistency of the toroidal-flux +constraint through order \(r^2\). This is the `r3_flux_constraint` construction +of {cite}`landreman2019highorder`. + +The implementation uses the short regular-coordinate relation as its primal +formula. It also evaluates the independent field-strength cancellation +formula retained by pyQSC; their agreement is reported instead of assumed. + +## Flux constraint + +Let \(\ell'=|G_0|/B_0\), let a prime denote +\(\mathrm d/\mathrm d\varphi\), and define + +\[ +\begin{aligned} +Q={}&-\frac{s_\psi B_0\ell'}{2G_0^2} +\left(\iota_N I_2+\frac{\mu_0p_2G_0}{B_0^2}\right) ++2(X_{2c}Y_{2s}-X_{2s}Y_{2c})\\ +&+\frac{s_\psi B_0}{2G_0} +\left(\ell'\kappa X_{20}-Z_{20}'\right)\\ +&+\frac{I_2}{4G_0} +\left[ +-\ell'\tau(X_{1c}^2+Y_{1s}^2+Y_{1c}^2) ++Y_{1c}X_{1c}'-X_{1c}Y_{1c}' +\right]. +\end{aligned} +\] + +The scalar coefficient at every toroidal sample is + +\[ +C=-\frac{Q}{2s_Gs_\psi}. +\] + +For the quasisymmetric flux-constraint construction, the nonzero Frenet-frame +coefficients are + +\[ +X_{3c1}=X_{1c}C,\qquad +Y_{3s1}=Y_{1s}C,\qquad +Y_{3c1}=Y_{1c}C. +\] + +\(X_{3s1}\), all \(Z_{3}\) coefficients, and all third-poloidal-harmonic +coefficients vanish in this construction. They remain explicit arrays in +`ThirdOrderData` so boundary assembly has a uniform representation and can be +extended without changing its interface. + +## Independent consistency checks + +The separately evaluated on-axis field correction must satisfy + +\[ +B_{0,\mathrm{cancel}}^{(2)}=2B_0C. +\] + +`flux_constraint_residual` is the maximum absolute residual of +\(Q+2s_Gs_\psi C=0\). `consistency_error` is the maximum disagreement between +\(C\) and \(B_{0,\mathrm{cancel}}^{(2)}/(2B_0)\). Both are scalar JAX arrays +and remain available inside compiled and differentiated workflows. + +## Helicity and boundary data + +The coefficient pairs are rotated from the winding Frenet frame into the +untwisted frame for first and third poloidal harmonics. The +ESSOS-compatible boundary map includes these terms multiplied by \(r^3\). +This is essential for QH axes: a coefficient that vanishes in the winding +frame can acquire a nonzero sine component after untwisting. + +## Validation + +Frozen upstream comparisons use the vacuum QA, finite-pressure/current, and QH +configurations of sections 5.1, 5.3, and 5.4 of +{cite}`landreman2019highorder`. Tests also cover both constraint residuals, +spectral coefficient derivatives, helical untwisting, r3 boundary assembly, +JIT, JVP, and centered finite differences. The upstream commit is recorded in +the refactor baseline. + +## Magnetic shear + +The generalized sigma equation gives the next rotational-transform term + +\[ +\iota(r)=\iota_0+r^2\iota_2+O(r^4) +\] + +as a periodic solvability condition {cite}`rodriguez2022weakly`. The source +derivation uses \(\epsilon=\sqrt{\psi}\), whereas the direct-construction +variables use \(r=\sqrt{2\psi/B_0}\). `solve_magnetic_shear` performs this +conversion explicitly before evaluating the third-order \(Z_{31}\), +\(X_{31}\), and \(Y_{31}\) relations. + +Writing \(\widetilde{\Lambda}\) for the inhomogeneous term and + +\[ +\mathcal E(\varphi)= +\exp\left(2\iota_N\int_0^\varphi\sigma(v)\,\mathrm dv\right), +\] + +the stellarator-symmetric result has the quotient form + +\[ +\iota_2= +\frac{B_0}{2} +\frac{\int \mathcal E\,\widetilde{\Lambda}\,\mathrm d\phi} +{\int \mathcal E +(X_{1c}^2+Y_{1c}^2+Y_{1s}^2)Y_{1s}^{-2}\,\mathrm d\phi}. +\] + +For a general asymmetric first-order solution, \(\sigma\) can have a nonzero +mean. The implementation separates the periodic and secular parts before +forming \(\mathcal E\), appends the analytically consistent endpoint, and uses +trapezoidal integration over one field period. This avoids applying a periodic +inverse derivative to a nonperiodic primitive. + +`B31c` is an explicit scalar input in the inverse-\(B^2\) convention. +`ShearData` records the numerator, denominator, integrating factor, +\(\widetilde{\Lambda}\), and third-order intermediates as well as `iota2`. +The current implementation is the standard-MHS specialization with +\(B_{31s}=0\), \(I_4=0\), and \(s_G=s_\psi=1\); other sign conventions are +rejected rather than silently extrapolated. + +Validation covers upstream pyQSC values for the section 5.1 vacuum QA, +section 5.3 finite-pressure/current, and section 5.4 QH configurations; an +asymmetric axis and nonzero sigma mean; resolution convergence; `B31c` +dependence; JIT, JVP, VJP, and differentiation through the complete near-axis +solve. diff --git a/docs/theory/third_order.md b/docs/theory/third_order.md new file mode 100644 index 0000000..dab2a03 --- /dev/null +++ b/docs/theory/third_order.md @@ -0,0 +1,11 @@ +# Third order + +Third order imposes the toroidal-flux constraint needed for pyQSC-compatible +boundary harmonics and exposes an explicit magnetic-shear calculation. +`order="r3"` produces first- and third-poloidal-harmonic coefficients and +two consistency residuals. `solve_magnetic_shear` is separate because +`B31c` is an independent input. + +The [third-order and shear derivation](third-order.md) records equations, +supported sign conventions, and upstream parity cases. Validation is in +`tests/physics/test_third_order.py` and `tests/physics/test_shear.py`. diff --git a/docs/tutorials/finite_beta_coils.md b/docs/tutorials/finite_beta_coils.md new file mode 100644 index 0000000..3c2cd0a --- /dev/null +++ b/docs/tutorials/finite_beta_coils.md @@ -0,0 +1,19 @@ +# Finite-beta coil target + +At finite current, coils must match the external vacuum jet, not the total +near-axis field. The target is + +\[ +(\boldsymbol B,D,H)_c=(\boldsymbol B,D,H)_{\mathrm{tot}} +-(\boldsymbol B,D,H)_p. +\] + +ESSOS PR [#46](https://github.com/uwplasma/ESSOS/pull/46) contains executable +stage-two and single-stage examples. The normalized loss uses smooth +least-squares blocks for 3 field, 5 STF-gradient, and 7 STF-Hessian +components. Tests compare actual coil-shape/current gradients and +axis/`etabar` gradients with finite differences. + +The formal radius must be supplied explicitly. Its asymptotic remainder +metadata and ratio to the curvature radius should be checked before accepting +a finite-beta target. diff --git a/docs/tutorials/first_order_qa.md b/docs/tutorials/first_order_qa.md new file mode 100644 index 0000000..9e19942 --- /dev/null +++ b/docs/tutorials/first_order_qa.md @@ -0,0 +1,14 @@ +# First-order QA + +This example constructs the standard three-field-period quasi-axisymmetric +case, prints convergence evidence and geometry metrics, and saves an axis +plot. Change parameters only at the top of the script. + +```{literalinclude} ../../examples/01_first_order_qa.py +:language: python +:linenos: +``` + +The reported transform is positive for the chosen negative `etabar` branch. +Always check `root_report.converged` and the residual before using derived +quantities in an objective. diff --git a/docs/tutorials/first_order_qh.md b/docs/tutorials/first_order_qh.md new file mode 100644 index 0000000..cf73371 --- /dev/null +++ b/docs/tutorials/first_order_qh.md @@ -0,0 +1,12 @@ +# First-order QH + +The QH example has nonzero Frenet-frame helicity. Consequently `iota` and +`iotaN` differ by an integer multiple of the field-period count. + +```{literalinclude} ../../examples/02_first_order_qh.py +:language: python +:linenos: +``` + +The printed helicity is a topology diagnostic computed from the magnetic +axis, not a user-supplied QA/QH label. diff --git a/docs/tutorials/optimize_axis.md b/docs/tutorials/optimize_axis.md new file mode 100644 index 0000000..766fde7 --- /dev/null +++ b/docs/tutorials/optimize_axis.md @@ -0,0 +1,14 @@ +# Optimize an axis + +This deliberately small deterministic example exercises bounded exploration, +exact \(B_{2c}\) elimination, local refinement, basin clustering, and status +semantics. + +```{literalinclude} ../../examples/06_optimize_axis_B20.py +:language: python +:linenos: +``` + +Increase the search budget, activate modes gradually, and retain doubled- and +quadrupled-grid verification for research searches. `best_found` never means +“globally optimal.” diff --git a/docs/tutorials/second_order_finite_beta.md b/docs/tutorials/second_order_finite_beta.md new file mode 100644 index 0000000..3aa4215 --- /dev/null +++ b/docs/tutorials/second_order_finite_beta.md @@ -0,0 +1,14 @@ +# Second-order finite beta and current + +This example activates pressure, current, and the complete r2 system. It +prints nonlinear and linear residuals, the matrix condition number, Mercier +quantity, \(B_{20}\) residual, and singular radius. + +```{literalinclude} ../../examples/03_second_order_finite_beta.py +:language: python +:linenos: +``` + +The demonstrated configuration is a regression case, not an optimized stable +stellarator. A negative Mercier value is therefore reported rather than +hidden. diff --git a/docs/tutorials/target_iota.md b/docs/tutorials/target_iota.md new file mode 100644 index 0000000..84d5893 --- /dev/null +++ b/docs/tutorials/target_iota.md @@ -0,0 +1,13 @@ +# Target rotational transform + +The inverse solve holds \(\iota\) fixed and replaces it in the sigma state by +the log-magnitude of a sign-preserving `etabar`. + +```{literalinclude} ../../examples/04_target_iota.py +:language: python +:linenos: +``` + +The forward reconstruction checks the achieved transform. If +`branch_fold` is true, the local inverse is ill-conditioned and the +[continuation workflow](../advanced/continuation.md) should be used. diff --git a/docs/tutorials/vacuum_coils.md b/docs/tutorials/vacuum_coils.md new file mode 100644 index 0000000..0c8e2ee --- /dev/null +++ b/docs/tutorials/vacuum_coils.md @@ -0,0 +1,19 @@ +# Vacuum coil target + +Coil optimization lives downstream in ESSOS so pyQSC_JAX has no circular +dependency. In vacuum, `formal_radius=None` makes the external target exactly +the total near-axis field jet. ESSOS normalizes the field, STF gradient, and +STF Hessian residual blocks by reference field and length scales. + +The executable stage-two integration is maintained in the +[ESSOS finite-beta integration PR](https://github.com/uwplasma/ESSOS/pull/46). +Its vacuum-reduction test confirms that the new objective is identical to the +total target when `I2 = 0`. + +For a local paired checkout, install pyQSC_JAX first and then ESSOS: + +```bash +python -m pip install -e /path/to/pyQSC_JAX +python -m pip install -e /path/to/ESSOS +python /path/to/ESSOS/examples/optimize_coils_for_near_axis_vacuum.py +``` diff --git a/docs/tutorials/vmec_export.md b/docs/tutorials/vmec_export.md new file mode 100644 index 0000000..638f699 --- /dev/null +++ b/docs/tutorials/vmec_export.md @@ -0,0 +1,17 @@ +# Export a fixed-boundary VMEC input + +The exporter does not require VMEC. It constructs the near-axis surface, +inverts cylindrical toroidal angle on a uniform grid, projects all four +boundary coefficient families with a two-dimensional FFT, writes a +deterministic `&INDATA` file, and returns accuracy and timing diagnostics. +If the requested surface does not admit a converged cylindrical-angle +inversion, `to_vmec` raises before creating the file. + +```{literalinclude} ../../examples/13_vmec_export.py +:language: python +``` + +For an asymptotic transform check, start at a small radius. Increasing the +radius tests finite-radius behavior of the truncated near-axis surface as well +as the exporter. The measured convergence is documented in +{doc}`../validation/vmec`. diff --git a/docs/tutorials/vmex_equilibrium.md b/docs/tutorials/vmex_equilibrium.md new file mode 100644 index 0000000..0625759 --- /dev/null +++ b/docs/tutorials/vmex_equilibrium.md @@ -0,0 +1,58 @@ +# Differentiate radial equilibrium quantities with VMEX + +`to_vmec` writes a conventional VMEC2000 input. `to_vmex_problem` instead +constructs an in-memory [VMEX](https://github.com/uwplasma/vmex) problem and +keeps the boundary and profile parameters as JAX leaves. + +Install the optional dependency: + +```bash +python -m pip install 'pyqsc-jax[vmex,plot]' +``` + +The equilibrium and implicit-adjoint compile are intentionally opt-in for the +direct script: + +```bash +PYQSC_RUN_VMEX=1 python examples/14_vmex_radial_profiles.py +``` + +Without that environment variable the script exits successfully after +printing the activation command. This keeps the all-examples gate bounded +when VMEX happens to be installed; the compatibility workflow separately runs +the real vacuum and finite-beta solves. + +The complete executable example is: + +```{literalinclude} ../../examples/14_vmex_radial_profiles.py +:language: python +``` + +The returned radial quantities are: + +- full-mesh rotational transform `iota`; +- the native-sign `iota_vmec` profile for convention audits; +- VMEX's differentiable quasisymmetry-ratio residual at each requested + normalized toroidal-flux surface; +- the canonical scalar magnetic well + \((V'(0)-V'(1))/V'(0)\); +- aspect ratio, volume, magnetic energy, and thermal energy. + +The call to `jax.value_and_grad` differentiates the converged fixed point by +VMEX's implicit adjoint. It does not unroll equilibrium iterations and does +not finite-difference a saved `wout`. + +For an outer pyQSC_JAX optimization, keep the discrete VMEX problem fixed and +traceably remap each candidate: + +```python +parameters = problem.parameters_for(candidate_solution) +objective = qsc.vmex_radial_quantities( + problem, + parameters, +).quasisymmetry.sum() +``` + +This updates the boundary Fourier arrays, toroidal flux, scalar pressure +profile, and current profile while preserving the VMEX resolution and +topology. diff --git a/docs/validation/essos.md b/docs/validation/essos.md new file mode 100644 index 0000000..fbe3e18 --- /dev/null +++ b/docs/validation/essos.md @@ -0,0 +1,35 @@ +# ESSOS integration + +ESSOS consumes pyQSC_JAX; pyQSC_JAX has no ESSOS dependency. Integration is +developed on ESSOS branch `refactor/pyqsc-external-field-jet` in draft PR +[#46](https://github.com/uwplasma/ESSOS/pull/46). + +The validated integration provides: + +- exact vacuum reduction to the total field jet; +- normalized smooth residual blocks for field, 5 STF-gradient, and 7 + STF-Hessian components; +- finite-pressure, exactly zero-current targets that subtract the plasma jet; +- actual coil-shape and current gradients checked against finite differences; +- axis and `etabar` gradients checked against finite differences; +- unchanged legacy-adapter smoke coverage; +- executable vacuum stage-two, finite-beta stage-two, and single-stage + examples. + +The focused integration suite passes 13 tests, including both optimization +scripts in clean subprocesses. The demonstrations use the nonplanar, +finite-pressure database stellarator ID 52521 with exactly \(I_2=0\), +\(\lvert\iota\rvert>2.8\), and RMS axis torsion above +\(0.85\ {\rm m}^{-1}\). + +At deliberately small example budgets, the stage-two normalized objective +decreases from \(6.7102\times10^7\) to \(4.4677\times10^3\). Its field, +gradient, and Hessian block sums all decrease independently: +\(4.4786\to0.7913\), \(2.0187\times10^4\to72.3078\), and +\(2.6833\times10^8\to1.7578\times10^4\). The single-stage objective decreases +from \(6.7102\times10^7\) to \(3.0052\times10^4\); its corresponding blocks +decrease to \(0.6382\), \(137.638\), and \(1.1966\times10^5\). The +single-stage solve keeps \(I_2=0\), reaches +\(\lvert\iota\rvert=2.843\), and lowers the weighted \(B_{20}\) residual from +0.2930 to 0.2514. These short runs establish execution and gradient coupling, +not device-quality global optima. diff --git a/docs/validation/literature_cases.md b/docs/validation/literature_cases.md new file mode 100644 index 0000000..10c7de5 --- /dev/null +++ b/docs/validation/literature_cases.md @@ -0,0 +1,49 @@ +# Literature cases + +The validation matrix is tied to primary literature: + +- Garren--Boozer existence and magnetic-field-strength papers for the + near-axis framework + ([DOI 10.1063/1.859916](https://doi.org/10.1063/1.859916), + [DOI 10.1063/1.859915](https://doi.org/10.1063/1.859915)); +- Landreman--Sengupta direct construction and high-order equations + ([arXiv:1809.10233](https://arxiv.org/abs/1809.10233), + [arXiv:1908.10253](https://arxiv.org/abs/1908.10253)); +- Landreman's near-axis figures of merit + ([arXiv:2012.00865](https://arxiv.org/abs/2012.00865)); +- Rodríguez, Sengupta, and Bhattacharjee for magnetic shear + ([DOI 10.1063/5.0076583](https://doi.org/10.1063/5.0076583)); +- Curvo, Ferreira, and Jorge for the configurable screening profile + ([DOI 10.1017/S002237782400165X](https://doi.org/10.1017/S002237782400165X)); +- Giuliani et al. for single-stage near-axis coil optimization + ([arXiv:2010.02033](https://arxiv.org/abs/2010.02033)). + +The [physics traceability table](../development/physics-traceability.md) +connects every block to code and independent tests. Numerical constants are +never copied from prose without a test or local, checksummed reference. + +## Public stellarator-database cases + +The README design gallery is traceable to the +[University of Wisconsin stellarator database](https://stellarator.physics.wisc.edu/) +and its [interactive application](https://stellarator.physics.wisc.edu/app). +The app reports transform magnitude, while the JSON downloads retain the +signed convention. The lead QA case uses the requested +\(\lvert\iota\rvert\ge0.3\) gate; the three QH/optimization cases use +\(\lvert\iota\rvert\ge0.4\). + +| bundled name | public source | role | +| --- | --- | --- | +| `database_qa_139524` | [ID 139524](https://stellarator.physics.wisc.edu/app/plot/139524) | one-period QA showcase with helicity zero and \(\lvert\iota\rvert>0.3\) | +| `database_example_3` | [ID 3](https://stellarator.physics.wisc.edu/app/plot/3) | documented API/download example | +| `database_low_b20_57409` | [ID 57409](https://stellarator.physics.wisc.edu/app/plot/57409) | low-\(B_{20}\) optimization seed | +| `b20_optimized_good` | derived from [ID 57409](https://stellarator.physics.wisc.edu/app/plot/57409) | constrained eight-mode refinement | +| `database_large_singularity_107579` | [ID 107579](https://stellarator.physics.wisc.edu/app/plot/107579) | large-singular-radius example | + +Every case is solved again rather than trusting displayed database metrics. +Regression tests require all four displayed cases to pass the full +configurable Curvo profile at the stated transform gate. ID 139524 must also +retain helicity zero, one field period, exactly zero \(I_2\), finite pressure, +and RMS axis torsion above \(1\ \mathrm{m}^{-1}\). The derived ID-57409 +refinement retains the source ID and URL in `ReferenceConfiguration` +metadata. diff --git a/docs/validation/plasma_field.md b/docs/validation/plasma_field.md new file mode 100644 index 0000000..2f9c678 --- /dev/null +++ b/docs/validation/plasma_field.md @@ -0,0 +1,60 @@ +# Plasma-field validation + +The surface-free current and field implementation is gated by tests of: + +- exact covariant/enclosed-current conversion; +- regular positive-volume radial scaling; +- vacuum, zero-current, and pressure-only limits; +- straight circular and sheared elliptical channels; +- circular toroidal finite-part and local-induction limits; +- Ampère's law, divergence, gradient symmetry, and Hessian symmetry/traces; +- full-torus periodicity and matching-length cancellation; +- angular and toroidal convergence; +- JIT/JVP and centered finite differences; +- reversible 5- and 7-component STF representations; +- independent resolved-volume Biot--Savart field comparisons with predicted + \(a^4|\log a|\) scaling. + +The Hessian uses interior-potential contact terms. Naively differentiating a +singular filament quadrature would omit those terms and is not treated as an +independent reference. See the [matched field derivation](../theory/plasma-field.md) +and `tests/physics/test_plasma_*.py`. + +## Pressure-only stellarator showcase + +The README and every finite-beta example use `plasma_stellarator`, the public +[Wisconsin stellarator-database configuration 52521](https://stellarator.physics.wisc.edu/app/plot/52521). +It has exactly `I2=0`, finite `p2=-28248.188`, four field periods, and a +strongly nonplanar magnetic axis. At `nphi=121`, its RMS torsion is +\(0.978864\ \mathrm{m}^{-1}\), so this case cannot silently regress to a +planar tokamak axis. It passes every configurable Curvo criterion with the +stricter \(\lvert\iota\rvert\ge0.4\). + +| quantity | value | +| --- | ---: | +| \(I_2\) | exactly 0 | +| \(p_2\) | \(-2.82482\times10^4\ \mathrm{Pa/m^2}\) | +| minimum/mean/maximum \(\lvert B_p\rvert/\lvert B_\mathrm{tot}\rvert\) | 0.00176059 / 0.00188960 / 0.00202583 | +| plasma \(\lvert B\rvert\) peak-to-peak / mean | 0.140368 | +| enclosed toroidal current | exactly 0 A | +| \(\lvert\iota\rvert\) | 2.809271 | +| RMS axis torsion | \(0.978864\ \mathrm{m}^{-1}\) | +| formal radius | 0.15 m | +| computed singular radius | 0.391608 m | +| formal/singular radius | 0.383 | +| estimated Hessian remainder | \(5.54\times10^{-3}\ \mathrm{T/m^2}\) | + +The zero current, finite pressure, nonzero torsion, nonzero plasma field, +angular variation, and radius margin are regression assertions. The +publication plot shows total, plasma, and external Frenet components; +plotting only their norms would hide both the constant-\(B_0\) total-field +construction and the cancellation of transverse plasma/external components. + +The pressure-only plasma contribution is about \(0.19\%\), not 30%. A scan +over screened database configurations and pressure multipliers found that +forcing a current-visible 30% fraction with \(I_2=0\) is incompatible with +the regular, criterion-passing near-axis examples considered here. The +earlier 30% showcase obtained its scale from finite \(I_2\) and appeared +tokamak-like, so it has been removed. `plasma_dominant_channel` is retained +only as a circular analytic current-normalization regression; it is not a +showcase stellarator or a recommended design. diff --git a/docs/validation/pyqsc_parity.md b/docs/validation/pyqsc_parity.md new file mode 100644 index 0000000..1f015f8 --- /dev/null +++ b/docs/validation/pyqsc_parity.md @@ -0,0 +1,40 @@ +# pyQSC parity + +Reference data were generated from landreman/pyQSC commit +`cd753596fd64babfb3832d2a676ca5b80b324b66` and stored locally with +license attribution and checksums. Tests never download regression data. + +The parity suite covers: + +- axis position, \(\varphi\), curvature, torsion, and helicity; +- sigma, transform, elongation, field and gradient; +- complete vacuum QA, finite-pressure/current, and QH r2 coefficients; +- \(B_{20}\), Mercier, Hessian, and singular radius; +- r3 boundary coefficients and flux-constraint checks; +- magnetic shear reference cases; +- historical adapter shapes, properties, `x`, `dofs`, and boundary behavior. + +pyQSC arrays are converted once at the compatibility boundary; canonical +pyQSC_JAX arrays keep the sample axis first. The upstream audit and exact +tolerances are implemented in `tests/regression`, `tests/physics`, and +`tests/compatibility`. + +The upstream project is available at +[landreman/pyQSC](https://github.com/landreman/pyQSC). + +## Independent high-resolution rerun + +At `nphi=121`, the current implementation was compared directly in one +process with the audited upstream checkout for vacuum QA, +finite-pressure/current, and QH cases. Across all three cases: + +- scalar transform and \(B_{20}\) diagnostics agreed within + \(8.1\times10^{-12}\); +- sigma, curvature, and torsion agreed within \(3.4\times10^{-14}\); +- the largest absolute difference among `X20`, `Y20`, and `B20` arrays was + \(1.11\times10^{-11}\). + +The comparison used upstream commit +`cd75359ea47548d5db7ccb458c100085c04ba1bc`, JAX 64-bit mode, and rebuilt +both implementations at the stated resolution rather than comparing with +interpolated frozen data. diff --git a/docs/validation/vmec.md b/docs/validation/vmec.md new file mode 100644 index 0000000..172771b --- /dev/null +++ b/docs/validation/vmec.md @@ -0,0 +1,101 @@ +# VMEC validation + +## Conversion algorithm + +`uniform_cylindrical_surface` solves the cylindrical-toroidal angle inversion +for the entire poloidal/toroidal grid with six vectorized JAX Newton +steps. Its derivative is obtained by JVP, so no finite-difference step is +selected. `vmec_boundary` reports the maximum residual, a dtype-aware default +tolerance of 100 machine epsilons, and a convergence flag. `to_vmec` raises +before writing if that flag is false; this prevents a large, non-star-shaped +or otherwise failed surface inversion from looking like a valid VMEC input. +The tolerance can be tightened explicitly. + +After the angle solve, `vmec_boundary` computes `RBC`, `RBS`, `ZBC`, and +`ZBS` simultaneously with a two-dimensional FFT. This replaces the legacy +scalar dense-root solve at every surface point and the direct Fourier +quadrature over every retained mode. + +At `nphi=61`, `ntheta=40`, `mpol=12`, and `ntor=14`, the new surface agrees +with the independent legacy root-based conversion to \(1.6\times10^{-14}\) m. +At radius 0.03 m, its retained-spectrum reconstruction errors are +\(3.42\times10^{-6}\) m in \(R\) and \(2.26\times10^{-6}\) m in \(Z\); +the maximum toroidal-angle residual is \(2.22\times10^{-16}\) rad. + +## Local equilibrium check + +The checked-in validation case was run with local VMEC2000 9.0 on macOS +arm64. It uses + +- `rc=[1.0, 0.045]`, `zs=[0.0, -0.045]`, `nfp=3`; +- `etabar=-0.9`, `B0=1`, `nphi=121`, order `r2`; +- export radius \(r=0.0025\) m; +- `ntheta=32`, `mpol=6`, `ntor=6`, `ns=31`. + +| Quantity | Near axis | VMEC | +| --- | ---: | ---: | +| on-axis \(\iota\) | 0.4183069102 | 0.4185430697 | +| relative difference | — | 0.05646% | +| `fsqr` | — | \(7.59\times10^{-11}\) | +| `fsqz` | — | \(4.21\times10^{-11}\) | +| `fsql` | — | \(2.40\times10^{-11}\) | + +The fixed +[`wout_qa_r0025.nc`](../../tests/reference/vmec/wout_qa_r0025.nc) +and its [manifest](../../tests/reference/vmec/manifest.json) are read on every +test run. The test checks the checksum, normal termination, force residuals, +and on-axis transform. Set +`PYQSC_VMEC_EXECUTABLE=/path/to/xvmec` to enable the second integration test, +which regenerates the input and reruns VMEC in a temporary directory. + +## Finite-pressure, zero-current database check + +A second opt-in local test exercises the user-facing finite-beta regime rather +than a finite-current channel. It uses Wisconsin stellarator-database QA +[ID 139524](https://stellarator.physics.wisc.edu/app/plot/139524), which has +finite \(p_2=-2.32744\times10^5\ \mathrm{Pa/m^2}\), exactly \(I_2=0\), +\(|\iota|=0.354802\), and RMS axis torsion +\(1.189\ \mathrm{m}^{-1}\). The export uses \(r=0.0015\) m, +`nphi=241`, `ntheta=40`, `mpol=8`, `ntor=8`, and radial stages +`ns=(31, 61)`. + +| Quantity | Near axis | VMEC | +| --- | ---: | ---: | +| on-axis \(\iota\) | \(-0.35480218\) | \(-0.35472615\) | +| relative difference | — | 0.02143% | +| on-axis pressure | 0.523673 Pa | 0.523673 Pa | +| enclosed current | exactly 0 A | exactly 0 A input | +| maximum force residual | — | \(9.97\times10^{-12}\) | + +This small radius is intentional: the transform comparison is an asymptotic +on-axis validation, while finite-beta radial behavior at practical boundary +radii is exercised by the VMEX integration. + +## Radius convergence + +The on-axis transform comparison is asymptotic. The same low-resolution VMEC +setup gives: + +| radius [m] | VMEC on-axis \(\iota\) | relative difference | +| ---: | ---: | ---: | +| 0.0025 | 0.41854307 | 0.0565% | +| 0.0050 | 0.41922558 | 0.2196% | +| 0.0100 | 0.42224025 | 0.9403% | +| 0.0200 | 0.43955604 | 5.0798% | +| 0.0300 | 0.49070165 | 17.3066% | + +The first three points show the expected approximately \(O(r^2)\) error. The +larger-radius results are intentionally retained: agreement at a small radius +does not imply that a truncated near-axis boundary is accurate arbitrarily far +from the axis. + +## Performance + +On the Apple M4 development machine, conversion at the higher +`40 × 61`, `mpol=12`, `ntor=14` resolution took 0.476 s including JAX +compilation and 5.20 ms after compilation. These are synchronized local +measurements, not cross-platform pass/fail promises. The unit test uses a +generous 0.5 s warm-call ceiling to catch a return to per-point root solving. +Raw data and the runnable command are in +`benchmarks/reports/2026-07-30-b20-vmec-apple-m4.json` and +`benchmarks/benchmark_vmec_export.py`. diff --git a/docs/validation/vmex_interface.md b/docs/validation/vmex_interface.md new file mode 100644 index 0000000..f2a5632 --- /dev/null +++ b/docs/validation/vmex_interface.md @@ -0,0 +1,85 @@ +# Differentiable VMEX integration + +The bridge was audited and exercised against +[uwplasma/VMEX](https://github.com/uwplasma/vmex) version 0.3.0 at commit +`2a40d7566be083070ea3ea534fa5d1fc44ad733a`. A compatibility CI job checks +the interface against current VMEX `main` on every pyQSC_JAX pull request. + +## Tested contract + +The integration uses VMEX's public APIs: + +- `VmecInput`; +- `implicit.params_from_input` and `implicit.run`; +- `implicit.iota_profile`; +- `optimize.QuasisymmetryRatioResidual.profile_state`; +- `optimize.magnetic_well`, `aspect_ratio`, and `volume`. + +The fixed-boundary solve and derived quantities remain inside the JAX +transformation. Unit tests also differentiate through +`problem.parameters_for(candidate_solution)`, which includes pyQSC_JAX's +surface conversion and pressure/current normalization. +The same diagnosed cylindrical-angle conversion used by `to_vmec` is used +in-memory. Concrete nonconvergence raises, while a traced optimization that +leaves the valid conversion domain receives nonfinite boundary leaves instead +of a silently incorrect equilibrium. + +The live compatibility test covers: + +| Case | Pressure | Current | Axis | Checked quantities | +| --- | ---: | ---: | --- | --- | +| vacuum QA | zero | zero | nonplanar QA | \(\iota(s)\), QS profile, well, implicit boundary gradient | +| `plasma_stellarator` (database ID 52521) | finite | exactly zero | nonplanar, RMS torsion \(0.979\ \mathrm{m}^{-1}\) | \(\iota(s)\), QS profile, well, nonzero thermal energy, implicit pressure/boundary gradients | + +Both live cases, including their implicit gradients, passed locally against +the exact current VMEX `main` commit above on 2026-07-30. The same source +checkout is installed by the compatibility CI job rather than relying on a +stale wheel. + +At the deliberately small `ns=7`, `mpol=4`, `ntor=2` smoke resolution, the +vacuum VMEX axis value in the pyQSC_JAX sign convention is `-0.421520`, +compared with the near-axis value `-0.420473`. This is a 0.249% difference; +the test gate is 0.8%. The purpose of this low-resolution job is cross-package +AD compatibility, not a production equilibrium accuracy claim. + +The strongly shaped finite-beta adjoint uses an explicit +`adjoint_tol=1e-8`; the public API retains VMEX's stricter \(10^{-11}\) +default unless the caller makes that tradeoff explicitly. VMEX enforces the +selected residual tolerance and raises rather than returning an unconverged +adjoint. + +The publication example records finite-beta values of: + +| Quantity | Value | +| --- | ---: | +| magnetic well | \(1.84561\times10^{-2}\) | +| \(\partial W/\partial p_\mathrm{scale}\) | \(+6.22590\times10^{-5}\) | +| \(\|\partial W/\partial RBC\|_2\) | \(9.78080\times10^1\) | + +These values are finite, nonzero regression evidence that the pressure and +boundary adjoint paths are active. + +## Coordinate convention + +VMEC's native transform for boundaries written by pyQSC_JAX has the opposite +sign from pyQSC_JAX's toroidal-angle convention. `VmexRadialQuantities` +therefore exposes both: + +- `iota_vmec`: VMEX native sign; +- `iota`: the sign aligned with `NearAxisSolution.iota`. + +The conversion is explicit and covered by tests. + +## Scope limits + +- The bridge differentiates VMEX's fixed-boundary equilibrium. It does not + claim an adjoint of the reconverged free-boundary NESTOR root. +- VMEX's traceable quasisymmetry profile currently requires stellarator + symmetry. Asymmetric equilibria can be solved with `qs_surfaces=()`. +- The magnetic well is VMEX's canonical endpoint scalar. No undocumented + radial “well profile” is invented. +- Pressure/current profiles are the simple near-axis-consistent profiles used + by `to_vmec`; users can replace the returned VMEX parameter leaves when a + different finite-radius profile model is intended. +- Production results require radial, angular, and Fourier-resolution studies + beyond the fast compatibility settings above. diff --git a/examples/01_first_order_qa.py b/examples/01_first_order_qa.py new file mode 100644 index 0000000..7f39e15 --- /dev/null +++ b/examples/01_first_order_qa.py @@ -0,0 +1,42 @@ +"""Construct and plot a first-order quasi-axisymmetric configuration.""" + +from pathlib import Path + +import matplotlib.pyplot as plt + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import plot_axis + +RC = [1.0, 0.045] +ZS = [0.0, -0.045] +NFP = 3 +ETABAR = -0.9 +NPHI = 31 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/01_first_order_qa.png") + +print("Solving first-order QA configuration...") +solution = qsc.Qsc( + rc=RC, + zs=ZS, + nfp=NFP, + etabar=ETABAR, + nphi=NPHI, + order="r1", +) +print("iota:", float(solution.iota)) +print("sigma residual:", float(solution.root_report.residual_norm)) +print("axis length [m]:", float(solution.axis_length)) +print("maximum elongation:", float(solution.elongation.max())) + +figure, _ = plot_axis(solution, label="first-order QA") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/02_first_order_qh.py b/examples/02_first_order_qh.py new file mode 100644 index 0000000..4546fac --- /dev/null +++ b/examples/02_first_order_qh.py @@ -0,0 +1,42 @@ +"""Construct and plot a first-order quasi-helically symmetric configuration.""" + +from pathlib import Path + +import matplotlib.pyplot as plt + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import plot_axis + +RC = [1.0, 0.265] +ZS = [0.0, -0.21] +NFP = 4 +ETABAR = -0.9 +NPHI = 31 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/02_first_order_qh.png") + +print("Solving first-order QH configuration...") +solution = qsc.Qsc( + rc=RC, + zs=ZS, + nfp=NFP, + etabar=ETABAR, + nphi=NPHI, + order="r1", +) +print("iota:", float(solution.iota)) +print("iota_N:", float(solution.iotaN)) +print("frame helicity:", int(solution.helicity)) +print("sigma residual:", float(solution.root_report.residual_norm)) + +figure, _ = plot_axis(solution, label="first-order QH", color="tab:orange") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/03_second_order_finite_beta.py b/examples/03_second_order_finite_beta.py new file mode 100644 index 0000000..88bdf85 --- /dev/null +++ b/examples/03_second_order_finite_beta.py @@ -0,0 +1,41 @@ +"""Solve the complete finite-pressure, zero-current second-order system.""" + +from pathlib import Path + +import matplotlib.pyplot as plt + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import plot_b20 + +CONFIGURATION = "plasma_stellarator" +NPHI = 61 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/03_second_order_finite_beta.png") + +print("Solving the finite-pressure, zero-current stellarator...") +solution = qsc.solve_configuration( + CONFIGURATION, + nphi=NPHI, + order="r2", +) +assert float(solution.inputs.I2) == 0.0 +assert float(solution.inputs.p2) != 0.0 +print("iota:", float(solution.iota)) +print("torsion RMS [1/m]:", float((solution.torsion**2).mean() ** 0.5)) +print("linear residual:", float(solution.linear_report.residual_norm)) +print("linear condition number:", float(solution.linear_report.matrix_condition_number)) +print("B20 weighted residual:", float(solution.B20_residual)) +print("Mercier D r^2:", float(solution.DMerc_times_r2)) +print("singular radius [m]:", float(solution.r_singularity)) + +figure, _ = plot_b20(solution, label="finite pressure, zero current", color="tab:red") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/04_target_iota.py b/examples/04_target_iota.py new file mode 100644 index 0000000..007859c --- /dev/null +++ b/examples/04_target_iota.py @@ -0,0 +1,54 @@ +"""Prescribe rotational transform and solve for etabar.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +AXIS = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) +TARGET_IOTA = 0.42 +ETABAR_SEED = -1.0 +NPHI = 31 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/04_target_iota.png") + +print("Solving inverse target-iota problem...") +inverse = qsc.solve( + axis=AXIS, + etabar=ETABAR_SEED, + iota=TARGET_IOTA, + solve_for="etabar", + nphi=NPHI, +) +forward = qsc.solve( + axis=AXIS, + etabar=inverse.inputs.etabar, + nphi=NPHI, +) +print("target iota:", TARGET_IOTA) +print("solved etabar:", float(inverse.inputs.etabar)) +print("forward-check iota:", float(forward.iota)) +print("response d(iota)/d(etabar):", float(inverse.response_derivative)) +print("branch fold:", bool(inverse.branch_fold)) + +figure, axis = plt.subplots(figsize=(6.0, 3.6)) +axis.plot( + np.asarray(inverse.varphi), + np.asarray(inverse.sigma), + linewidth=2, +) +axis.set_xlabel("Boozer toroidal angle [rad]") +axis.set_ylabel(r"$\sigma$") +axis.set_title("Target-iota periodic solution") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/05_optimal_B2c.py b/examples/05_optimal_B2c.py new file mode 100644 index 0000000..3335811 --- /dev/null +++ b/examples/05_optimal_B2c.py @@ -0,0 +1,53 @@ +"""Eliminate the affine B2c subproblem exactly.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +RC = [1.0, 0.155, 0.0102] +ZS = [0.0, 0.154, 0.0111] +NFP = 2 +ETABAR = 0.64 +INITIAL_B2C = 0.0 +NPHI = 61 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/05_optimal_B2c.png") + +print("Solving the initial r2 configuration...") +initial = qsc.Qsc( + rc=RC, + zs=ZS, + nfp=NFP, + etabar=ETABAR, + B2c=INITIAL_B2C, + nphi=NPHI, + order="r2", +) +result = qsc.optimize_B2c(initial) +optimized = result.solution +print("initial B2c:", INITIAL_B2C) +print("optimal B2c:", float(result.B2c_optimal)) +print("initial weighted L2:", float(qsc.b20_diagnostics(initial).weighted_l2)) +print("optimal weighted L2:", float(result.diagnostics.weighted_l2)) +print("affine reconstruction error:", float(result.affine_reconstruction_error)) + +angle = np.asarray(initial.varphi * NFP / (2 * np.pi)) +figure, axis = plt.subplots(figsize=(6.2, 3.8)) +axis.plot(angle, np.asarray(initial.B20_anomaly), label="initial", linewidth=2) +axis.plot(angle, np.asarray(optimized.B20_anomaly), label="optimal B2c", linewidth=2) +axis.set_xlabel("Boozer angle / field period") +axis.set_ylabel(r"$B_{20}-\langle B_{20}\rangle$ [T/m$^2$]") +axis.legend() +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/06_optimize_axis_B20.py b/examples/06_optimize_axis_B20.py new file mode 100644 index 0000000..c889be8 --- /dev/null +++ b/examples/06_optimize_axis_B20.py @@ -0,0 +1,59 @@ +"""Refine a screened database axis into a nearly constant-B20 stellarator.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +STOCK_CONFIGURATION = "database_low_b20_57409" +OPTIMIZED_CONFIGURATION = "b20_optimized_good" +NPHI = 121 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/06_optimize_axis_B20.png") + +print("Eliminating B2c exactly for database configuration 57409...") +stock = qsc.optimize_B2c(qsc.solve_configuration(STOCK_CONFIGURATION, nphi=NPHI)) +print("Loading the independently optimized Fourier axis...") +optimized = qsc.solve_configuration(OPTIMIZED_CONFIGURATION, nphi=NPHI) +optimized_diagnostics = qsc.b20_diagnostics(optimized) +verification = qsc.verify_B20_resolution(optimized, multipliers=(1, 2)) +improvement = float(stock.diagnostics.weighted_l2 / optimized_diagnostics.weighted_l2) + +print("stock exact-B2c weighted L2:", float(stock.diagnostics.weighted_l2)) +print("optimized-axis weighted L2:", float(optimized_diagnostics.weighted_l2)) +print("optimized-axis dense maximum:", float(optimized_diagnostics.grid_maximum)) +print("improvement factor:", improvement) +print("nphi verification:", np.asarray(verification.resolutions)) +print("verified weighted L2:", np.asarray(verification.weighted_l2)) +criteria = qsc.Criteria.from_curvo_2025(minimum_abs_iota=0.4) +print("Curvo profile passed:", criteria.evaluate(optimized).passed) +print("|iota|:", abs(float(optimized.iota))) +print("singular radius:", float(optimized.r_singularity)) + +stock_angle = np.asarray(stock.solution.varphi * stock.solution.inputs.axis.nfp / (2 * np.pi)) +optimized_angle = np.asarray(optimized.varphi * optimized.inputs.axis.nfp / (2 * np.pi)) +figure, axes = plt.subplots(1, 2, figsize=(9.4, 3.8)) +axes[0].plot(stock_angle, np.asarray(stock.diagnostics.anomaly), linewidth=2) +axes[0].set_title(r"database ID 57409 + exact $B_{2c}$") +axes[0].set_xlabel("Boozer angle / field period") +axes[0].set_ylabel(r"$B_{20}-\langle B_{20}\rangle$ [T/m$^2$]") +axes[1].plot( + optimized_angle, + np.asarray(optimized_diagnostics.anomaly), + linewidth=2, + color="tab:green", +) +axes[1].set_title(f"optimized axis ({improvement:,.0f}x smaller)") +axes[1].set_xlabel("Boozer angle / field period") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/07_good_stellarator_search.py b/examples/07_good_stellarator_search.py new file mode 100644 index 0000000..a498584 --- /dev/null +++ b/examples/07_good_stellarator_search.py @@ -0,0 +1,82 @@ +"""Apply the Curvo screen during a reproducible zero-current stellarator search.""" + +from pathlib import Path + +import jax.numpy as jnp +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +CONFIGURATION = "plasma_stellarator" +I2 = 0.0 +NPHI = 31 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/07_good_stellarator_search.png") + +configuration = qsc.get_configuration(CONFIGURATION) +axis = qsc.Axis( + rc=configuration.rc, + zs=configuration.zs, + nfp=configuration.nfp, +) +variable_indices = qsc.stellarator_symmetric_variable_indices(axis, modes=(3,)) +criteria = qsc.Criteria.from_curvo_2025(minimum_abs_iota=0.4) +problem = qsc.AxisSearchProblem( + axis=axis, + variable_indices=variable_indices, + lower_bounds=jnp.asarray([0.002, 0.002]), + upper_bounds=jnp.asarray([0.008, 0.008]), + etabar=configuration.etabar, + I2=I2, + p2=configuration.p2, + nphi=NPHI, + criteria=criteria, + selector="maximum_singular_radius", +) +options = qsc.AxisSearchOptions( + coarse_samples=2, + local_starts=1, + maximum_iterations=2, + verification_multipliers=(1, 2), + verification_tail_tolerance=1.0, +) + +print("Running criteria-aware search from database ID 52521...") +result = qsc.search_axis(problem, options=options) +if result.best is None or result.best.criteria_report is None: + raise RuntimeError("The criteria-aware example did not produce a report.") +report = result.best.criteria_report +print("status:", result.status) +print("criteria passed:", report.passed) +print("I2:", float(result.best.solution.inputs.I2)) +print("p2:", float(result.best.solution.inputs.p2)) +print("iota:", float(result.best.solution.iota)) +print("singular radius [m]:", float(result.best.solution.r_singularity)) +print("torsion RMS [1/m]:", float((result.best.solution.torsion**2).mean() ** 0.5)) +for evaluation in report.evaluations: + print( + evaluation.name, + "value=", + float(evaluation.value), + "margin=", + float(evaluation.margin), + evaluation.units, + ) + +names = [evaluation.name for evaluation in report.evaluations] +margins = np.asarray([evaluation.margin for evaluation in report.evaluations]) +figure, axis = plt.subplots(figsize=(7.4, 4.0)) +axis.barh(names, margins, color=np.where(margins >= 0, "tab:green", "tab:red")) +axis.axvline(0, color="black", linewidth=1) +axis.set_xlabel("signed criterion margin") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/08_boundary_and_coordinates.py b/examples/08_boundary_and_coordinates.py new file mode 100644 index 0000000..b67b80c --- /dev/null +++ b/examples/08_boundary_and_coordinates.py @@ -0,0 +1,68 @@ +"""Generate an available-order boundary through the compatibility adapter.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +from pyqsc_jax.near_axis import near_axis + +RC = [1.0, 0.155, 0.0102] +ZS = [0.0, 0.154, 0.0111] +NFP = 2 +ETABAR = 0.64 +B2C = -0.00322 +RADIUS = 0.05 +NTHETA = 24 +NPHI = 49 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/08_boundary_and_coordinates.png") + +print("Constructing an r2 compatibility surface...") +field = near_axis( + rc=RC, + zs=ZS, + nfp=NFP, + etabar=ETABAR, + B2c=B2C, + nphi=31, + order="r2", +) +x, y, z, radius = field.get_boundary( + r=RADIUS, + ntheta=NTHETA, + nphi=NPHI, +) +print("boundary array shape:", x.shape) +print("R range [m]:", float(radius.min()), float(radius.max())) +print("Z range [m]:", float(z.min()), float(z.max())) + +figure = plt.figure(figsize=(8.0, 3.8)) +axis_3d = figure.add_subplot(1, 2, 1, projection="3d") +axis_3d.plot_surface( + np.asarray(x), + np.asarray(y), + np.asarray(z), + cmap="viridis", + alpha=0.8, + linewidth=0, +) +axis_3d.set_box_aspect((1, 1, 1)) +axis_3d.set_xlabel("x [m]") +axis_3d.set_ylabel("y [m]") +axis_3d.set_zlabel("z [m]") +axis_cross = figure.add_subplot(1, 2, 2) +axis_cross.plot(np.asarray(radius[:, 0]), np.asarray(z[:, 0]), linewidth=2) +axis_cross.set_aspect("equal") +axis_cross.set_xlabel("R [m]") +axis_cross.set_ylabel("Z [m]") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/09_total_field_jet.py b/examples/09_total_field_jet.py new file mode 100644 index 0000000..5eb92f9 --- /dev/null +++ b/examples/09_total_field_jet.py @@ -0,0 +1,50 @@ +"""Inspect the total on-axis field, gradient, and Hessian.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +CONFIGURATION = "qa" +NPHI = 61 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/09_total_field_jet.png") + +print("Solving the total-field-jet reference configuration...") +solution = qsc.solve_configuration(CONFIGURATION, nphi=NPHI) +jet = solution.field_jet +if jet is None: + raise RuntimeError("The r2 solution did not produce a total field jet.") +print("field reconstruction error:", float(jet.maximum_field_error)) +print("gradient reconstruction error:", float(jet.maximum_gradient_error)) +print("maximum divergence:", float(jet.maximum_divergence)) +print("maximum Hessian derivative asymmetry:", float(jet.maximum_derivative_asymmetry)) +print("maximum gradient of divergence:", float(jet.maximum_divergence_gradient)) +print("minimum Hessian scale length [m]:", float(jet.L_grad_grad_B.min())) + +angle = np.asarray(solution.varphi * solution.inputs.axis.nfp / (2 * np.pi)) +hessian = np.asarray(jet.hessian) +figure, axes = plt.subplots(2, 1, figsize=(7.0, 6.0), sharex=True) +axes[0].plot(angle, np.asarray(solution.B_axis[:, 0]), label=r"$B_x$") +axes[0].plot(angle, np.asarray(solution.B_axis[:, 1]), label=r"$B_y$") +axes[0].plot(angle, np.asarray(solution.B_axis[:, 2]), label=r"$B_z$") +axes[0].set_ylabel("field [T]") +axes[0].legend(ncol=3) +axes[1].plot(angle, hessian[:, 0, 0, 0], label=r"$\partial_{xx}B_x$") +axes[1].plot(angle, hessian[:, 0, 1, 1], label=r"$\partial_{yy}B_x$") +axes[1].plot(angle, hessian[:, 0, 2, 2], label=r"$\partial_{zz}B_x$") +axes[1].set_xlabel("Boozer angle / field period") +axes[1].set_ylabel(r"Hessian component [T/m$^2$]") +axes[1].legend(ncol=3) +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/10_plasma_external_field_jet.py b/examples/10_plasma_external_field_jet.py new file mode 100644 index 0000000..7b02ac3 --- /dev/null +++ b/examples/10_plasma_external_field_jet.py @@ -0,0 +1,58 @@ +"""Separate total, plasma, and external vacuum field jets.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import plot_field_split_components + +CONFIGURATION = "plasma_stellarator" +FORMAL_RADIUS = 0.15 +NPHI = 61 +ANGULAR_RESOLUTION = 96 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/10_plasma_external_field_jet.png") + +print("Solving the finite-pressure, zero-current stellarator...") +solution = qsc.solve_configuration(CONFIGURATION, nphi=NPHI) +assert float(solution.inputs.I2) == 0.0 +assert float(solution.inputs.p2) != 0.0 +print("torsion RMS [1/m]:", float((solution.torsion**2).mean() ** 0.5)) +print("Separating the surface-free plasma and external jets...") +result = qsc.plasma_hessian_on_axis( + solution, + formal_radius=FORMAL_RADIUS, + angular_resolution=ANGULAR_RESOLUTION, +) +enclosed_current = result.field.field.current_source.enclosed_toroidal_current +plasma_fraction = ( + (result.field.field.field**2).sum(axis=-1) / (solution.B_axis**2).sum(axis=-1) +) ** 0.5 +print("enclosed toroidal current [A]:", float(enclosed_current)) +print("minimum |B_plasma| / |B_total|:", float(plasma_fraction.min())) +print("mean |B_plasma| / |B_total|:", float(plasma_fraction.mean())) +plasma_norm = np.linalg.norm(np.asarray(result.field.field.field), axis=-1) +print("plasma |B| peak-to-peak / mean:", float(np.ptp(plasma_norm) / np.mean(plasma_norm))) +print("formal radius / singular radius:", float(FORMAL_RADIUS / solution.r_singularity)) +print("external gradient STF components:", result.field.external_gradient_independent.shape[-1]) +print("external Hessian STF components:", result.external_hessian_independent.shape[-1]) +print("maximum external-gradient trace:", float(result.field.maximum_external_trace)) +print("maximum external-Hessian trace:", float(result.maximum_external_trace)) +print("estimated plasma-field remainder [T]:", float(result.field.field.estimated_field_remainder)) +print( + "estimated plasma-Hessian remainder [T/m^2]:", + float(result.estimated_hessian_remainder), +) + +figure, _ = plot_field_split_components(result, solution) +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/11_continuation_scan.py b/examples/11_continuation_scan.py new file mode 100644 index 0000000..256bc40 --- /dev/null +++ b/examples/11_continuation_scan.py @@ -0,0 +1,51 @@ +"""Trace an etabar branch with pseudo-arclength continuation.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +AXIS = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) +ETABAR_START = -1.0 +ETABAR_NEXT = -0.98 +NUMBER_OF_POINTS = 12 +NPHI = 31 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/11_continuation_scan.png") + +print("Tracing the first-order etabar branch...") +result = qsc.continue_etabar_branch( + axis=AXIS, + etabar_start=ETABAR_START, + etabar_next=ETABAR_NEXT, + num_points=NUMBER_OF_POINTS, + nphi=NPHI, +) +etabar = np.asarray(result.etabar) +iota = np.asarray(result.iota) +fold = np.asarray(result.fold_detected) +print("status:", result.status) +print("points:", len(result.solutions)) +print("etabar interval:", float(etabar.min()), float(etabar.max())) +print("iota interval:", float(iota.min()), float(iota.max())) +print("detected fold samples:", np.flatnonzero(fold)) + +figure, axis = plt.subplots(figsize=(6.0, 3.8)) +axis.plot(etabar, iota, "-o", label="corrected branch") +if fold.any(): + axis.scatter(etabar[fold], iota[fold], marker="x", s=80, label="fold diagnostic") +axis.set_xlabel(r"$\bar{\eta}$ [m$^{-1}$]") +axis.set_ylabel(r"$\iota$") +axis.legend() +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/12_autodiff_check.py b/examples/12_autodiff_check.py new file mode 100644 index 0000000..5d66890 --- /dev/null +++ b/examples/12_autodiff_check.py @@ -0,0 +1,53 @@ +"""Compare an implicit-solve JVP with a centered finite difference.""" + +from pathlib import Path + +import jax +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +AXIS = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) +ETABAR = -0.9 +NPHI = 31 +FINITE_DIFFERENCE_STEP = 1.0e-5 +SAVE_OUTPUT = True +SHOW_FIGURE = False +OUTPUT = Path("examples/output/12_autodiff_check.png") + + +def iota_from_etabar(etabar): + """Return one converged implicit-solve output.""" + + return qsc.solve(axis=AXIS, etabar=etabar, nphi=NPHI).iota + + +print("Differentiating rotational transform through the converged sigma solve...") +value, tangent = jax.jvp(iota_from_etabar, (ETABAR,), (1.0,)) +upper = iota_from_etabar(ETABAR + FINITE_DIFFERENCE_STEP) +lower = iota_from_etabar(ETABAR - FINITE_DIFFERENCE_STEP) +finite_difference = (upper - lower) / (2 * FINITE_DIFFERENCE_STEP) +relative_error = abs(tangent - finite_difference) / max(abs(finite_difference), 1.0e-14) +print("iota:", float(value)) +print("JVP d(iota)/d(etabar):", float(tangent)) +print("finite-difference derivative:", float(finite_difference)) +print("relative error:", float(relative_error)) + +figure, axis = plt.subplots(figsize=(5.5, 3.6)) +axis.bar( + ("JAX implicit JVP", "centered difference"), + np.asarray((tangent, finite_difference)), + color=("tab:blue", "tab:orange"), +) +axis.set_ylabel(r"$d\iota/d\bar{\eta}$ [m]") +axis.set_title(f"relative error = {float(relative_error):.2e}") +figure.tight_layout() +if SAVE_OUTPUT: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/13_vmec_export.py b/examples/13_vmec_export.py new file mode 100644 index 0000000..342a5c6 --- /dev/null +++ b/examples/13_vmec_export.py @@ -0,0 +1,77 @@ +"""Export a diagnosed VMEC boundary without invoking VMEC.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +RADIUS = 0.005 +NPHI = 61 +NTHETA = 40 +MPOL = 12 +NTOR = 14 +SAVE_OUTPUT = True +SHOW_FIGURE = False +VMEC_INPUT = Path("examples/output/input.pyqsc_jax_qa") +FIGURE_OUTPUT = Path("examples/output/13_vmec_export.png") + +solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=NPHI, + order="r2", +) +print("Writing the fixed-boundary VMEC input...") +export = qsc.to_vmec( + solution, + VMEC_INPUT, + r=RADIUS, + ntheta=NTHETA, + mpol=MPOL, + ntor=NTOR, +) +boundary = export.boundary +print("input:", export.path) +print("conversion seconds (compile + first execution):", export.conversion_seconds) +print("toroidal-angle inversion converged:", bool(boundary.toroidal_angle_converged)) +print("maximum toroidal-angle residual:", float(boundary.maximum_toroidal_angle_residual)) +print("toroidal-angle tolerance:", float(boundary.toroidal_angle_tolerance)) +print("maximum R reconstruction error [m]:", float(boundary.maximum_R_reconstruction_error)) +print("maximum Z reconstruction error [m]:", float(boundary.maximum_Z_reconstruction_error)) +print("run with: xvmec", export.path.name) + +figure, axes = plt.subplots(1, 2, figsize=(9.4, 3.8)) +phi_indices = np.linspace(0, NPHI - 1, 7, dtype=int) +for index in phi_indices: + axes[0].plot( + np.asarray(boundary.R[:, index]), + np.asarray(boundary.Z[:, index]), + linewidth=1.3, + ) +axes[0].set_aspect("equal") +axes[0].set_xlabel("R [m]") +axes[0].set_ylabel("Z [m]") +axes[0].set_title("uniform cylindrical-toroidal sections") +mode_amplitude = np.sqrt( + np.asarray(boundary.RBC) ** 2 + + np.asarray(boundary.RBS) ** 2 + + np.asarray(boundary.ZBC) ** 2 + + np.asarray(boundary.ZBS) ** 2 +) +axes[1].semilogy(np.sort(mode_amplitude.ravel())[::-1], marker=".", linewidth=1) +axes[1].set_xlabel("coefficient rank") +axes[1].set_ylabel("combined Fourier amplitude [m]") +axes[1].set_title("VMEC boundary spectrum") +figure.tight_layout() +if SAVE_OUTPUT: + FIGURE_OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(FIGURE_OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", FIGURE_OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/14_vmex_radial_profiles.py b/examples/14_vmex_radial_profiles.py new file mode 100644 index 0000000..97fe7fb --- /dev/null +++ b/examples/14_vmex_radial_profiles.py @@ -0,0 +1,81 @@ +"""Solve differentiable radial equilibrium quantities with optional VMEX.""" + +import os +from pathlib import Path + +import jax +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +CONFIGURATION = "qa" +RADIUS = 0.02 +QS_SURFACES = (0.25, 0.5, 0.75, 1.0) +OUTPUT = Path("examples/output/14_vmex_radial_profiles.png") +SAVE_FIGURE = True +SHOW_FIGURE = False + +if os.environ.get("PYQSC_RUN_VMEX") != "1": + print("Set PYQSC_RUN_VMEX=1 to run the optional VMEX equilibrium example.") + raise SystemExit(0) + +print("Constructing the near-axis boundary...") +solution = qsc.solve_configuration(CONFIGURATION, nphi=31) +try: + problem = qsc.to_vmex_problem( + solution, + r=RADIUS, + qs_surfaces=QS_SURFACES, + ntheta=8, + mpol=3, + ntor=2, + ns_array=(7,), + ftol=1.0e-7, + max_iterations=1200, + multigrid=False, + ) +except ImportError as error: + print(error) + print("Skipping the optional VMEX equilibrium example.") + raise SystemExit(0) from None + +print("Solving the converged fixed-boundary equilibrium with VMEX...") +result = problem.solve() +quantities = result.quantities +print("VMEX version:", problem.vmex_version) +print("iota(s):", np.asarray(quantities.iota)) +print("quasisymmetry profile:", np.asarray(quantities.quasisymmetry)) +print("magnetic well:", float(quantities.magnetic_well)) + +print("Differentiating magnetic well through the converged VMEX fixed point...") +well, gradient = jax.value_and_grad( + lambda parameters: qsc.vmex_radial_quantities(problem, parameters).magnetic_well +)(problem.parameters) +print("magnetic well:", float(well)) +print("d(well)/d(pres_scale):", float(gradient.pres_scale)) +print("||d(well)/d(RBC)||:", float(np.linalg.norm(np.asarray(gradient.rbc)))) + +figure, axes = plt.subplots(1, 2, figsize=(9.4, 3.8)) +axes[0].plot(np.asarray(quantities.s), np.asarray(quantities.iota), "-o") +axes[0].axhline(float(solution.iota), color="black", linestyle="--", label="near axis") +axes[0].set_xlabel(r"normalized toroidal flux $s$") +axes[0].set_ylabel(r"$\iota(s)$") +axes[0].legend() +axes[1].semilogy( + np.asarray(quantities.qs_surfaces), + np.asarray(quantities.quasisymmetry), + "-o", +) +axes[1].set_xlabel(r"normalized toroidal flux $s$") +axes[1].set_ylabel("VMEX QS residual") +figure.tight_layout() + +if SAVE_FIGURE: + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(OUTPUT, dpi=180, bbox_inches="tight") + print("saved:", OUTPUT) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_B20_optimization.py b/examples/publication/figure_B20_optimization.py new file mode 100644 index 0000000..47e9a93 --- /dev/null +++ b/examples/publication/figure_B20_optimization.py @@ -0,0 +1,154 @@ +"""Publication figure for exact B2c elimination and full-axis B20 optimization.""" + +import json +import subprocess +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import plot_surface_3d + +STOCK_CONFIGURATION = "database_low_b20_57409" +OPTIMIZED_CONFIGURATION = "b20_optimized_good" +NPHI = 121 +SURFACE_RADIUS = 0.075 +OUTPUT_STEM = Path("examples/output/publication/B20_optimization") +README_PNG = Path("docs/_static/B20_optimization.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +plt.style.use("seaborn-v0_8-whitegrid") +print("Refining the traceable low-B20 database configuration 57409...") +stock = qsc.optimize_B2c(qsc.solve_configuration(STOCK_CONFIGURATION, nphi=NPHI)) +optimized = qsc.solve_configuration(OPTIMIZED_CONFIGURATION, nphi=NPHI) +optimized_diagnostics = qsc.b20_diagnostics(optimized) +verification = qsc.verify_B20_resolution(optimized, multipliers=(1, 2)) +improvement = float(stock.diagnostics.weighted_l2 / optimized_diagnostics.weighted_l2) +stock_angle = np.asarray(stock.solution.varphi * stock.solution.inputs.axis.nfp / (2 * np.pi)) +optimized_angle = np.asarray(optimized.varphi * optimized.inputs.axis.nfp / (2 * np.pi)) + +figure = plt.figure(figsize=(13.2, 8.2)) +surface_axis = figure.add_subplot(2, 2, 1, projection="3d") +plot_surface_3d( + optimized, + radius=SURFACE_RADIUS, + ntheta=36, + ax=surface_axis, + cmap="viridis", +) +surface_axis.view_init(elev=24, azim=38) +surface_axis.set_title( + "database-seeded optimized QH surface\n" + rf"$|\iota|={abs(float(optimized.iota)):.3f}$, " + rf"$r_\mathrm{{sing}}={float(optimized.r_singularity):.3f}$ m" +) + +stock_axis = figure.add_subplot(2, 2, 2) +stock_axis.plot(stock_angle, np.asarray(stock.diagnostics.anomaly), linewidth=2.2) +stock_axis.set_title(r"database ID 57409 + exact $B_{2c}$") +stock_axis.set_xlabel("Boozer angle / field period") +stock_axis.set_ylabel(r"$B_{20}-\langle B_{20}\rangle$ [T/m$^2$]") + +optimized_axis = figure.add_subplot(2, 2, 3) +optimized_axis.plot( + optimized_angle, + np.asarray(optimized_diagnostics.anomaly), + linewidth=2.2, + color="tab:green", +) +optimized_axis.set_title("optimized Fourier axis") +optimized_axis.set_xlabel("Boozer angle / field period") +optimized_axis.set_ylabel(r"$B_{20}-\langle B_{20}\rangle$ [T/m$^2$]") +diagnostic_names = ("weighted $L^2$", "dense maximum", "peak-to-peak") +stock_values = ( + float(stock.diagnostics.weighted_l2), + float(stock.diagnostics.grid_maximum), + float(stock.diagnostics.peak_to_peak), +) +optimized_values = ( + float(optimized_diagnostics.weighted_l2), + float(optimized_diagnostics.grid_maximum), + float(optimized_diagnostics.peak_to_peak), +) +summary_axis = figure.add_subplot(2, 2, 4) +summary_axis.bar( + np.arange(3) - 0.18, + stock_values, + width=0.36, + label=r"exact $B_{2c}$ only", +) +summary_axis.bar( + np.arange(3) + 0.18, + optimized_values, + width=0.36, + label="axis + $B_{2c}$", + color="tab:green", +) +summary_axis.set_xticks(np.arange(3), diagnostic_names, rotation=12) +summary_axis.set_yscale("log") +summary_axis.set_ylabel(r"$B_{20}$ nonuniformity [T/m$^2$]") +summary_axis.legend(fontsize=8) +clearance = float(optimized.r_singularity) / SURFACE_RADIUS +summary_axis.text( + 0.04, + 0.05, + f"{improvement:,.0f}x lower residual\n" + "all Curvo criteria pass\n" + rf"shown surface: $r={SURFACE_RADIUS:.3f}$ m" + f" ({clearance:.1f}x inside singular radius)", + transform=summary_axis.transAxes, + fontweight="bold", + bbox={"facecolor": "white", "edgecolor": "0.7", "alpha": 0.9}, +) +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata = { + "git_commit": commit, + "stock_configuration": STOCK_CONFIGURATION, + "optimized_configuration": OPTIMIZED_CONFIGURATION, + "nphi": NPHI, + "stock_exact_B2c": float(stock.B2c_optimal), + "optimized_B2c": float(optimized.inputs.B2c), + "stock_weighted_l2": float(stock.diagnostics.weighted_l2), + "optimized_weighted_l2": float(optimized_diagnostics.weighted_l2), + "optimized_grid_maximum": float(optimized_diagnostics.grid_maximum), + "optimized_peak_to_peak": float(optimized_diagnostics.peak_to_peak), + "optimized_singular_radius": float(optimized.r_singularity), + "surface_radius": SURFACE_RADIUS, + "singular_radius_clearance": clearance, + "improvement_factor": improvement, + "verification_nphi": np.asarray(verification.resolutions).tolist(), + "verification_weighted_l2": np.asarray(verification.weighted_l2).tolist(), + "optimizer": "bounded scipy least_squares after exact B2c elimination", + "source_database_id": 57409, + "source_url": "https://stellarator.physics.wisc.edu/app/plot/57409", + "curvo_profile_minimum_abs_iota": 0.4, + "curvo_profile_passed": qsc.Criteria.from_curvo_2025(minimum_abs_iota=0.4) + .evaluate(optimized) + .passed, +} +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=120, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_QA_QH_branches.py b/examples/publication/figure_QA_QH_branches.py new file mode 100644 index 0000000..d69f024 --- /dev/null +++ b/examples/publication/figure_QA_QH_branches.py @@ -0,0 +1,92 @@ +"""Publication figure comparing QA and QH etabar branches.""" + +import json +import subprocess +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +BRANCHES = { + "QA": { + "axis": qsc.Axis( + rc=[1.0, -0.06883207, 0.0017516185, 0.023231717], + zs=[0.0, -0.28447896, 0.074662544, 0.07483574], + nfp=1, + ), + "etabar": np.linspace(-1.15, -0.4, 13), + "color": "tab:blue", + "database_id": 139524, + }, + "QH": { + "axis": qsc.Axis( + rc=[1.0, -0.53677857, -0.046455786, -0.0070183445], + zs=[0.0, -0.5888703, -0.04447083, -0.009581006], + nfp=4, + ), + "etabar": np.linspace(0.8, 1.8, 13), + "color": "tab:orange", + "database_id": 3, + }, +} +NPHI = 31 +OUTPUT_STEM = Path("examples/output/publication/QA_QH_branches") +README_PNG = Path("docs/_static/QA_QH_branches.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +plt.style.use("seaborn-v0_8-whitegrid") +figure, axes = plt.subplots(1, 2, figsize=(10.0, 3.9)) +metadata = {"nphi": NPHI, "branches": {}} +print("Sampling QA and QH first-order branches...") +for name, branch in BRANCHES.items(): + etabar_values = branch["etabar"] + solutions = [ + qsc.solve(axis=branch["axis"], etabar=etabar, nphi=NPHI) for etabar in etabar_values + ] + iota = np.asarray([float(solution.iota) for solution in solutions]) + elongation = np.asarray([float(solution.elongation.max()) for solution in solutions]) + axes[0].plot(etabar_values, iota, "-o", color=branch["color"], label=name) + axes[1].plot(etabar_values, elongation, "-o", color=branch["color"], label=name) + metadata["branches"][name] = { + "source_database_id": branch["database_id"], + "etabar": etabar_values.tolist(), + "iota": iota.tolist(), + "maximum_elongation": elongation.tolist(), + } +axes[0].set_xlabel(r"$\bar{\eta}$ [m$^{-1}$]") +axes[0].set_ylabel(r"$\iota$") +axes[0].axhline(0.0, color="0.3", linewidth=0.8) +axes[1].set_xlabel(r"$\bar{\eta}$ [m$^{-1}$]") +axes[1].set_ylabel("maximum elongation") +for axis in axes: + axis.legend() +figure.suptitle("QA and QH topology respond differently to the same design variable") +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata["git_commit"] = commit +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=130, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_axis_and_surfaces.py b/examples/publication/figure_axis_and_surfaces.py new file mode 100644 index 0000000..0ef3310 --- /dev/null +++ b/examples/publication/figure_axis_and_surfaces.py @@ -0,0 +1,91 @@ +"""Publication figure comparing QA and QH near-axis surfaces.""" + +import json +import subprocess +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +from pyqsc_jax.near_axis import near_axis + +CONFIGURATIONS = { + "QA": { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + }, + "QH": { + "rc": [1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + "zs": [0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + "nfp": 4, + "etabar": 1.569, + "B2c": 0.1348, + }, +} +RADIUS = 0.05 +NTHETA = 36 +NPHI = 121 +OUTPUT_STEM = Path("examples/output/publication/axis_and_surfaces") +README_PNG = Path("docs/_static/axis_and_surfaces.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +plt.style.use("seaborn-v0_8-whitegrid") +figure = plt.figure(figsize=(11.0, 4.6)) +metadata = {"formal_surface_radius": RADIUS, "configurations": CONFIGURATIONS} + +print("Rendering QA and QH available-order surfaces...") +for panel, (name, parameters) in enumerate(CONFIGURATIONS.items(), start=1): + solution = near_axis(**parameters, nphi=61, order="r2") + x, y, z, _ = solution.get_boundary(r=RADIUS, ntheta=NTHETA, nphi=NPHI) + axis = figure.add_subplot(1, 2, panel, projection="3d") + axis.plot_surface( + np.asarray(x), + np.asarray(y), + np.asarray(z), + cmap="viridis" if name == "QA" else "plasma", + linewidth=0, + antialiased=True, + alpha=0.88, + ) + axis.plot( + np.asarray(x).mean(axis=0), + np.asarray(y).mean(axis=0), + np.asarray(z).mean(axis=0), + color="black", + linewidth=2.0, + ) + axis.set_title(f"{name}: $\\iota={float(solution.iota):.3f}$") + axis.set_xlabel("x [m]") + axis.set_ylabel("y [m]") + axis.set_zlabel("z [m]") + axis.set_box_aspect((1, 1, 0.5)) +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata["git_commit"] = commit +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=120, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_convergence.py b/examples/publication/figure_convergence.py new file mode 100644 index 0000000..2bb176b --- /dev/null +++ b/examples/publication/figure_convergence.py @@ -0,0 +1,82 @@ +"""Publication figure for spectral resolution convergence.""" + +import json +import subprocess +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +CONFIGURATION = "database_qa_139524" +RESOLUTIONS = (15, 31, 61, 121) +OUTPUT_STEM = Path("examples/output/publication/convergence") +README_PNG = Path("docs/_static/convergence.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +plt.style.use("seaborn-v0_8-whitegrid") +print("Solving database QA ID 139524 on successively finer grids...") +solutions = [qsc.solve_configuration(CONFIGURATION, nphi=resolution) for resolution in RESOLUTIONS] +reference = solutions[-1] +iota_errors = np.asarray( + [abs(float(solution.iota - reference.iota)) for solution in solutions[:-1]] +) +b20_values = np.asarray( + [float(qsc.b20_diagnostics(solution).weighted_l2) for solution in solutions] +) +b20_errors = np.abs(b20_values[:-1] - b20_values[-1]) +floor = np.finfo(float).eps + +figure, axes = plt.subplots(1, 2, figsize=(9.8, 3.9)) +axes[0].loglog( + RESOLUTIONS[:-1], + np.maximum(iota_errors, floor), + "-o", + linewidth=2, +) +axes[0].set_xlabel("toroidal grid points") +axes[0].set_ylabel(r"$|\iota_N-\iota_{121}|$") +axes[1].loglog( + RESOLUTIONS[:-1], + np.maximum(b20_errors, floor), + "-o", + linewidth=2, + color="tab:red", +) +axes[1].set_xlabel("toroidal grid points") +axes[1].set_ylabel(r"$|R_{B20,N}-R_{B20,121}|$") +figure.suptitle(r"Resolution audit: $\iota$ converges early, $B_{20}$ needs a finer grid") +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata = { + "git_commit": commit, + "configuration": CONFIGURATION, + "resolutions": RESOLUTIONS, + "iota": [float(solution.iota) for solution in solutions], + "B20_weighted_l2": b20_values.tolist(), +} +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=120, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_core_performance.py b/examples/publication/figure_core_performance.py new file mode 100644 index 0000000..7bccdb4 --- /dev/null +++ b/examples/publication/figure_core_performance.py @@ -0,0 +1,145 @@ +"""Publication figure for measured JIT, JVP, and VMAP performance.""" + +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +REPOSITORY = Path(__file__).resolve().parents[2] +REPORT = REPOSITORY / "benchmarks" / "reports" / "2026-07-29-apple-m4.json" +OUTPUT_STEM = Path("examples/output/publication/core_performance") +README_PNG = Path("docs/_static/core_performance.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +report = json.loads(REPORT.read_text(encoding="utf-8")) +rows = sorted(report["results"], key=lambda row: row["nphi"]) +nphi = np.asarray([row["nphi"] for row in rows]) +cold_ms = np.asarray([1.0e3 * row["cold_compile_and_execute_seconds"] for row in rows]) +warm_ms = np.asarray( + [ + 1.0e3 + * row.get( + "warm_median_seconds", + float(np.median(row["warm_samples_seconds"])), + ) + for row in rows + ] +) +speedups = cold_ms / warm_ms +finest = rows[-1] +batch_size = int(finest["vmap_batch_size"]) +latencies_ms = np.asarray( + [ + 1.0e3 + * finest.get( + "warm_median_seconds", + float(np.median(finest["warm_samples_seconds"])), + ), + 1.0e3 + * finest.get( + "jvp_warm_median_seconds", + float(np.median(finest["jvp_warm_samples_seconds"])), + ), + 1.0e3 + * finest.get( + "vmap_warm_median_seconds", + float(np.median(finest["vmap_warm_samples_seconds"])), + ), + ] +) + +plt.style.use("seaborn-v0_8-whitegrid") +figure, axes = plt.subplots(1, 2, figsize=(10.8, 4.2)) +width = 5.5 +axes[0].bar( + nphi - width / 2, + cold_ms, + width=width, + label="compile + first execution", + color="#34495e", +) +axes[0].bar( + nphi + width / 2, + warm_ms, + width=width, + label="warm median", + color="#1abc9c", +) +axes[0].set_yscale("log") +axes[0].set_xlabel("toroidal grid points") +axes[0].set_ylabel("wall time [ms, log scale]") +axes[0].set_xticks(nphi) +axes[0].legend(frameon=True) +for x_value, warm_value, speedup in zip(nphi, warm_ms, speedups, strict=True): + axes[0].annotate( + f"{speedup:,.0f}×", + (x_value + width / 2, warm_value), + xytext=(0, 6), + textcoords="offset points", + ha="center", + va="bottom", + color="#087f6d", + fontweight="bold", + ) + +labels = ("solve", "solve + JVP", f"VMAP × {batch_size}") +colors = ("#3498db", "#9b59b6", "#f39c12") +bars = axes[1].bar(labels, latencies_ms, color=colors, width=0.62) +axes[1].set_ylabel("warm wall time [ms]") +axes[1].set_title(f"Composable transforms at $n_\\phi={finest['nphi']}$") +axes[1].bar_label(bars, fmt="%.3f ms", padding=4) +per_case_us = 1.0e3 * latencies_ms[-1] / batch_size +axes[1].text( + 0.05, + 0.92, + f"{batch_size} batched solutions\n{per_case_us:.1f} μs per case", + transform=axes[1].transAxes, + ha="left", + va="top", + color="#9c640c", + fontweight="bold", +) +axes[1].set_ylim(0.0, 1.25 * float(latencies_ms.max())) + +figure.suptitle("JAX compilation is amortized; derivatives and batches stay fast") +figure.text( + 0.5, + -0.01, + ( + f"{report.get('hardware', report['platform'])} • " + f"JAX {report['jax']} {report['jax_backend']} " + f"• x64={report['jax_enable_x64']} • synchronized medians" + ), + ha="center", + fontsize=8, + color="0.35", +) +figure.tight_layout() + +metadata = { + "report": str(REPORT.relative_to(REPOSITORY)), + "report_git_commit": report["git_commit"], + "cold_to_warm_speedup": dict( + zip((str(value) for value in nphi), speedups.tolist(), strict=True) + ), + "finest_nphi": int(finest["nphi"]), + "finest_warm_latency_ms": latencies_ms.tolist(), + "finest_vmap_per_case_microseconds": float(per_case_us), +} +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=130, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_optimizer_comparison.py b/examples/publication/figure_optimizer_comparison.py new file mode 100644 index 0000000..d51f5bd --- /dev/null +++ b/examples/publication/figure_optimizer_comparison.py @@ -0,0 +1,136 @@ +"""Publication figure comparing B20 optimizer strategies and the final refinement.""" + +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc + +REPOSITORY = Path(__file__).resolve().parents[2] +REPORT = REPOSITORY / "benchmarks" / "reports" / "2026-07-30-b20-optimizers-apple-m4.json" +OUTPUT_STEM = Path("examples/output/publication/optimizer_comparison") +README_PNG = Path("docs/_static/optimizer_comparison.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +report = json.loads(REPORT.read_text(encoding="utf-8")) +rows = report["results"] +production = qsc.solve_configuration("b20_optimized_good", nphi=report["nphi"]) +production_residual = float(qsc.b20_diagnostics(production).weighted_l2) +short_names = ( + r"exact $B_{2c}$", + "L-BFGS-B", + "least squares", + "differential\nevolution", + "multistart LM", + "staged 8-mode\nrefinement", +) +residuals = np.asarray([row["weighted_l2"] for row in rows] + [production_residual]) +colors = ("#7f8c8d", "#3498db", "#2ecc71", "#e74c3c", "#9b59b6", "#117864") + +plt.style.use("seaborn-v0_8-whitegrid") +figure, axes = plt.subplots(1, 2, figsize=(11.2, 4.45)) +bars = axes[0].bar(short_names, residuals, color=colors, width=0.72) +axes[0].set_yscale("log") +axes[0].set_ylabel(r"weighted $\|P B_{20}\|_2$ [T m$^{-2}$]") +axes[0].set_title("Same dense target, visibly different outcomes") +axes[0].tick_params(axis="x", rotation=20) +axes[0].bar_label( + bars, + labels=[f"{value:.1e}" for value in residuals], + padding=3, + fontsize=8, +) +axes[0].text( + 0.98, + 0.36, + f"final improvement\n{residuals[0] / production_residual:,.1e}×", + transform=axes[0].transAxes, + ha="right", + va="bottom", + color="#117864", + fontweight="bold", + bbox={"boxstyle": "round,pad=0.3", "facecolor": "white", "alpha": 0.85}, +) + +screened_rows = rows[1:] +evaluations = np.asarray([row["function_evaluations"] for row in screened_rows]) +screened_residuals = np.asarray([row["weighted_l2"] for row in screened_rows]) +seconds = np.asarray([row["seconds"] for row in screened_rows]) +point_colors = colors[1:5] +point_labels = ("L-BFGS-B", "least squares", "low-budget DE", "multistart LM") +sizes = 75.0 + 55.0 * np.log10(1.0 + seconds / seconds.min()) +for evaluation, residual, size, color, label, elapsed in zip( + evaluations, + screened_residuals, + sizes, + point_colors, + point_labels, + seconds, + strict=True, +): + axes[1].scatter( + evaluation, + residual, + s=size, + color=color, + edgecolor="white", + linewidth=1.0, + label=f"{label} • {elapsed:.2f} s", + zorder=3, + ) +axes[1].set_yscale("log") +axes[1].set_xlabel("objective evaluations") +axes[1].set_ylabel(r"weighted $\|P B_{20}\|_2$") +axes[1].set_title("Budget, basin coverage, and local accuracy") +axes[1].legend(fontsize=8, frameon=True) +axes[1].annotate( + "best local screen", + (evaluations[1], screened_residuals[1]), + xytext=(15, 30), + textcoords="offset points", + ha="left", + arrowprops={"arrowstyle": "->", "color": "#117864"}, + color="#117864", + fontweight="bold", +) + +figure.suptitle(r"Screen broadly, refine locally: $B_{20}$ becomes nearly constant") +figure.text( + 0.5, + -0.015, + ( + "Database ID 57409 • nphi=121 • exact B2c elimination in every axis " + "evaluation • final point uses staged Fourier continuation" + ), + ha="center", + fontsize=8, + color="0.35", +) +figure.tight_layout() + +metadata = { + "report": str(REPORT.relative_to(REPOSITORY)), + "report_git_commit": report["git_commit"], + "source_database_id": report["source_database_id"], + "production_configuration": "b20_optimized_good", + "production_weighted_l2": production_residual, + "production_improvement_over_exact_B2c": float(residuals[0] / production_residual), +} +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=130, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_plasma_external_jet.py b/examples/publication/figure_plasma_external_jet.py new file mode 100644 index 0000000..6bb9730 --- /dev/null +++ b/examples/publication/figure_plasma_external_jet.py @@ -0,0 +1,132 @@ +"""Publication figure for the total/plasma/external 3+5+7 field jet.""" + +import json +import subprocess +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import field_split_frenet_components, plot_surface_3d + +CONFIGURATION = "plasma_stellarator" +FORMAL_RADIUS = 0.15 +DISPLAY_RADIUS = 0.12 +NPHI = 121 +ANGULAR_RESOLUTION = 128 +OUTPUT_STEM = Path("examples/output/publication/plasma_external_jet") +README_PNG = Path("docs/_static/plasma_external_jet.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +plt.style.use("seaborn-v0_8-whitegrid") +print("Evaluating the surface-free plasma/external jet...") +solution = qsc.solve_configuration(CONFIGURATION, nphi=NPHI) +assert float(solution.inputs.I2) == 0.0 +assert float(solution.inputs.p2) != 0.0 +result = qsc.plasma_hessian_on_axis( + solution, + formal_radius=FORMAL_RADIUS, + angular_resolution=ANGULAR_RESOLUTION, +) +plasma_fraction = ( + (result.field.field.field**2).sum(axis=-1) / (solution.B_axis**2).sum(axis=-1) +) ** 0.5 +components = np.asarray(field_split_frenet_components(result, solution)) +angle = np.asarray(solution.varphi * solution.inputs.axis.nfp / (2 * np.pi)) +plasma_norm = np.linalg.norm(np.asarray(result.field.field.field), axis=1) +external_norm = np.linalg.norm(np.asarray(result.field.external_field), axis=1) + +figure = plt.figure(figsize=(12.8, 8.0)) +surface_axis = figure.add_subplot(2, 2, 1, projection="3d") +plot_surface_3d( + solution, + radius=DISPLAY_RADIUS, + ntheta=32, + ax=surface_axis, + cmap="plasma", +) +surface_axis.view_init(elev=25, azim=35) +surface_axis.set_title( + r"database stellarator: finite $p_2$, exactly $I_2=0$" + "\n" + rf"$|\iota|={abs(float(solution.iota)):.3f}$, " + rf"$\tau_\mathrm{{rms}}={float((solution.torsion**2).mean() ** 0.5):.3f}$ m$^{{-1}}$" +) + +fraction_axis = figure.add_subplot(2, 2, 2) +fraction_percent = 100.0 * np.asarray(plasma_fraction) +fraction_axis.plot(angle, fraction_percent, color="tab:red", linewidth=2.3) +fraction_axis.fill_between(angle, 0.0, fraction_percent, alpha=0.18, color="tab:red") +fraction_axis.set_xlabel("Boozer angle / field period") +fraction_axis.set_ylabel(r"$100|B_\mathrm{plasma}|/|B_\mathrm{total}|$ [%]") +fraction_axis.set_title("pressure-driven plasma contribution") + +tangent_axis = figure.add_subplot(2, 2, 3) +for field_index, label in enumerate(("total", "plasma", "external")): + tangent_axis.plot(angle, components[field_index, 0], label=label) +tangent_axis.set_xlabel("Boozer angle / field period") +tangent_axis.set_ylabel(r"$B_t$ [T]") +tangent_axis.set_title("tangent field component") +tangent_axis.legend(ncol=3) + +transverse_axis = figure.add_subplot(2, 2, 4) +transverse_axis.plot(angle, components[1, 1], label=r"plasma $B_n$") +transverse_axis.plot(angle, components[2, 1], "--", label=r"external $B_n$") +transverse_axis.plot(angle, components[1, 2], label=r"plasma $B_b$") +transverse_axis.plot(angle, components[2, 2], "--", label=r"external $B_b$") +transverse_axis.set_xlabel("Boozer angle / field period") +transverse_axis.set_ylabel("transverse field [T]") +transverse_axis.set_title("plasma and external parts cancel to the total") +transverse_axis.legend(ncol=2, fontsize=8) +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata = { + "git_commit": commit, + "configuration": CONFIGURATION, + "formal_radius": FORMAL_RADIUS, + "display_radius": DISPLAY_RADIUS, + "nphi": NPHI, + "angular_resolution": ANGULAR_RESOLUTION, + "I2": float(solution.inputs.I2), + "p2": float(solution.inputs.p2), + "torsion_minimum": float(solution.torsion.min()), + "torsion_maximum": float(solution.torsion.max()), + "torsion_rms": float((solution.torsion**2).mean() ** 0.5), + "enclosed_current_amperes": float(result.field.field.current_source.enclosed_toroidal_current), + "minimum_plasma_field_fraction": float(plasma_fraction.min()), + "mean_plasma_field_fraction": float(plasma_fraction.mean()), + "maximum_plasma_field_fraction": float(plasma_fraction.max()), + "plasma_norm_peak_to_peak_over_mean": float(np.ptp(plasma_norm) / np.mean(plasma_norm)), + "external_norm_peak_to_peak_over_mean": float(np.ptp(external_norm) / np.mean(external_norm)), + "singular_radius": float(solution.r_singularity), + "formal_radius_to_singular_radius": float(FORMAL_RADIUS / solution.r_singularity), + "maximum_external_gradient_trace": float(result.field.maximum_external_trace), + "maximum_external_hessian_trace": float(result.maximum_external_trace), + "field_remainder": float(result.field.field.estimated_field_remainder), + "hessian_remainder": float(result.estimated_hessian_remainder), +} +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=120, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_stellarator_gallery.py b/examples/publication/figure_stellarator_gallery.py new file mode 100644 index 0000000..ae07e8f --- /dev/null +++ b/examples/publication/figure_stellarator_gallery.py @@ -0,0 +1,142 @@ +"""Publication gallery of independently screened stellarator designs.""" + +import json +import subprocess +from pathlib import Path + +import matplotlib.pyplot as plt + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import plot_surface_3d + +CASES = ( + ( + "database_qa_139524", + r"QA • database ID 139524 • 1 field period", + 0.03, + "viridis", + 32, + 55, + 0.3, + ), + ( + "database_example_3", + r"QH • database ID 3 • 4 field periods", + 0.075, + "magma", + 23, + 35, + 0.4, + ), + ( + "b20_optimized_good", + r"optimization • nearly constant $B_{20}$", + 0.075, + "cividis", + 28, + 28, + 0.4, + ), + ( + "database_large_singularity_107579", + r"robust surface • large $r_\mathrm{sing}$", + 0.15, + "plasma", + 25, + 40, + 0.4, + ), +) +NPHI = 121 +OUTPUT_STEM = Path("examples/output/publication/stellarator_gallery") +README_PNG = Path("docs/_static/stellarator_gallery.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +plt.style.use("seaborn-v0_8-whitegrid") +figure = plt.figure(figsize=(12.0, 9.2)) +metadata = {"nphi": NPHI, "configurations": {}} + +print("Rendering the bundled stellarator gallery...") +for panel, ( + name, + title, + radius, + cmap, + elevation, + azimuth, + minimum_abs_iota, +) in enumerate(CASES, start=1): + solution = qsc.solve_configuration(name, nphi=NPHI) + criteria = qsc.Criteria.from_curvo_2025( + minimum_abs_iota=minimum_abs_iota, + ) + assert criteria.evaluate(solution).passed + axis = figure.add_subplot(2, 2, panel, projection="3d") + plot_surface_3d( + solution, + radius=radius, + ntheta=36, + ax=axis, + cmap=cmap, + ) + axis.view_init(elev=elevation, azim=azimuth) + if name == "b20_optimized_good": + diagnostic = rf"$\|P B_{{20}}\|_2={float(solution.B20_residual):.2e}$" + elif name == "database_large_singularity_107579": + diagnostic = rf"$r_\mathrm{{sing}}={float(solution.r_singularity):.3f}$ m" + elif name == "database_qa_139524": + diagnostic = rf"helicity $=0$, $|\iota|={abs(float(solution.iota)):.3f}$" + else: + diagnostic = rf"$|\iota|={abs(float(solution.iota)):.3f}$" + axis.set_title( + title + + "\n" + + diagnostic + + ( + "" + if name == "database_large_singularity_107579" + else rf", $r_\mathrm{{sing}}={float(solution.r_singularity):.3f}$ m" + ) + + "\nCurvo profile: pass" + ) + metadata["configurations"][name] = { + "surface_radius": radius, + "iota": float(solution.iota), + "B20_residual": float(solution.B20_residual), + "singular_radius": float(solution.r_singularity), + "surface_to_singular_radius": radius / float(solution.r_singularity), + "curvo_profile_minimum_abs_iota": minimum_abs_iota, + "curvo_profile_passed": True, + "source_database_id": qsc.get_configuration(name).source_database_id, + "source_url": qsc.get_configuration(name).source_url, + "I2": float(solution.inputs.I2), + "p2": float(solution.inputs.p2), + "torsion_rms": float((solution.torsion**2).mean() ** 0.5), + } +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata["git_commit"] = commit +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=130, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_vmec_validation.py b/examples/publication/figure_vmec_validation.py new file mode 100644 index 0000000..e33e49b --- /dev/null +++ b/examples/publication/figure_vmec_validation.py @@ -0,0 +1,88 @@ +"""Publication figure for VMEC export timing and on-axis-iota convergence.""" + +import json +import subprocess +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +REPORT = ( + Path(__file__).resolve().parents[2] + / "benchmarks" + / "reports" + / "2026-07-30-b20-vmec-apple-m4.json" +) +OUTPUT_STEM = Path("examples/output/publication/vmec_validation") +README_PNG = Path("docs/_static/vmec_validation.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +print("Loading the frozen VMEC validation report...") +report = json.loads(REPORT.read_text(encoding="utf-8")) +convergence = report["vmec_radius_convergence"] +radii = np.asarray([row["radius"] for row in convergence["results"]]) +errors = np.asarray([row["relative_iota_error"] for row in convergence["results"]]) +timing = report["vmec_export"] + +plt.style.use("seaborn-v0_8-whitegrid") +figure, axes = plt.subplots(1, 2, figsize=(9.8, 3.8)) +axes[0].loglog(radii, errors, "o-", linewidth=2.2, label="VMEC") +axes[0].loglog( + radii, + errors[0] * (radii / radii[0]) ** 2, + "--", + linewidth=1.7, + label=r"$O(r^2)$", +) +axes[0].set_xlabel("export radius [m]") +axes[0].set_ylabel("relative on-axis iota error") +axes[0].set_title("near-axis convergence") +axes[0].legend() +cold_ms = 1000 * timing["cold_compile_and_execute_seconds"] +warm_ms = 1000 * timing["warm_median_seconds"] +axes[1].bar(("compile + execute", "warm execute"), (cold_ms, warm_ms), color=("0.45", "tab:blue")) +axes[1].set_yscale("log") +axes[1].set_ylabel("conversion time [ms]") +axes[1].set_title("vectorized JAX exporter") +axes[1].text( + 0.5, + 0.07, + f"{timing['maximum_R_reconstruction_error_m'] * 1e6:.2f} µm max R error\n" + f"{timing['maximum_Z_reconstruction_error_m'] * 1e6:.2f} µm max Z error", + ha="center", + transform=axes[1].transAxes, +) +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata = { + "git_commit": commit, + "source_report": str(REPORT), + "near_axis_iota": convergence["near_axis_iota"], + "radius_convergence": convergence["results"], + "vmec_version": convergence["vmec_version"], + "export": timing, +} +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=120, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/examples/publication/figure_vmex_radial_profiles.py b/examples/publication/figure_vmex_radial_profiles.py new file mode 100644 index 0000000..1d46b3f --- /dev/null +++ b/examples/publication/figure_vmex_radial_profiles.py @@ -0,0 +1,167 @@ +"""Publication figure for the differentiable pyQSC_JAX-to-VMEX bridge.""" + +import json +import os +import subprocess +from pathlib import Path + +import jax +import matplotlib.pyplot as plt +import numpy as np + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import plot_surface_3d + +CASES = ( + ("qa", "vacuum QA"), + ("plasma_stellarator", r"finite $p_2$, $I_2=0$"), +) +RADIUS = 0.02 +QS_SURFACES = (0.2, 0.4, 0.6, 0.8, 1.0) +OUTPUT_STEM = Path("examples/output/publication/vmex_radial_profiles") +README_PNG = Path("docs/_static/vmex_radial_profiles.png") +SAVE_OUTPUT = True +SHOW_FIGURE = False + +if os.environ.get("PYQSC_RUN_VMEX") != "1": + print("Set PYQSC_RUN_VMEX=1 to regenerate the optional VMEX publication figure.") + raise SystemExit(0) + +print("Solving vacuum and finite-beta radial equilibria with VMEX...") +records = [] +try: + for configuration, label in CASES: + near_axis = qsc.solve_configuration(configuration, nphi=31) + problem = qsc.to_vmex_problem( + near_axis, + r=RADIUS, + qs_surfaces=QS_SURFACES, + ntheta=8, + mpol=3, + ntor=2, + ns_array=(7,), + ftol=1.0e-7, + max_iterations=1200, + adjoint_tol=1.0e-8, + multigrid=False, + ) + equilibrium = problem.solve() + records.append((label, near_axis, problem, equilibrium.quantities)) +except ImportError as error: + print(error) + print("Skipping the optional VMEX publication figure.") + raise SystemExit(0) from None + +finite_label, finite_axis, finite_problem, finite_quantities = records[-1] +print("Differentiating the finite-beta magnetic well...") +well, gradient = jax.value_and_grad( + lambda parameters: ( + qsc.vmex_radial_quantities( + finite_problem, + parameters, + ).magnetic_well + ) +)(finite_problem.parameters) + +plt.style.use("seaborn-v0_8-whitegrid") +figure = plt.figure(figsize=(12.8, 8.0)) +surface_axis = figure.add_subplot(2, 2, 1, projection="3d") +plot_surface_3d( + finite_axis, + radius=0.08, + ntheta=32, + ax=surface_axis, + cmap="magma", +) +surface_axis.view_init(elev=25, azim=35) +surface_axis.set_title(r"VMEX source boundary" "\n" r"finite $p_2$, exactly $I_2=0$") + +iota_axis = figure.add_subplot(2, 2, 2) +qs_axis = figure.add_subplot(2, 2, 3) +well_axis = figure.add_subplot(2, 2, 4) +for label, near_axis, _problem, quantities in records: + iota_axis.plot( + np.asarray(quantities.s), + np.asarray(quantities.iota), + "-o", + label=label, + ) + iota_axis.scatter( + [0], + [float(near_axis.iota)], + marker="x", + s=55, + ) + qs_axis.semilogy( + np.asarray(quantities.qs_surfaces), + np.asarray(quantities.quasisymmetry), + "-o", + label=label, + ) +iota_axis.set_xlabel(r"normalized toroidal flux $s$") +iota_axis.set_ylabel(r"$\iota(s)$ in pyQSC convention") +iota_axis.set_title("on-axis construction → radial equilibrium") +iota_axis.legend() +qs_axis.set_xlabel(r"normalized toroidal flux $s$") +qs_axis.set_ylabel("VMEX QS residual") +qs_axis.set_title("differentiable quasisymmetry profile") +qs_axis.legend() + +wells = [float(item[3].magnetic_well) for item in records] +well_axis.bar([item[0] for item in records], wells, color=("tab:blue", "tab:red")) +well_axis.axhline(0, color="black", linewidth=1) +well_axis.set_ylabel(r"$(V'(0)-V'(1))/V'(0)$") +well_axis.set_title("magnetic well + implicit gradient") +well_axis.text( + 0.04, + 0.06, + rf"$\partial W/\partial p_\mathrm{{scale}}={float(gradient.pres_scale):+.2e}$" + "\n" + rf"$\|\partial W/\partial RBC\|={float(np.linalg.norm(np.asarray(gradient.rbc))):.2e}$", + transform=well_axis.transAxes, + fontsize=9, +) +figure.tight_layout() + +try: + repository = Path(__file__).resolve().parents[2] + commit = subprocess.check_output( + ("git", "-C", str(repository), "rev-parse", "HEAD"), + text=True, + ).strip() +except (OSError, subprocess.CalledProcessError): + commit = "unavailable" +metadata = { + "git_commit": commit, + "vmex_version": finite_problem.vmex_version, + "vmex_validated_commit": finite_problem.validated_commit, + "radius": RADIUS, + "adjoint_tolerance": finite_problem.adjoint_tol, + "cases": { + label: { + "iota": np.asarray(quantities.iota).tolist(), + "quasisymmetry": np.asarray(quantities.quasisymmetry).tolist(), + "magnetic_well": float(quantities.magnetic_well), + "thermal_energy": float(quantities.thermal_energy), + } + for label, _near_axis, _problem, quantities in records + }, + "finite_beta_well": float(well), + "finite_beta_well_gradient_pres_scale": float(gradient.pres_scale), + "finite_beta_well_gradient_rbc_norm": float(np.linalg.norm(np.asarray(gradient.rbc))), +} +if SAVE_OUTPUT: + OUTPUT_STEM.parent.mkdir(parents=True, exist_ok=True) + for suffix in ("png", "svg", "pdf"): + figure.savefig(OUTPUT_STEM.with_suffix(f".{suffix}"), dpi=220, bbox_inches="tight") + README_PNG.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(README_PNG, dpi=120, bbox_inches="tight") + OUTPUT_STEM.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print("saved:", OUTPUT_STEM) +if SHOW_FIGURE: + plt.show() +else: + plt.close(figure) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..634245b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,85 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "pyqsc-jax" +description = "Differentiable near-axis stellarator construction and plasma-coil field jets in JAX" +readme = "README.md" +license = "MIT" +license-files = ["LICENSE", "NOTICE"] +authors = [{ name = "UW Plasma", email = "rogerio.jorge@wisc.edu" }] +requires-python = ">=3.12" +dynamic = ["version"] +dependencies = ["jax", "solvax"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering :: Physics", + "Typing :: Typed", +] + +[project.urls] +Homepage = "https://github.com/uwplasma/pyQSC_JAX" +Documentation = "https://pyqsc-jax.readthedocs.io" +Issues = "https://github.com/uwplasma/pyQSC_JAX/issues" +Source = "https://github.com/uwplasma/pyQSC_JAX" + +[project.optional-dependencies] +plot = ["matplotlib"] +vmex = ["vmex"] +docs = [ + "furo", + "myst-parser", + "sphinx", + "sphinx-copybutton", + "sphinxcontrib-bibtex", +] +dev = [ + "build", + "numpy", + "pytest", + "pytest-cov", + "ruff", + "scipy", + "twine", +] + +[tool.hatch.version] +path = "src/pyqsc_jax/__init__.py" + +[tool.hatch.build.targets.wheel] +packages = ["src/pyqsc_jax"] + +[tool.ruff] +line-length = 100 +src = ["src", "tests"] + +[tool.ruff.lint] +select = ["B", "E", "F", "I", "UP"] + +[tool.pytest.ini_options] +addopts = "-ra" +testpaths = ["tests"] +pythonpath = ["src"] +markers = [ + "integration: cross-package or optional-runtime integration tests", + "literature: regression tests tied to published configurations", + "numerics: numerical convergence and conditioning tests", + "physics: independent physical identity tests", + "slow: tests excluded from the default quick developer loop", +] + +[tool.coverage.run] +branch = true +source = ["pyqsc_jax"] + +[tool.coverage.report] +fail_under = 95 +show_missing = true diff --git a/pyqsc_jax/__init__.py b/pyqsc_jax/__init__.py deleted file mode 100644 index 8b13789..0000000 --- a/pyqsc_jax/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/pyqsc_jax/near_axis.py b/pyqsc_jax/near_axis.py deleted file mode 100644 index c6b8e6a..0000000 --- a/pyqsc_jax/near_axis.py +++ /dev/null @@ -1,505 +0,0 @@ -import jax -import jax.numpy as jnp -from jax import jit, jacfwd, grad, vmap, tree_util, lax -from functools import partial - -class near_axis(): - def __init__(self, rc=jnp.array([1, 0.1]), zs=jnp.array([0, 0.1]), etabar=1.0, - B0=1, sigma0=0, I2=0, nphi=31, spsi=1, sG=1, nfp=2, order='r1', B2c=0, p2=0): - assert nphi % 2 == 1, 'nphi must be odd' - self.rc = jnp.array(rc) - self.zs = jnp.array(zs) - self.etabar = etabar - self.nphi = nphi - self.sigma0 = sigma0 - self.I2 = I2 - self.spsi = spsi - self.sG = sG - self.B0 = B0 - self.nfp = nfp - self.order = order - self.B2c = B2c - self.p2 = p2 - - self._dofs = jnp.concatenate((jnp.ravel(self.rc), jnp.ravel(self.zs), jnp.array([etabar]))) - - self.phi = jnp.linspace(0, 2 * jnp.pi / self.nfp, self.nphi, endpoint=False) - self.nfourier = max(len(self.rc), len(self.zs)) - - parameters = self.calculate(self.rc, self.zs, self.etabar) - (self.R0, self.Z0, self.sigma, self.elongation, self.B_axis, self.grad_B_axis, self.axis_length, self.iota, self.iotaN, self.G0, - self.helicity, self.X1c_untwisted, self.X1s_untwisted, self.Y1s_untwisted, self.Y1c_untwisted, - self.normal_R, self.normal_phi, self.normal_z, self.binormal_R, self.binormal_phi, self.binormal_z, - self.L_grad_B, self.inv_L_grad_B, self.torsion, self.curvature, self.varphi, self.R0p, self.Z0p) = parameters - - @property - def dofs(self): - return self._dofs - - @dofs.setter - def dofs(self, new_dofs): - self._dofs = jnp.array(new_dofs) - self.rc = self._dofs[:self.nfourier] - self.zs = self._dofs[self.nfourier:2*self.nfourier] - self.etabar = self._dofs[-1] - parameters = self.calculate(self.rc, self.zs, self.etabar) - (self.R0, self.Z0, self.sigma, self.elongation, self.B_axis, self.grad_B_axis, self.axis_length, self.iota, self.iotaN, self.G0, - self.helicity, self.X1c_untwisted, self.X1s_untwisted, self.Y1s_untwisted, self.Y1c_untwisted, - self.normal_R, self.normal_z, self.normal_phi, self.binormal_R, self.binormal_z, self.binormal_phi, - self.L_grad_B, self.inv_L_grad_B, self.torsion, self.curvature, self.varphi, self.R0p, self.Z0p) = parameters - - @property - def x(self): - return self._dofs - - @x.setter - def x(self, new_x): - self.dofs = new_x - - def _tree_flatten(self): - children = (self.rc, self.zs, self.etabar, self.B0, self.sigma0, self.I2) - aux_data = {"nphi": self.nphi, "spsi": self.spsi, "sG": self.sG, - "nfp": self.nfp, "order": self.order, "B2c": self.B2c, "p2": self.p2} - return (children, aux_data) - - @classmethod - def _tree_unflatten(cls, aux_data, children): - return cls(*children, **aux_data) - - @partial(jit, static_argnames=['self']) - def B_covariant(self, points): - r, theta, phi = points - Br = 0 - Btheta = r*r*self.I2 - Bphi = self.G0 - return jnp.array([Br, Btheta, Bphi]) - - @partial(jit, static_argnames=['self']) - def B_contravariant(self, points): - r, theta, phi = points - jac = self.jacobian(points) - AbsB = self.AbsB(points) - Bphi = r*AbsB/jac - return jnp.array([0, self.iotaN * Bphi, Bphi]) - - @partial(jit, static_argnames=['self']) - def AbsB(self, points): - r, theta, phi = points - return self.B0*(1 + r*self.etabar*jnp.cos(theta)) - - @partial(jit, static_argnames=['self']) - def jacobian(self, points): - r, theta, phi = points - AbsB = self.AbsB(points) - return r*self.B0*(self.G0+self.iota*self.I2)/(AbsB*AbsB) - - @partial(jit, static_argnames=['self']) - def calculate(self, rc, zs, etabar): - phi = self.phi - nphi = self.nphi - nfp = self.nfp - nfourier = self.nfourier - spsi = self.spsi - sG = self.sG - B0 = self.B0 - sigma0 = self.sigma0 - I2 = self.I2 - d_phi = phi[1] - phi[0] - - n_values = jnp.arange(nfourier) * nfp - - @jit - def compute_terms(jn): - n = n_values[jn] - sinangle = jnp.sin(n * phi) - cosangle = jnp.cos(n * phi) - return jnp.array([rc[jn] * cosangle, zs[jn] * sinangle, - rc[jn] * (-n * sinangle), zs[jn] * (n * cosangle), - rc[jn] * (-n * n * cosangle), zs[jn] * (-n * n * sinangle), - rc[jn] * (n * n * n * sinangle), zs[jn] * (-n * n * n * cosangle)]) - - @jit - def spectral_diff_matrix_jax(): - n=nphi - xmin=0 - xmax=2 * jnp.pi / nfp - h = 2 * jnp.pi / n - kk = jnp.arange(1, n) - n_half = n // 2 - topc = 1 / jnp.sin(jnp.arange(1, n_half + 1) * h / 2) - temp = jnp.concatenate((topc, jnp.flip(topc[:n_half]))) - col1 = jnp.concatenate((jnp.array([0]), 0.5 * ((-1) ** kk) * temp)) - row1 = -col1 - vals = jnp.concatenate((row1[-1:0:-1], col1)) - a, b = jnp.ogrid[0:len(col1), len(row1)-1:-1:-1] - return 2 * jnp.pi / (xmax - xmin) * vals[a + b] - - @jit - def determine_helicity(normal_cylindrical): - x_positive = normal_cylindrical[:, 0] >= 0 - z_positive = normal_cylindrical[:, 2] >= 0 - quadrant = 1 * x_positive * z_positive + 2 * (~x_positive) * z_positive \ - + 3 * (~x_positive) * (~z_positive) + 4 * x_positive * (~z_positive) - quadrant = jnp.append(quadrant, quadrant[0]) - delta_quadrant = quadrant[1:] - quadrant[:-1] - increment = jnp.sum((quadrant[:-1] == 4) & (quadrant[1:] == 1)) - decrement = jnp.sum((quadrant[:-1] == 1) & (quadrant[1:] == 4)) - return (jnp.sum(delta_quadrant) + increment - decrement) * spsi * sG - - summed_values = jnp.sum(jax.vmap(compute_terms)(jnp.arange(nfourier)), axis=0) - - R0, Z0, R0p, Z0p, R0pp, Z0pp, R0ppp, Z0ppp = summed_values - d_l_d_phi = jnp.sqrt(R0 * R0 + R0p * R0p + Z0p * Z0p) - d2_l_d_phi2 = (R0 * R0p + R0p * R0pp + Z0p * Z0pp) / d_l_d_phi - B0_over_abs_G0 = nphi / jnp.sum(d_l_d_phi) - abs_G0_over_B0 = 1 / B0_over_abs_G0 - d_l_d_varphi = abs_G0_over_B0 - G0 = sG * abs_G0_over_B0 * B0 - - d_r_d_phi_cylindrical = jnp.stack([R0p, R0, Z0p]).T - d2_r_d_phi2_cylindrical = jnp.stack([R0pp - R0, 2 * R0p, Z0pp]).T - d3_r_d_phi3_cylindrical = jnp.stack([R0ppp - 3 * R0p, 3 * R0pp - R0, Z0ppp]).T - - d_tangent_d_l_cylindrical = (-d_r_d_phi_cylindrical * d2_l_d_phi2[:, None] / d_l_d_phi[:, None] \ - +d2_r_d_phi2_cylindrical) / (d_l_d_phi[:, None] * d_l_d_phi[:, None]) - curvature = jnp.sqrt(jnp.sum(d_tangent_d_l_cylindrical**2, axis=1)) - axis_length = jnp.sum(d_l_d_phi) * d_phi * nfp - varphi = jnp.concatenate([jnp.zeros(1), jnp.cumsum(d_l_d_phi[:-1] + d_l_d_phi[1:])]) * (0.5 * d_phi * 2 * jnp.pi / axis_length) - - tangent_cylindrical = d_r_d_phi_cylindrical / d_l_d_phi[:, None] - normal_cylindrical = d_tangent_d_l_cylindrical / curvature[:, None] - binormal_cylindrical = jnp.cross(tangent_cylindrical, normal_cylindrical) - - torsion_numerator = jnp.sum(d_r_d_phi_cylindrical * jnp.cross(d2_r_d_phi2_cylindrical, d3_r_d_phi3_cylindrical), axis=1) - torsion_denominator = jnp.sum(jnp.cross(d_r_d_phi_cylindrical, d2_r_d_phi2_cylindrical)**2, axis=1) - torsion = torsion_numerator / torsion_denominator - - d_d_phi = spectral_diff_matrix_jax() - d_varphi_d_phi = B0_over_abs_G0 * d_l_d_phi - d_d_varphi = d_d_phi / d_varphi_d_phi[:, None] - helicity = determine_helicity(normal_cylindrical) - - @jit - def replace_first_element(x, new_value): - return jnp.concatenate([jnp.array([new_value]), x[1:]]) - - @jit - def sigma_equation_residual(x): - iota = x[0] - sigma = replace_first_element(x, sigma0) - etaOcurv2 = etabar**2 / curvature**2 - return jnp.matmul(d_d_varphi, sigma) \ - + (iota + helicity * nfp) * (etaOcurv2**2 + 1 + sigma**2) \ - - 2 * etaOcurv2 * (-spsi * torsion + I2 / B0) * G0 / B0 - - @jit - def sigma_equation_jacobian(x): - iota = x[0] - sigma = replace_first_element(x, sigma0) - etaOcurv2 = etabar**2 / curvature**2 - jac = d_d_varphi + (iota + helicity * nfp) * 2 * jnp.diag(sigma) - return jac.at[:, 0].set(etaOcurv2**2 + 1 + sigma**2) - - @partial(jit, static_argnums=(1,)) - def newton(x0, niter=5): - def body_fun(i, x): - residual = sigma_equation_residual(x) - jacobian = sigma_equation_jacobian(x) - step = jax.scipy.linalg.solve(jacobian, -residual) - return x + step - x = jax.lax.fori_loop(0, niter, body_fun, x0) - return x - - x0 = jnp.full(nphi, sigma0) - x0 = replace_first_element(x0, 0.) - sigma = newton(x0) - iota = sigma[0] - iotaN = iota + helicity * nfp - sigma = replace_first_element(sigma, sigma0) - - X1c = etabar / curvature - Y1s = sG * spsi * curvature / etabar - Y1c = sG * spsi * curvature * sigma / etabar - p = + X1c * X1c + Y1s * Y1s + Y1c * Y1c - q = - X1c * Y1s - elongation = (p + jnp.sqrt(p * p - 4 * q * q)) / (2 * jnp.abs(q)) - - B_axis_cylindrical = sG * B0 * tangent_cylindrical.T - B_x = jnp.cos(phi) * B_axis_cylindrical[0] - jnp.sin(phi) * B_axis_cylindrical[1] - B_y = jnp.sin(phi) * B_axis_cylindrical[0] + jnp.cos(phi) * B_axis_cylindrical[1] - B_z = B_axis_cylindrical[2] - B_axis = jnp.array([B_x, B_y, B_z]) - - d_X1c_d_varphi = -etabar / curvature**2 - d_Y1s_d_varphi = jnp.matmul(d_d_varphi, Y1s) - d_Y1c_d_varphi = jnp.matmul(d_d_varphi, Y1c) - t = tangent_cylindrical.transpose() - n = normal_cylindrical.transpose() - b = binormal_cylindrical.transpose() - d_X1c_d_varphi = jnp.matmul(d_d_varphi, X1c) - d_Y1s_d_varphi = jnp.matmul(d_d_varphi, Y1s) - d_Y1c_d_varphi = jnp.matmul(d_d_varphi, Y1c) - factor = spsi * B0 / d_l_d_varphi - tn = sG * B0 * curvature - nt = tn - bb = factor * (X1c * d_Y1s_d_varphi - iotaN * X1c * Y1c) - nn = factor * (d_X1c_d_varphi * Y1s + iotaN * X1c * Y1c) - bn = factor * (-sG * spsi * d_l_d_varphi * torsion - iotaN * X1c * X1c) - nb = factor * (d_Y1c_d_varphi * Y1s - d_Y1s_d_varphi * Y1c + sG * spsi * d_l_d_varphi * torsion + iotaN * (Y1s * Y1s + Y1c * Y1c)) - tt = 0 - nablaB = jnp.array([[ - nn * n[i] * n[j] \ - + bn * b[i] * n[j] + nb * n[i] * b[j] \ - + bb * b[i] * b[j] \ - + tn * t[i] * n[j] + nt * n[i] * t[j] \ - + tt * t[i] * t[j] - for i in range(3)] for j in range(3)]) - cosphi = jnp.cos(phi) - sinphi = jnp.sin(phi) - grad_B_axis = jnp.array([ - [cosphi**2*nablaB[0, 0] - cosphi*sinphi*(nablaB[0, 1] + nablaB[1, 0]) + - sinphi**2*nablaB[1, 1], cosphi**2*nablaB[0, 1] - sinphi**2*nablaB[1, 0] + - cosphi*sinphi*(nablaB[0, 0] - nablaB[1, 1]), cosphi*nablaB[0, 2] - - sinphi*nablaB[1, 2]], [-(sinphi**2*nablaB[0, 1]) + cosphi**2*nablaB[1, 0] + - cosphi*sinphi*(nablaB[0, 0] - nablaB[1, 1]), sinphi**2*nablaB[0, 0] + - cosphi*sinphi*(nablaB[0, 1] + nablaB[1, 0]) + cosphi**2*nablaB[1, 1], - sinphi*nablaB[0, 2] + cosphi*nablaB[1, 2]], - [cosphi*nablaB[2, 0] - sinphi*nablaB[2, 1], sinphi*nablaB[2, 0] + cosphi*nablaB[2, 1], - nablaB[2, 2]] - ]) - - grad_B_colon_grad_B = tn * tn + nt * nt \ - + bb * bb + nn * nn \ - + nb * nb + bn * bn \ - + tt * tt - L_grad_B = self.B0 * jnp.sqrt(2 / grad_B_colon_grad_B) - inv_L_grad_B = 1.0 / L_grad_B - - X1c_untwisted = jnp.where(helicity == 0, X1c, X1c * jnp.cos(-helicity * nfp * varphi)) - X1s_untwisted = jnp.where(helicity == 0, 0 * X1c, X1c * jnp.sin(-helicity * nfp * varphi)) - Y1s_untwisted = jnp.where(helicity == 0, Y1s, Y1s * jnp.cos(-helicity * nfp * varphi) + Y1c * jnp.sin(-helicity * nfp * varphi)) - Y1c_untwisted = jnp.where(helicity == 0, Y1c, Y1s * (-jnp.sin(-helicity * nfp * varphi)) + Y1c * jnp.cos(-helicity * nfp * varphi)) - - normal_R = normal_cylindrical[:,0] - normal_phi = normal_cylindrical[:,1] - normal_z = normal_cylindrical[:,2] - binormal_R = binormal_cylindrical[:,0] - binormal_phi = binormal_cylindrical[:,1] - binormal_z = binormal_cylindrical[:,2] - - return (R0, Z0, sigma, elongation, B_axis, grad_B_axis, axis_length, iota, iotaN, G0, - helicity, X1c_untwisted, X1s_untwisted, Y1s_untwisted, Y1c_untwisted, - normal_R, normal_phi, normal_z, binormal_R, binormal_phi, binormal_z, - L_grad_B, inv_L_grad_B, torsion, curvature, varphi, R0p, Z0p) - - @jit - def residual_phi0_of_theta_varphi_func(self, phi_0, r, theta, varphi): - X_at_this_theta = r * (self.X1c_untwisted * jnp.cos(theta) + self.X1s_untwisted * jnp.sin(theta)) - Y_at_this_theta = r * (self.Y1c_untwisted * jnp.cos(theta) + self.Y1s_untwisted * jnp.sin(theta)) - _, _, phi = self.Frenet_to_cylindrical_1_point(phi_0, X_at_this_theta, Y_at_this_theta) - nu0 = self.interpolated_array_at_point(self.varphi-self.phi, phi_0) - X1c = self.interpolated_array_at_point(self.X1c_untwisted, phi_0) - X1s = self.interpolated_array_at_point(self.X1s_untwisted, phi_0) - Y1c = self.interpolated_array_at_point(self.Y1c_untwisted, phi_0) - Y1s = self.interpolated_array_at_point(self.Y1s_untwisted, phi_0) - bR = self.interpolated_array_at_point(self.binormal_R, phi_0) - bZ = self.interpolated_array_at_point(self.binormal_z, phi_0) - nR = self.interpolated_array_at_point(self.normal_R, phi_0) - nZ = self.interpolated_array_at_point(self.normal_z, phi_0) - R0 = self.interpolated_array_at_point(self.R0, phi_0) - R0p = self.interpolated_array_at_point(self.R0p, phi_0) - Z0p = self.interpolated_array_at_point(self.Z0p, phi_0) - nu1c = X1c * (bR * Z0p - bZ * R0p)/R0 + Y1c * (nZ * R0p - nR * Z0p)/R0 - nu1s = X1s * (bR * Z0p - bZ * R0p)/R0 + Y1s * (nZ * R0p - nR * Z0p)/R0 - nu = nu0 + r * (nu1c * jnp.cos(theta) + nu1s * jnp.sin(theta)) - return phi + nu - varphi - - @jit - def phi_of_theta_varphi(self, r, theta, varphi): - residual = partial(self.residual_phi0_of_theta_varphi_func, theta=theta, r=r, varphi=varphi) - - def internal_newton(f, x0): - def body_fun(i, x): - res = f(x) - jac = grad(f)(x) - return x - res / jac - return jax.lax.fori_loop(0, 5, body_fun, x0) - - phi_on_axis = lax.custom_root(residual, varphi, internal_newton, lambda g, y: y / g(1.0)) - X_at_this_theta = r * (self.X1c_untwisted * jnp.cos(theta) + self.X1s_untwisted * jnp.sin(theta)) - Y_at_this_theta = r * (self.Y1c_untwisted * jnp.cos(theta) + self.Y1s_untwisted * jnp.sin(theta)) - _, _, phi_off_axis = self.Frenet_to_cylindrical_1_point(phi_on_axis, X_at_this_theta, Y_at_this_theta) - return phi_off_axis - - @jit - def interpolated_array_at_point(self,array,point): - sp=jnp.interp(jnp.array([point]), jnp.append(self.phi,2*jnp.pi/self.nfp), jnp.append(array,array[0]), period=2*jnp.pi/self.nfp)[0] - return sp - - @jit - def Frenet_to_cylindrical_residual_func(self,phi0, phi_target, X_at_this_theta, Y_at_this_theta): - sinphi0 = jnp.sin(phi0) - cosphi0 = jnp.cos(phi0) - R0_at_phi0 = self.interpolated_array_at_point(self.R0,phi0) - X_at_phi0 = self.interpolated_array_at_point(X_at_this_theta,phi0) - Y_at_phi0 = self.interpolated_array_at_point(Y_at_this_theta,phi0) - normal_R = self.interpolated_array_at_point(self.normal_R,phi0) - normal_phi = self.interpolated_array_at_point(self.normal_phi,phi0) - binormal_R = self.interpolated_array_at_point(self.binormal_R,phi0) - binormal_phi = self.interpolated_array_at_point(self.binormal_phi,phi0) - normal_x = normal_R * cosphi0 - normal_phi * sinphi0 - normal_y = normal_R * sinphi0 + normal_phi * cosphi0 - binormal_x = binormal_R * cosphi0 - binormal_phi * sinphi0 - binormal_y = binormal_R * sinphi0 + binormal_phi * cosphi0 - total_x = R0_at_phi0 * cosphi0 + X_at_phi0 * normal_x + Y_at_phi0 * binormal_x - total_y = R0_at_phi0 * sinphi0 + X_at_phi0 * normal_y + Y_at_phi0 * binormal_y - Frenet_to_cylindrical_residual = jnp.arctan2(total_y, total_x) - phi_target - Frenet_to_cylindrical_residual = jnp.where(Frenet_to_cylindrical_residual > jnp.pi, Frenet_to_cylindrical_residual - 2 * jnp.pi, Frenet_to_cylindrical_residual) - Frenet_to_cylindrical_residual = jnp.where(Frenet_to_cylindrical_residual <-jnp.pi, Frenet_to_cylindrical_residual + 2 * jnp.pi, Frenet_to_cylindrical_residual) - return Frenet_to_cylindrical_residual - - @jit - def Frenet_to_cylindrical_1_point(self, phi0, X_at_this_theta, Y_at_this_theta): - sinphi0 = jnp.sin(phi0) - cosphi0 = jnp.cos(phi0) - R0_at_phi0 = self.interpolated_array_at_point(self.R0,phi0) - z0_at_phi0 = self.interpolated_array_at_point(self.Z0,phi0) - X_at_phi0 = self.interpolated_array_at_point(X_at_this_theta,phi0) - Y_at_phi0 = self.interpolated_array_at_point(Y_at_this_theta,phi0) - normal_R = self.interpolated_array_at_point(self.normal_R,phi0) - normal_phi = self.interpolated_array_at_point(self.normal_phi,phi0) - normal_z = self.interpolated_array_at_point(self.normal_z,phi0) - binormal_R = self.interpolated_array_at_point(self.binormal_R,phi0) - binormal_phi = self.interpolated_array_at_point(self.binormal_phi,phi0) - binormal_z = self.interpolated_array_at_point(self.binormal_z,phi0) - normal_x = normal_R * cosphi0 - normal_phi * sinphi0 - normal_y = normal_R * sinphi0 + normal_phi * cosphi0 - binormal_x = binormal_R * cosphi0 - binormal_phi * sinphi0 - binormal_y = binormal_R * sinphi0 + binormal_phi * cosphi0 - total_x = R0_at_phi0 * cosphi0 + X_at_phi0 * normal_x + Y_at_phi0 * binormal_x - total_y = R0_at_phi0 * sinphi0 + X_at_phi0 * normal_y + Y_at_phi0 * binormal_y - total_z = z0_at_phi0 + X_at_phi0 * normal_z + Y_at_phi0 * binormal_z - total_R = jnp.sqrt(total_x * total_x + total_y * total_y) - total_phi=jnp.arctan2(total_y, total_x) - return total_R, total_z, total_phi - - @partial(jit, static_argnames=['ntheta']) - def Frenet_to_cylindrical(self, r, ntheta=20, phi_is_varphi=False): - nphi_conversion = self.nphi - theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False) - phi_conversion = jnp.linspace(0, 2 * jnp.pi / self.nfp, nphi_conversion, endpoint=False) - - def compute_for_theta(theta_j): - costheta = jnp.cos(theta_j) - sintheta = jnp.sin(theta_j) - X_at_this_theta = r * (self.X1c_untwisted * costheta + self.X1s_untwisted * sintheta) - Y_at_this_theta = r * (self.Y1c_untwisted * costheta + self.Y1s_untwisted * sintheta) - - def compute_for_phi(phi_target): - - def internal_newton(f, x0): - def body_fun(i, x): - res = f(x) - jac = grad(f)(x) - return x - res / jac - return jax.lax.fori_loop(0, 5, body_fun, x0) - - def residual(z): - return jax.lax.cond( - phi_is_varphi, - lambda _: self.residual_phi0_of_theta_varphi_func( - z, r=r, theta=theta_j, varphi=phi_target - ), - lambda _: self.Frenet_to_cylindrical_residual_func( - z, phi_target=phi_target, - X_at_this_theta=X_at_this_theta, - Y_at_this_theta=Y_at_this_theta - ), - operand=None - ) - - phi0_solution = lax.custom_root(residual, phi_target, internal_newton, lambda g, y: y / g(1.0)) - - final_R, final_Z, _ = self.Frenet_to_cylindrical_1_point(phi0_solution, X_at_this_theta, Y_at_this_theta) - return final_R, final_Z, phi0_solution - - return vmap(compute_for_phi)(phi_conversion) - - R_2D, Z_2D, phi0_2D = vmap(compute_for_theta)(theta) - return R_2D, Z_2D, phi0_2D - - - @partial(jit, static_argnames=['mpol', 'ntor']) - def to_Fourier(self, R_2D, Z_2D, nfp, mpol, ntor): - ntheta, nphi_conversion = R_2D.shape - theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False) - phi_conversion = jnp.linspace(0, 2 * jnp.pi / nfp, nphi_conversion, endpoint=False) - - phi2d, theta2d = jnp.meshgrid(phi_conversion, theta, indexing='xy') - factor = 2 / (ntheta * nphi_conversion) - - def compute_RBC_ZBS(m, n): - angle = m * theta2d - n * nfp * phi2d - sinangle, cosangle = jnp.sin(angle), jnp.cos(angle) - - factor2 = jax.lax.cond( - (ntheta % 2 == 0) & (m == (ntheta / 2)), - lambda _: factor / 2, lambda _: factor, - operand=None) - - factor2 = jax.lax.cond( - (nphi_conversion % 2 == 0) & (abs(n) == (nphi_conversion / 2)), - lambda _: factor2 / 2, lambda _: factor2, - operand=None) - - return jnp.sum(R_2D * cosangle * factor2), jnp.sum(Z_2D * sinangle * factor2) - - m_vals = jnp.arange(mpol + 1) - n_vals = jnp.concatenate([jnp.array([1]), jnp.arange(-ntor, ntor + 1)]) if mpol == 0 else jnp.arange(-ntor, ntor + 1) - RBC, ZBS = vmap(lambda n: vmap(lambda m: compute_RBC_ZBS(m, n))(m_vals))(n_vals) - - RBC = RBC.at[ntor, 0].set(jnp.sum(R_2D) / (ntheta * nphi_conversion)) - ZBS = ZBS.at[:ntor, 0].set(0) - RBC = RBC.at[:ntor, 0].set(0) - return RBC, ZBS - - @partial(jit, static_argnames=['ntheta_fourier', 'mpol', 'ntor', 'ntheta', 'nphi', 'phi_is_varphi']) - def get_boundary(self, r=0.1, ntheta=30, nphi=120, ntheta_fourier=20, mpol=5, ntor=5, phi_is_varphi=False, phi_offset=0.0): - R_2D, Z_2D, _ = self.Frenet_to_cylindrical(r, ntheta=ntheta_fourier, phi_is_varphi=phi_is_varphi) - RBC, ZBS = self.to_Fourier(R_2D, Z_2D, self.nfp, mpol=mpol, ntor=ntor) - - theta1D = jnp.linspace(0, 2 * jnp.pi, ntheta) - phi1D = jnp.linspace(0, 2 * jnp.pi, nphi) + phi_offset - - phi2D_original, theta2D = jnp.meshgrid(phi1D, theta1D, indexing='ij') - - phi2D = jax.lax.cond( - phi_is_varphi, - lambda _: vmap(lambda theta_row, varphi_row: vmap(lambda theta, varphi: self.phi_of_theta_varphi(r, theta, varphi))(theta_row, varphi_row))(theta2D, phi2D_original), - lambda _: phi2D_original, - operand=None - ) - - def compute_RZ(m, n): - angle = m * theta2D - n * self.nfp * phi2D_original - return RBC[n + ntor, m] * jnp.cos(angle), ZBS[n + ntor, m] * jnp.sin(angle) - - m_vals = jnp.arange(mpol + 1) - n_vals = jnp.arange(-ntor, ntor + 1) - - R_2Dnew, Z_2Dnew = vmap(lambda m: vmap(lambda n: compute_RZ(m, n))(n_vals))(m_vals) - R_2Dnew, Z_2Dnew = R_2Dnew.sum(axis=(0, 1)), Z_2Dnew.sum(axis=(0, 1)) - - x_2D_plot = R_2Dnew.T * jnp.cos(phi2D.T) - y_2D_plot = R_2Dnew.T * jnp.sin(phi2D.T) - z_2D_plot = Z_2Dnew.T - return x_2D_plot, y_2D_plot, z_2D_plot, R_2Dnew.T - - @partial(jit, static_argnames=['self']) - def B_mag(self, r, theta, phi): - return self.B0*(1 + r * self.etabar * jnp.cos(theta - (self.iota - self.iotaN) * phi)) - - -tree_util.register_pytree_node(near_axis, - near_axis._tree_flatten, - near_axis._tree_unflatten) diff --git a/setup.py b/setup.py deleted file mode 100644 index 709934f..0000000 --- a/setup.py +++ /dev/null @@ -1,17 +0,0 @@ -from setuptools import setup, find_packages - -setup( - name="pyqsc_jax", - version="0.1.0", - author="Uw_Plasma", - packages=find_packages(), - install_requires=[ - "jax>=0.4.0", - "jaxlib>=0.4.0", - ], - classifiers=[ - "Programming Language :: Python :: 3", - "License :: OSI Approved :: MIT License", - ], - python_requires=">=3.9", -) diff --git a/src/pyqsc_jax/__init__.py b/src/pyqsc_jax/__init__.py new file mode 100644 index 0000000..a5732ca --- /dev/null +++ b/src/pyqsc_jax/__init__.py @@ -0,0 +1,181 @@ +"""Differentiable near-axis stellarator construction in JAX.""" + +from pyqsc_jax.axis import Axis +from pyqsc_jax.axis_optimization import ( + AxisSearchCandidate, + AxisSearchContinuation, + AxisSearchOptions, + AxisSearchProblem, + AxisSearchResult, + LocalLeastSquaresReport, + continue_axis_search, + search_axis, + stellarator_symmetric_variable_indices, +) +from pyqsc_jax.configurations import ( + REFERENCE_CONFIGURATIONS, + ReferenceConfiguration, + available_configurations, + get_configuration, + solve_configuration, +) +from pyqsc_jax.continuation import ContinuationResult, continue_etabar_branch +from pyqsc_jax.criteria import Criteria, CriteriaReport, CriterionEvaluation +from pyqsc_jax.diagnostics import mercier_diagnostics +from pyqsc_jax.field import total_field_jet +from pyqsc_jax.first_order import Qsc, solve +from pyqsc_jax.models import ( + FieldJet, + InverseSolveDiagnostics, + LinearSolveReport, + MercierDiagnostics, + NearAxisInputs, + NearAxisSolution, + RootSolveReport, + SecondOrderData, + ShearData, + SingularityDiagnostics, + ThirdOrderData, +) +from pyqsc_jax.near_axis import near_axis +from pyqsc_jax.optimize import ( + B2cOptimizationResult, + B20Diagnostics, + B20ResolutionVerification, + b20_diagnostics, + optimal_B2c_value, + optimize_B2c, + verify_B20_resolution, +) +from pyqsc_jax.plasma import ( + PlasmaCurrentSource, + PlasmaFieldData, + PlasmaGradientData, + PlasmaHessianData, + covariant_current_from_enclosed, + elliptical_channel_gradient, + enclosed_current_from_covariant, + evaluate_weighted_current, + matched_plasma_field_kernel, + pack_symmetric_trace_free_rank2, + pack_symmetric_trace_free_rank3, + plasma_current_source, + plasma_field_on_axis, + plasma_gradient_on_axis, + plasma_hessian_on_axis, + project_symmetric_trace_free_rank2, + project_symmetric_trace_free_rank3, + regularized_axis_integral, + unpack_symmetric_trace_free_rank2, + unpack_symmetric_trace_free_rank3, +) +from pyqsc_jax.second_order import SecondOrderResiduals, second_order_residuals +from pyqsc_jax.shear import solve_magnetic_shear +from pyqsc_jax.singularity import singularity_diagnostics +from pyqsc_jax.solvers import RootSolveOptions +from pyqsc_jax.third_order import solve_third_order +from pyqsc_jax.vmec import ( + VmecBoundary, + VmecExport, + VmecInputParameters, + to_vmec, + uniform_cylindrical_surface, + vmec_boundary, +) +from pyqsc_jax.vmex import ( + VMEX_VALIDATED_COMMIT, + VmexEquilibrium, + VmexProblem, + VmexRadialQuantities, + solve_vmex, + to_vmex_problem, + vmex_parameters_from_solution, + vmex_radial_quantities, +) + +__all__ = [ + "Axis", + "AxisSearchCandidate", + "AxisSearchContinuation", + "AxisSearchOptions", + "AxisSearchProblem", + "AxisSearchResult", + "B2cOptimizationResult", + "B20Diagnostics", + "B20ResolutionVerification", + "ContinuationResult", + "Criteria", + "CriteriaReport", + "CriterionEvaluation", + "FieldJet", + "InverseSolveDiagnostics", + "LinearSolveReport", + "LocalLeastSquaresReport", + "MercierDiagnostics", + "NearAxisInputs", + "NearAxisSolution", + "PlasmaCurrentSource", + "PlasmaFieldData", + "PlasmaGradientData", + "PlasmaHessianData", + "Qsc", + "REFERENCE_CONFIGURATIONS", + "ReferenceConfiguration", + "RootSolveOptions", + "RootSolveReport", + "SecondOrderData", + "SecondOrderResiduals", + "ShearData", + "SingularityDiagnostics", + "ThirdOrderData", + "VmecBoundary", + "VmecExport", + "VmecInputParameters", + "VmexEquilibrium", + "VmexProblem", + "VmexRadialQuantities", + "VMEX_VALIDATED_COMMIT", + "near_axis", + "b20_diagnostics", + "available_configurations", + "continue_etabar_branch", + "continue_axis_search", + "covariant_current_from_enclosed", + "enclosed_current_from_covariant", + "elliptical_channel_gradient", + "evaluate_weighted_current", + "get_configuration", + "matched_plasma_field_kernel", + "pack_symmetric_trace_free_rank2", + "pack_symmetric_trace_free_rank3", + "mercier_diagnostics", + "second_order_residuals", + "search_axis", + "solve", + "solve_configuration", + "solve_magnetic_shear", + "singularity_diagnostics", + "stellarator_symmetric_variable_indices", + "solve_third_order", + "total_field_jet", + "to_vmec", + "to_vmex_problem", + "uniform_cylindrical_surface", + "vmec_boundary", + "vmex_parameters_from_solution", + "vmex_radial_quantities", + "solve_vmex", + "optimal_B2c_value", + "optimize_B2c", + "plasma_current_source", + "plasma_field_on_axis", + "plasma_gradient_on_axis", + "plasma_hessian_on_axis", + "project_symmetric_trace_free_rank2", + "project_symmetric_trace_free_rank3", + "regularized_axis_integral", + "unpack_symmetric_trace_free_rank2", + "unpack_symmetric_trace_free_rank3", + "verify_B20_resolution", +] +__version__ = "0.2.0.dev0" diff --git a/src/pyqsc_jax/axis.py b/src/pyqsc_jax/axis.py new file mode 100644 index 0000000..c092787 --- /dev/null +++ b/src/pyqsc_jax/axis.py @@ -0,0 +1,153 @@ +"""Fourier representations of the magnetic axis.""" + +from dataclasses import dataclass, field +from typing import Any, ClassVar + +import jax +import jax.numpy as jnp + +ArrayLike = Any + + +def _as_coefficient_array(values: ArrayLike, dtype: jnp.dtype) -> jax.Array: + array = jnp.asarray(values, dtype=dtype) + if array.ndim != 1: + raise ValueError("Axis Fourier coefficients must be one-dimensional.") + return array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class Axis: + """Immutable cylindrical Fourier representation of a closed magnetic axis. + + The convention over one field period is + + .. math:: + + R(\\phi) = \\sum_n R_{cn}\\cos(n n_{fp}\\phi) + + R_{sn}\\sin(n n_{fp}\\phi), + + Z(\\phi) = \\sum_n Z_{cn}\\cos(n n_{fp}\\phi) + + Z_{sn}\\sin(n n_{fp}\\phi). + + All four coefficient arrays are padded to the same length. ``nfp`` is + static pytree metadata, while the coefficient arrays are differentiable + pytree leaves. + """ + + rc: jax.Array + zs: jax.Array + nfp: int = field(default=1, metadata={"static": True}) + rs: jax.Array = () + zc: jax.Array = () + + coefficient_order: ClassVar[tuple[str, ...]] = ("rc", "rs", "zc", "zs") + + def __post_init__(self) -> None: + if not isinstance(self.nfp, int) or isinstance(self.nfp, bool) or self.nfp < 1: + raise ValueError("nfp must be a positive integer.") + + raw = tuple(jnp.asarray(values) for values in (self.rc, self.rs, self.zc, self.zs)) + dtype = jnp.result_type(jnp.asarray(0.0), *(array.dtype for array in raw)) + arrays = tuple(_as_coefficient_array(values, dtype) for values in raw) + nfourier = max(array.size for array in arrays) + if nfourier < 1: + raise ValueError("At least one axis Fourier coefficient is required.") + + padded = tuple(jnp.pad(array, (0, nfourier - array.size)) for array in arrays) + object.__setattr__(self, "rc", padded[0]) + object.__setattr__(self, "rs", padded[1]) + object.__setattr__(self, "zc", padded[2]) + object.__setattr__(self, "zs", padded[3]) + + @classmethod + def stellarator_symmetric(cls, *, rc: ArrayLike, zs: ArrayLike, nfp: int = 1) -> "Axis": + """Construct an axis for which ``R`` is even and ``Z`` is odd.""" + + return cls(rc=rc, zs=zs, nfp=nfp) + + @classmethod + def from_dofs(cls, dofs: ArrayLike, *, nfp: int) -> "Axis": + """Construct from four equal blocks ordered as ``rc, rs, zc, zs``.""" + + array = jnp.asarray(dofs) + if array.ndim != 1 or array.size % 4: + raise ValueError( + "Axis dofs must be a one-dimensional array with length divisible by 4." + ) + nfourier = array.size // 4 + rc, rs, zc, zs = jnp.split(array, (nfourier, 2 * nfourier, 3 * nfourier)) + return cls(rc=rc, rs=rs, zc=zc, zs=zs, nfp=nfp) + + @property + def nfourier(self) -> int: + """Number of retained Fourier modes, including mode zero.""" + + return self.rc.size + + @property + def dofs(self) -> jax.Array: + """Coefficient vector ordered as ``rc, rs, zc, zs``.""" + + return jnp.concatenate((self.rc, self.rs, self.zc, self.zs)) + + @property + def stellarator_symmetry_residual(self) -> jax.Array: + """Largest coefficient forbidden by stellarator symmetry.""" + + return jnp.maximum(jnp.max(jnp.abs(self.rs)), jnp.max(jnp.abs(self.zc))) + + def with_dofs(self, dofs: ArrayLike) -> "Axis": + """Return a new axis with packed coefficients and the same ``nfp``.""" + + return type(self).from_dofs(dofs, nfp=self.nfp) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class AxisSamples: + """Axis coordinates and their first three cylindrical-angle derivatives.""" + + phi: jax.Array + R: jax.Array + Z: jax.Array + d_R_d_phi: jax.Array + d_Z_d_phi: jax.Array + d2_R_d_phi2: jax.Array + d2_Z_d_phi2: jax.Array + d3_R_d_phi3: jax.Array + d3_Z_d_phi3: jax.Array + + +def evaluate_axis(axis: Axis, phi: ArrayLike) -> AxisSamples: + """Evaluate an axis and analytic derivatives at arbitrary toroidal angles.""" + + phi = jnp.asarray(phi) + mode = jnp.arange(axis.nfourier, dtype=phi.dtype) * axis.nfp + angle = phi[..., None] * mode + cosine = jnp.cos(angle) + sine = jnp.sin(angle) + + R = jnp.sum(axis.rc * cosine + axis.rs * sine, axis=-1) + Z = jnp.sum(axis.zc * cosine + axis.zs * sine, axis=-1) + d_R = jnp.sum(mode * (-axis.rc * sine + axis.rs * cosine), axis=-1) + d_Z = jnp.sum(mode * (-axis.zc * sine + axis.zs * cosine), axis=-1) + mode2 = mode * mode + d2_R = jnp.sum(-mode2 * (axis.rc * cosine + axis.rs * sine), axis=-1) + d2_Z = jnp.sum(-mode2 * (axis.zc * cosine + axis.zs * sine), axis=-1) + mode3 = mode2 * mode + d3_R = jnp.sum(mode3 * (axis.rc * sine - axis.rs * cosine), axis=-1) + d3_Z = jnp.sum(mode3 * (axis.zc * sine - axis.zs * cosine), axis=-1) + + return AxisSamples( + phi=phi, + R=R, + Z=Z, + d_R_d_phi=d_R, + d_Z_d_phi=d_Z, + d2_R_d_phi2=d2_R, + d2_Z_d_phi2=d2_Z, + d3_R_d_phi3=d3_R, + d3_Z_d_phi3=d3_Z, + ) diff --git a/src/pyqsc_jax/axis_optimization.py b/src/pyqsc_jax/axis_optimization.py new file mode 100644 index 0000000..06a9b45 --- /dev/null +++ b/src/pyqsc_jax/axis_optimization.py @@ -0,0 +1,745 @@ +"""Branch-aware global-to-local optimization of magnetic-axis coefficients.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from math import sqrt +from typing import Literal + +import jax +import jax.numpy as jnp + +from pyqsc_jax.axis import Axis +from pyqsc_jax.criteria import Criteria, CriteriaReport +from pyqsc_jax.first_order import solve +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.optimize import ( + B20Diagnostics, + B20ResolutionVerification, + optimize_B2c, + verify_B20_resolution, +) + +SearchStatus = Literal[ + "verified_zero", + "best_found", + "no_feasible_candidate", + "branch_fold", + "ill_conditioned", + "solver_failure", + "verification_failure", +] + +Selector = Literal[ + "maximum_singular_radius", + "maximum_minimum_L_grad_B", + "maximum_minimum_L_grad_grad_B", + "minimum_maximum_elongation", + "minimum_axis_sobolev_norm", + "target_axis_length", + "maximum_criteria_margin", +] + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class LocalLeastSquaresReport: + """Convergence and damping evidence for one local refinement.""" + + initial_residual_norm: jax.Array + residual_norm: jax.Array + gradient_infinity_norm: jax.Array + step_norm: jax.Array + damping: jax.Array + jacobian_condition_number: jax.Array + iterations: jax.Array + accepted_steps: jax.Array + rejected_steps: jax.Array + function_evaluations: jax.Array + converged: jax.Array + finite: jax.Array + + +@dataclass(frozen=True) +class AxisSearchProblem: + """Physical inputs, bounded axis variables, and branch policy.""" + + axis: Axis + variable_indices: tuple[int, ...] + lower_bounds: jax.Array + upper_bounds: jax.Array + etabar: float + B0: float = 1.0 + sigma0: float = 0.0 + I2: float = 0.0 + p2: float = 0.0 + B2s: float = 0.0 + nphi: int = 31 + sG: int = 1 + spsi: int = 1 + solve_for: str = "iota" + target_iota: float | None = None + criteria: Criteria | None = None + selector: Selector = "maximum_singular_radius" + target_axis_length: float | None = None + + +@dataclass(frozen=True) +class AxisSearchOptions: + """Deterministic exploration, local solve, clustering, and verification policy.""" + + coarse_samples: int = 32 + local_starts: int = 4 + maximum_iterations: int = 20 + initial_damping: float = 1.0e-3 + residual_tolerance: float = 1.0e-9 + gradient_tolerance: float = 1.0e-8 + step_tolerance: float = 1.0e-10 + primary_tolerance: float = 1.0e-6 + verified_zero_tolerance: float = 1.0e-8 + verification_relative_tolerance: float = 5.0e-3 + verification_tail_tolerance: float = 1.0e-6 + verification_multipliers: tuple[int, ...] = (1, 2, 4) + basin_distance_tolerance: float = 1.0e-3 + maximum_linear_condition_number: float = 1.0e14 + + +@dataclass(frozen=True) +class AxisSearchCandidate: + """One distinct locally refined basin and its independent checks.""" + + variables: jax.Array + solution: NearAxisSolution + diagnostics: B20Diagnostics + local_report: LocalLeastSquaresReport + criteria_report: CriteriaReport | None + verification: B20ResolutionVerification + primary_residual: float + selector_value: float + feasible: bool + + +@dataclass(frozen=True) +class AxisSearchResult: + """Reproducible search result with explicit global-certificate semantics.""" + + status: SearchStatus + best: AxisSearchCandidate | None + basins: tuple[AxisSearchCandidate, ...] + coarse_variables: jax.Array + coarse_residuals: jax.Array + coarse_feasible: jax.Array + search_budget: int + local_starts_attempted: int + distinct_basins: int + global_certificate: bool + message: str + + +@dataclass(frozen=True) +class AxisSearchContinuation: + """Warm-started low-to-high Fourier search stages.""" + + stages: tuple[AxisSearchResult, ...] + final: AxisSearchResult | None + complete: bool + + +def stellarator_symmetric_variable_indices( + axis: Axis, + *, + modes: tuple[int, ...] | None = None, +) -> tuple[int, ...]: + """Return packed indices for nonconstant ``rc`` and ``zs`` coefficients.""" + + if modes is None: + modes = tuple(range(1, axis.nfourier)) + if any( + not isinstance(mode, int) or isinstance(mode, bool) or mode < 1 or mode >= axis.nfourier + for mode in modes + ): + raise ValueError("modes must contain valid positive retained Fourier modes.") + zs_offset = 3 * axis.nfourier + return tuple(modes) + tuple(zs_offset + mode for mode in modes) + + +def _first_primes(count: int) -> tuple[int, ...]: + primes: list[int] = [] + candidate = 2 + while len(primes) < count: + if all(candidate % prime for prime in primes if prime <= sqrt(candidate)): + primes.append(candidate) + candidate += 1 + return tuple(primes) + + +def _radical_inverse(index: int, base: int) -> float: + result = 0.0 + fraction = 1.0 / base + while index: + result += fraction * (index % base) + index //= base + fraction /= base + return result + + +def _halton_box(count: int, dimension: int, dtype: jnp.dtype) -> jax.Array: + primes = _first_primes(dimension) + values = [[_radical_inverse(sample + 1, base) for base in primes] for sample in range(count)] + return jnp.asarray(values, dtype=dtype) + + +def _validate_problem(problem: AxisSearchProblem, options: AxisSearchOptions) -> None: + indices = problem.variable_indices + if not indices or len(set(indices)) != len(indices): + raise ValueError("variable_indices must be nonempty and unique.") + if any( + not isinstance(index, int) + or isinstance(index, bool) + or index < 0 + or index >= problem.axis.dofs.size + for index in indices + ): + raise ValueError("variable_indices contains an invalid packed-axis index.") + lower = jnp.asarray(problem.lower_bounds) + upper = jnp.asarray(problem.upper_bounds) + if lower.shape != (len(indices),) or upper.shape != (len(indices),): + raise ValueError("lower_bounds and upper_bounds must match variable_indices.") + if not bool(jnp.all(jnp.isfinite(lower) & jnp.isfinite(upper) & (lower < upper))): + raise ValueError("Every variable bound must be finite and strictly ordered.") + initial = problem.axis.dofs[jnp.asarray(indices)] + if not bool(jnp.all((initial > lower) & (initial < upper))): + raise ValueError("The initial axis variables must lie strictly inside their bounds.") + if problem.nphi < 3 or problem.nphi % 2 != 1: + raise ValueError("nphi must be an odd integer >= 3.") + if problem.solve_for == "iota": + if problem.target_iota is not None: + raise ValueError("target_iota requires solve_for='etabar' or solve_for='I2'.") + elif problem.solve_for in ("etabar", "I2"): + if problem.target_iota is None: + raise ValueError("target_iota is required for an inverse solve.") + else: + raise ValueError("solve_for must be 'iota', 'etabar', or 'I2'.") + if problem.selector == "target_axis_length" and problem.target_axis_length is None: + raise ValueError("target_axis_length is required by the selected policy.") + if problem.selector == "maximum_criteria_margin" and problem.criteria is None: + raise ValueError("maximum_criteria_margin requires a criteria profile.") + if problem.selector not in ( + "maximum_singular_radius", + "maximum_minimum_L_grad_B", + "maximum_minimum_L_grad_grad_B", + "minimum_maximum_elongation", + "minimum_axis_sobolev_norm", + "target_axis_length", + "maximum_criteria_margin", + ): + raise ValueError("Unknown canonical selector.") + + positive_integers = ( + options.coarse_samples, + options.local_starts, + options.maximum_iterations, + ) + if any( + not isinstance(value, int) or isinstance(value, bool) or value < 1 + for value in positive_integers + ): + raise ValueError("Search sample, start, and iteration counts must be positive integers.") + positive_values = ( + options.initial_damping, + options.residual_tolerance, + options.gradient_tolerance, + options.step_tolerance, + options.primary_tolerance, + options.verified_zero_tolerance, + options.verification_relative_tolerance, + options.verification_tail_tolerance, + options.basin_distance_tolerance, + options.maximum_linear_condition_number, + ) + if any(value <= 0 for value in positive_values): + raise ValueError("Search tolerances, damping, and condition limit must be positive.") + if not options.verification_multipliers: + raise ValueError("verification_multipliers must be nonempty.") + + +def _physical_from_internal( + internal: jax.Array, + lower: jax.Array, + upper: jax.Array, +) -> jax.Array: + midpoint = 0.5 * (lower + upper) + half_width = 0.5 * (upper - lower) + return midpoint + half_width * jnp.tanh(internal) + + +def _internal_from_physical( + physical: jax.Array, + lower: jax.Array, + upper: jax.Array, +) -> jax.Array: + midpoint = 0.5 * (lower + upper) + half_width = 0.5 * (upper - lower) + return jnp.arctanh((physical - midpoint) / half_width) + + +def _solve_axis_candidate( + problem: AxisSearchProblem, + variables: jax.Array, +) -> tuple[NearAxisSolution, B20Diagnostics]: + indices = jnp.asarray(problem.variable_indices) + axis = problem.axis.with_dofs(problem.axis.dofs.at[indices].set(variables)) + keyword_arguments = { + "axis": axis, + "etabar": problem.etabar, + "B0": problem.B0, + "sigma0": problem.sigma0, + "I2": problem.I2, + "p2": problem.p2, + "B2c": 0.0, + "B2s": problem.B2s, + "nphi": problem.nphi, + "order": "r2", + "sG": problem.sG, + "spsi": problem.spsi, + "solve_for": problem.solve_for, + } + if problem.solve_for != "iota": + keyword_arguments["iota"] = problem.target_iota + optimized = optimize_B2c(solve(**keyword_arguments)) + return optimized.solution, optimized.diagnostics + + +def _projected_residual( + problem: AxisSearchProblem, + variables: jax.Array, +) -> jax.Array: + solution, diagnostics = _solve_axis_candidate(problem, variables) + weights = solution.geometry.d_l_d_phi + return jnp.sqrt(weights / jnp.sum(weights)) * diagnostics.anomaly / solution.inputs.B0 + + +def _coarse_metrics( + problem: AxisSearchProblem, + variables: jax.Array, +) -> tuple[jax.Array, jax.Array]: + solution, diagnostics = _solve_axis_candidate(problem, variables) + finite = ( + jnp.isfinite(diagnostics.weighted_l2) + & solution.root_report.converged + & solution.linear_report.converged + & solution.geometry.diagnostics.frenet_valid + & solution.geometry.diagnostics.cylindrical_coordinates_valid + ) + return jnp.where(finite, diagnostics.weighted_l2, jnp.inf), finite + + +def _levenberg_marquardt( + residual_function, + initial: jax.Array, + options: AxisSearchOptions, +) -> tuple[jax.Array, LocalLeastSquaresReport]: + residual_and_jacobian = jax.jit( + lambda value: (residual_function(value), jax.jacrev(residual_function)(value)) + ) + residual = jax.jit(residual_function) + internal = initial + current_residual, jacobian = residual_and_jacobian(internal) + initial_norm = jnp.linalg.norm(current_residual) + damping = options.initial_damping + accepted = 0 + rejected = 0 + function_evaluations = 1 + converged = False + step_norm = jnp.asarray(jnp.inf, dtype=initial.dtype) + gradient_norm = jnp.asarray(jnp.inf, dtype=initial.dtype) + + for _iteration in range(1, options.maximum_iterations + 1): + normal_matrix = jacobian.T @ jacobian + gradient = jacobian.T @ current_residual + gradient_norm = jnp.linalg.norm(gradient, ord=jnp.inf) + diagonal_scale = jnp.maximum(jnp.diag(normal_matrix), 1.0) + step = -jnp.linalg.solve( + normal_matrix + damping * jnp.diag(diagonal_scale), + gradient, + ) + step_norm = jnp.linalg.norm(step) + candidate = internal + step + candidate_residual = residual(candidate) + function_evaluations += 1 + current_cost = 0.5 * jnp.sum(current_residual**2) + candidate_cost = 0.5 * jnp.sum(candidate_residual**2) + predicted_reduction = -gradient @ step - 0.5 * step @ normal_matrix @ step + ratio = (current_cost - candidate_cost) / jnp.maximum( + predicted_reduction, + jnp.finfo(current_cost.dtype).tiny, + ) + accept = bool( + jnp.isfinite(candidate_cost) + & jnp.isfinite(step_norm) + & (candidate_cost < current_cost) + & (ratio > 0) + ) + if accept: + internal = candidate + current_residual, jacobian = residual_and_jacobian(internal) + function_evaluations += 1 + accepted += 1 + damping = jnp.maximum( + damping * jnp.maximum(1 / 3, 1 - (2 * ratio - 1) ** 3), + jnp.finfo(current_cost.dtype).eps, + ) + else: + rejected += 1 + damping = damping * 2 + + residual_norm = jnp.linalg.norm(current_residual) + converged = bool( + (residual_norm <= options.residual_tolerance) + | (gradient_norm <= options.gradient_tolerance) + | (step_norm <= options.step_tolerance * (1 + jnp.linalg.norm(internal))) + ) + if converged: + break + + normal_eigenvalues = jnp.maximum( + jnp.linalg.eigvalsh(jacobian.T @ jacobian), + 0, + ) + largest_eigenvalue = normal_eigenvalues[-1] + smallest_eigenvalue = normal_eigenvalues[0] + rank_deficient = smallest_eigenvalue <= (jnp.finfo(jacobian.dtype).eps * largest_eigenvalue) + condition_number = jnp.where( + rank_deficient, + jnp.inf, + jnp.sqrt(largest_eigenvalue / smallest_eigenvalue), + ) + condition_number = jnp.where(jnp.isnan(condition_number), jnp.inf, condition_number) + residual_converged = jnp.linalg.norm(current_residual) <= options.residual_tolerance + derivative_finite = jnp.all(jnp.isfinite(jacobian)) + finite = ( + jnp.all(jnp.isfinite(internal)) + & jnp.all(jnp.isfinite(current_residual)) + & (derivative_finite | residual_converged) + ) + return internal, LocalLeastSquaresReport( + initial_residual_norm=initial_norm, + residual_norm=jnp.linalg.norm(current_residual), + gradient_infinity_norm=gradient_norm, + step_norm=step_norm, + damping=jnp.asarray(damping), + jacobian_condition_number=condition_number, + iterations=jnp.asarray(_iteration), + accepted_steps=jnp.asarray(accepted), + rejected_steps=jnp.asarray(rejected), + function_evaluations=jnp.asarray(function_evaluations), + converged=jnp.asarray(converged), + finite=finite, + ) + + +def _axis_sobolev_norm(axis: Axis) -> jax.Array: + modes = jnp.arange(axis.nfourier, dtype=axis.dofs.dtype) + weights = (1 + modes**2) ** 2 + return jnp.sqrt(jnp.sum(weights * (axis.rc**2 + axis.rs**2 + axis.zc**2 + axis.zs**2))) + + +def _selector_value( + problem: AxisSearchProblem, + solution: NearAxisSolution, + criteria_report: CriteriaReport | None, +) -> float: + if problem.selector == "maximum_singular_radius": + value = solution.r_singularity + elif problem.selector == "maximum_minimum_L_grad_B": + value = jnp.min(solution.L_grad_B) + elif problem.selector == "maximum_minimum_L_grad_grad_B": + value = jnp.min(solution.L_grad_grad_B) + elif problem.selector == "minimum_maximum_elongation": + value = -jnp.max(solution.elongation) + elif problem.selector == "minimum_axis_sobolev_norm": + value = -_axis_sobolev_norm(solution.inputs.axis) + elif problem.selector == "target_axis_length": + value = -jnp.abs(solution.axis_length - problem.target_axis_length) + else: + normalized_margins = [ + evaluation.margin / jnp.maximum(jnp.abs(evaluation.threshold), 1.0) + for evaluation in criteria_report.evaluations + ] + value = jnp.min(jnp.stack(normalized_margins)) + return float(value) + + +def _candidate_is_distinct( + variables: jax.Array, + candidates: list[AxisSearchCandidate], + lower: jax.Array, + upper: jax.Array, + tolerance: float, +) -> bool: + scaled = (variables - lower) / (upper - lower) + return all( + float(jnp.linalg.norm(scaled - (candidate.variables - lower) / (upper - lower))) > tolerance + for candidate in candidates + ) + + +def _copy_retained_axis_modes(source: Axis, target: Axis) -> Axis: + if source.nfp != target.nfp: + raise ValueError("Fourier continuation stages must use the same nfp.") + retained = min(source.nfourier, target.nfourier) + return Axis( + rc=target.rc.at[:retained].set(source.rc[:retained]), + rs=target.rs.at[:retained].set(source.rs[:retained]), + zc=target.zc.at[:retained].set(source.zc[:retained]), + zs=target.zs.at[:retained].set(source.zs[:retained]), + nfp=target.nfp, + ) + + +def _verification_passes( + verification: B20ResolutionVerification, + options: AxisSearchOptions, +) -> bool: + weighted_zero = verification.weighted_l2 <= options.verified_zero_tolerance + maximum_zero = verification.grid_maximum <= options.verified_zero_tolerance + relative_l2_stable = ( + verification.relative_weighted_l2_change <= options.verification_relative_tolerance + ) | weighted_zero + relative_maximum_stable = ( + verification.relative_grid_maximum_change <= options.verification_relative_tolerance + ) | maximum_zero + spectral_tail_resolved = ( + verification.fourier_tail_ratio[-1] <= options.verification_tail_tolerance + ) | maximum_zero[-1] + return bool( + jnp.all(weighted_zero) + & jnp.all(maximum_zero) + & jnp.all(relative_l2_stable) + & jnp.all(relative_maximum_stable) + & spectral_tail_resolved + ) + + +def search_axis( + problem: AxisSearchProblem, + *, + options: AxisSearchOptions | None = None, + seeds: tuple[jax.Array, ...] = (), +) -> AxisSearchResult: + """Explore, refine, cluster, and independently verify bounded axis candidates.""" + + if options is None: + options = AxisSearchOptions() + _validate_problem(problem, options) + lower = jnp.asarray(problem.lower_bounds, dtype=problem.axis.dofs.dtype) + upper = jnp.asarray(problem.upper_bounds, dtype=problem.axis.dofs.dtype) + initial = problem.axis.dofs[jnp.asarray(problem.variable_indices)] + for seed in seeds: + seed_array = jnp.asarray(seed) + if seed_array.shape != initial.shape or not bool( + jnp.all(jnp.isfinite(seed_array) & (seed_array > lower) & (seed_array < upper)) + ): + raise ValueError("Every explicit seed must be finite, in bounds, and correctly sized.") + + samples = lower + (upper - lower) * _halton_box( + options.coarse_samples, + len(problem.variable_indices), + initial.dtype, + ) + explicit = jnp.stack((initial, *(jnp.asarray(seed) for seed in seeds))) + coarse_variables = jnp.concatenate((explicit, samples), axis=0) + coarse_residuals, coarse_feasible = jax.jit( + jax.vmap(lambda variables: _coarse_metrics(problem, variables)) + )(coarse_variables) + + finite_indices = [ + int(index) for index in jnp.argsort(coarse_residuals) if bool(coarse_feasible[index]) + ] + selected_indices = finite_indices[: options.local_starts] + if not selected_indices: + return AxisSearchResult( + status="solver_failure", + best=None, + basins=(), + coarse_variables=coarse_variables, + coarse_residuals=coarse_residuals, + coarse_feasible=coarse_feasible, + search_budget=int(coarse_variables.shape[0]), + local_starts_attempted=0, + distinct_basins=0, + global_certificate=False, + message="No coarse candidate produced a finite converged near-axis solve.", + ) + + residual_function = lambda internal: _projected_residual( # noqa: E731 + problem, + _physical_from_internal(internal, lower, upper), + ) + candidates: list[AxisSearchCandidate] = [] + local_evaluations = 0 + ill_conditioned_count = 0 + for index in selected_indices: + start = coarse_variables[index] + internal_start = _internal_from_physical(start, lower, upper) + internal, local_report = _levenberg_marquardt( + residual_function, + internal_start, + options, + ) + local_evaluations += int(local_report.function_evaluations) + variables = _physical_from_internal(internal, lower, upper) + if not bool(local_report.finite): + continue + solution, diagnostics = _solve_axis_candidate(problem, variables) + if solution.linear_report.matrix_condition_number > options.maximum_linear_condition_number: + ill_conditioned_count += 1 + continue + criteria_report = None if problem.criteria is None else problem.criteria.evaluate(solution) + feasible = bool( + solution.root_report.converged + & solution.linear_report.converged + & solution.geometry.diagnostics.frenet_valid + & solution.geometry.diagnostics.cylindrical_coordinates_valid + ) and (criteria_report is None or criteria_report.passed) + if not _candidate_is_distinct( + variables, + candidates, + lower, + upper, + options.basin_distance_tolerance, + ): + continue + verification = verify_B20_resolution( + solution, + multipliers=options.verification_multipliers, + ) + candidates.append( + AxisSearchCandidate( + variables=variables, + solution=solution, + diagnostics=diagnostics, + local_report=local_report, + criteria_report=criteria_report, + verification=verification, + primary_residual=float(diagnostics.weighted_l2), + selector_value=_selector_value(problem, solution, criteria_report), + feasible=feasible, + ) + ) + + search_budget = int(coarse_variables.shape[0]) + local_evaluations + if not candidates: + status: SearchStatus = ( + "ill_conditioned" + if ill_conditioned_count == len(selected_indices) + else "solver_failure" + ) + return AxisSearchResult( + status=status, + best=None, + basins=(), + coarse_variables=coarse_variables, + coarse_residuals=coarse_residuals, + coarse_feasible=coarse_feasible, + search_budget=search_budget, + local_starts_attempted=len(selected_indices), + distinct_basins=0, + global_certificate=False, + message="Local refinement did not produce a finite, acceptable candidate.", + ) + + feasible = [candidate for candidate in candidates if candidate.feasible] + if not feasible: + best = min(candidates, key=lambda candidate: candidate.primary_residual) + return AxisSearchResult( + status="no_feasible_candidate", + best=best, + basins=tuple(candidates), + coarse_variables=coarse_variables, + coarse_residuals=coarse_residuals, + coarse_feasible=coarse_feasible, + search_budget=search_budget, + local_starts_attempted=len(selected_indices), + distinct_basins=len(candidates), + global_certificate=False, + message="The search found converged basins, but none passed the hard criteria.", + ) + + primary_feasible = [ + candidate + for candidate in feasible + if candidate.primary_residual <= options.primary_tolerance + ] + if primary_feasible: + best = max(primary_feasible, key=lambda candidate: candidate.selector_value) + else: + best = min(feasible, key=lambda candidate: candidate.primary_residual) + + if best.solution.inverse is not None and bool(best.solution.branch_fold): + status = "branch_fold" + message = "The selected inverse-solve candidate lies at a detected branch fold." + elif best.primary_residual <= options.verified_zero_tolerance: + if _verification_passes(best.verification, options): + status = "verified_zero" + message = ( + "The independently verified nonnegative B20 residual reaches its " + "global lower bound of zero within the requested tolerances." + ) + else: + status = "verification_failure" + message = "The nominal zero failed independent resolution or spectral checks." + else: + status = "best_found" + message = ( + "This is the best basin found within the reported search budget; " + "there is no global-minimum certificate." + ) + return AxisSearchResult( + status=status, + best=best, + basins=tuple(candidates), + coarse_variables=coarse_variables, + coarse_residuals=coarse_residuals, + coarse_feasible=coarse_feasible, + search_budget=search_budget, + local_starts_attempted=len(selected_indices), + distinct_basins=len(candidates), + global_certificate=status == "verified_zero", + message=message, + ) + + +def continue_axis_search( + stages: tuple[AxisSearchProblem, ...], + *, + options: AxisSearchOptions | None = None, + initial_seeds: tuple[jax.Array, ...] = (), +) -> AxisSearchContinuation: + """Warm-start successively richer Fourier problems from the prior best basin.""" + + if not stages: + raise ValueError("At least one Fourier-continuation stage is required.") + results: list[AxisSearchResult] = [] + previous_solution: NearAxisSolution | None = None + for stage_index, stage in enumerate(stages): + if previous_solution is not None: + stage = replace( + stage, + axis=_copy_retained_axis_modes(previous_solution.inputs.axis, stage.axis), + ) + result = search_axis( + stage, + options=options, + seeds=initial_seeds if stage_index == 0 else (), + ) + results.append(result) + if result.best is None: + break + previous_solution = result.best.solution + complete = len(results) == len(stages) and all(result.best is not None for result in results) + return AxisSearchContinuation( + stages=tuple(results), + final=results[-1] if results else None, + complete=complete, + ) diff --git a/src/pyqsc_jax/configurations.py b/src/pyqsc_jax/configurations.py new file mode 100644 index 0000000..a8b32c9 --- /dev/null +++ b/src/pyqsc_jax/configurations.py @@ -0,0 +1,274 @@ +"""Named, immutable near-axis reference configurations.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from pyqsc_jax.first_order import Qsc +from pyqsc_jax.models import NearAxisSolution + + +@dataclass(frozen=True) +class ReferenceConfiguration: + """A documented set of inputs used in regression and example workflows.""" + + name: str + description: str + rc: tuple[float, ...] + zs: tuple[float, ...] + nfp: int + etabar: float + B0: float = 1.0 + B2c: float = 0.0 + I2: float = 0.0 + p2: float = 0.0 + order: str = "r2" + source_database_id: int | None = None + source_url: str | None = None + + def parameters(self, **overrides: Any) -> dict[str, Any]: + """Return fresh keyword arguments, with explicit caller overrides.""" + + parameters: dict[str, Any] = { + "rc": self.rc, + "zs": self.zs, + "nfp": self.nfp, + "etabar": self.etabar, + "B0": self.B0, + "B2c": self.B2c, + "I2": self.I2, + "p2": self.p2, + "order": self.order, + } + parameters.update(overrides) + return parameters + + def solve(self, **overrides: Any) -> NearAxisSolution: + """Solve this configuration with optional input overrides.""" + + return Qsc(**self.parameters(**overrides)) + + +REFERENCE_CONFIGURATIONS = ( + ReferenceConfiguration( + name="qa", + description="Vacuum quasi-axisymmetric reference used for pyQSC parity.", + rc=(1.0, 0.155, 0.0102), + zs=(0.0, 0.154, 0.0111), + nfp=2, + etabar=0.64, + B2c=-0.00322, + ), + ReferenceConfiguration( + name="qh", + description="Vacuum quasi-helically symmetric reference with nonzero frame helicity.", + rc=(1.0, 0.17, 0.01804, 0.001409, 5.877e-5), + zs=(0.0, 0.1581, 0.01820, 0.001548, 7.772e-5), + nfp=4, + etabar=1.569, + B2c=0.1348, + ), + ReferenceConfiguration( + name="finite_pressure_current", + description="Finite-pressure/current r2 reference used throughout plasma-jet validation.", + rc=(1.0, 0.09), + zs=(0.0, -0.09), + nfp=2, + etabar=0.95, + B2c=-0.7, + I2=0.9, + p2=-600000.0, + ), + ReferenceConfiguration( + name="b20_optimized_qa", + description=( + "High-order QA axis refined by exact B2c elimination and bounded " + "least-squares minimization of the dense B20 anomaly." + ), + rc=( + 1.0038581971135636, + 0.15364113635027635, + 0.019066219667598507, + 0.0029880230763716796, + 0.00040790711310020265, + 0.00004610617355685316, + 4.102693907398271e-06, + 5.154300428457222e-07, + 4.8802742243232844e-08, + 7.301132037525988e-09, + ), + zs=( + 0.0, + -0.14606422619055753, + -0.020369711261517005, + -0.003031761163968464, + -0.00039984211636995954, + -0.00005034787759957322, + -4.174376962124085e-06, + -4.557462755956434e-07, + -8.173481495049928e-08, + -3.732477282851326e-09, + ), + nfp=2, + etabar=-0.6783912804454629, + B0=1.006541121335688, + B2c=0.9427285320639192, + ), + ReferenceConfiguration( + name="database_example_3", + description=( + "Curvo stellarator-database configuration 3, retained as the " + "documented download/API example." + ), + rc=(1.0, -0.53677857, -0.046455786, -0.0070183445), + zs=(0.0, -0.5888703, -0.04447083, -0.009581006), + nfp=4, + etabar=1.4014399, + B2c=-0.7512066, + p2=-74635.375, + order="r3", + source_database_id=3, + source_url="https://stellarator.physics.wisc.edu/app/plot/3", + ), + ReferenceConfiguration( + name="database_qa_139524", + description=( + "Curvo stellarator-database configuration 139524: a one-field-period " + "finite-pressure QA stellarator with |iota| above 0.3, zero current, " + "nonzero torsion, and positive margins for the named design screen." + ), + rc=(1.0, -0.06883207, 0.0017516185, 0.023231717), + zs=(0.0, -0.28447896, 0.074662544, 0.07483574), + nfp=1, + etabar=-0.7771866, + B2c=-1.8120022, + I2=0.0, + p2=-232743.75, + order="r3", + source_database_id=139524, + source_url="https://stellarator.physics.wisc.edu/app/plot/139524", + ), + ReferenceConfiguration( + name="database_low_b20_57409", + description=( + "Low-B20 Curvo stellarator-database configuration 57409, used as " + "the traceable seed for the bundled constrained refinement." + ), + rc=(1.0, -0.51677144, -0.009499784, -0.005914526), + zs=(0.0, -0.5420635, -0.012225689, -0.0059485724), + nfp=4, + etabar=-1.3295174, + B2c=-0.7577404, + p2=-23501.281, + order="r3", + source_database_id=57409, + source_url="https://stellarator.physics.wisc.edu/app/plot/57409", + ), + ReferenceConfiguration( + name="b20_optimized_good", + description=( + "Eight-mode QH refinement of database configuration 57409. Exact " + "B2c elimination and staged bounded least-squares flatten B20 while " + "retaining every Curvo Table-3 margin with |iota| above 0.4." + ), + rc=( + 1.0, + -0.5039436500066075, + -0.043965867334349464, + -0.00654919263283669, + -3.047898633100702e-06, + 2.7519909478191202e-05, + 6.071324626161564e-06, + 8.46692978641554e-07, + 5.99743474381913e-08, + ), + zs=( + 0.0, + -0.5050020451312105, + -0.045010140721391825, + -0.006585874304053245, + -1.5550078876295003e-05, + 2.6653739413147386e-05, + 5.978859527401134e-06, + 8.327402553082346e-07, + 5.9109366390399247e-08, + ), + nfp=4, + etabar=-1.3295174, + B2c=-1.132420959333329, + p2=-23501.281, + order="r3", + source_database_id=57409, + source_url="https://stellarator.physics.wisc.edu/app/plot/57409", + ), + ReferenceConfiguration( + name="database_large_singularity_107579", + description=( + "Curvo stellarator-database configuration 107579, selected for its " + "large independently recomputed singular radius." + ), + rc=(1.0, -0.54465365, 0.005908036, 0.0054288576), + zs=(0.0, 0.540987, -0.009328544, -0.003327578), + nfp=3, + etabar=-0.95917624, + B2c=0.10977107, + p2=-19211.062, + order="r3", + source_database_id=107579, + source_url="https://stellarator.physics.wisc.edu/app/plot/107579", + ), + ReferenceConfiguration( + name="plasma_dominant_channel", + description=( + "Circular finite-pressure/current channel for which the matched plasma " + "field exceeds 30 percent of the total field at formal radius 0.2." + ), + rc=(1.0,), + zs=(0.0,), + nfp=1, + etabar=0.5, + I2=4.2, + p2=-100000.0, + ), + ReferenceConfiguration( + name="plasma_stellarator", + description=( + "Curvo stellarator-database configuration 52521: a strongly " + "nonplanar finite-pressure, zero-current stellarator used for " + "the plasma/external and VMEX showcases." + ), + rc=(1.0, -0.5415884, 0.029195854, 0.0048646266), + zs=(0.0, -0.57113713, 0.029922731, 0.0041398546), + nfp=4, + etabar=1.1396117, + B2c=-0.050057083, + I2=0.0, + p2=-28248.188, + order="r3", + source_database_id=52521, + source_url="https://stellarator.physics.wisc.edu/app/plot/52521", + ), +) + + +def available_configurations() -> tuple[str, ...]: + """Return the stable names of bundled reference configurations.""" + + return tuple(configuration.name for configuration in REFERENCE_CONFIGURATIONS) + + +def get_configuration(name: str) -> ReferenceConfiguration: + """Return one named reference configuration.""" + + for configuration in REFERENCE_CONFIGURATIONS: + if configuration.name == name: + return configuration + available = ", ".join(available_configurations()) + raise ValueError(f"Unknown configuration {name!r}. Available configurations: {available}.") + + +def solve_configuration(name: str, **overrides: Any) -> NearAxisSolution: + """Solve a named reference configuration.""" + + return get_configuration(name).solve(**overrides) diff --git a/src/pyqsc_jax/continuation.py b/src/pyqsc_jax/continuation.py new file mode 100644 index 0000000..90c5d60 --- /dev/null +++ b/src/pyqsc_jax/continuation.py @@ -0,0 +1,229 @@ +"""Pseudo-arclength continuation of first-order near-axis branches.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +import jax +import jax.numpy as jnp + +from pyqsc_jax.axis import Axis +from pyqsc_jax.first_order import first_order_solution, sigma_equation, solve +from pyqsc_jax.inverse import parameter_response_derivative +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.solvers import DEFAULT_ROOT_OPTIONS, RootSolveOptions, implicit_dense_root + +ArrayLike = Any + + +@dataclass(frozen=True) +class ContinuationResult: + """A variable-length sequence of corrected pseudo-arclength points.""" + + solutions: tuple[NearAxisSolution, ...] + tangents: jax.Array + response_derivative: jax.Array + fold_detected: jax.Array + requested_points: int + status: str + + @property + def etabar(self) -> jax.Array: + """Etabar values along the corrected branch.""" + + return jnp.stack(tuple(solution.inputs.etabar for solution in self.solutions)) + + @property + def iota(self) -> jax.Array: + """Rotational-transform values along the corrected branch.""" + + return jnp.stack(tuple(solution.iota for solution in self.solutions)) + + @property + def sigma(self) -> jax.Array: + """Periodic sigma samples along the corrected branch.""" + + return jnp.stack(tuple(solution.sigma for solution in self.solutions)) + + @property + def complete(self) -> bool: + """Whether all requested points converged.""" + + return self.status == "complete" and len(self.solutions) == self.requested_points + + +def continue_etabar_branch( + *, + axis: Axis, + etabar_start: ArrayLike, + etabar_next: ArrayLike, + num_points: int = 25, + step_size: float | None = None, + B0: ArrayLike = 1.0, + sigma0: ArrayLike = 0.0, + I2: ArrayLike = 0.0, + nphi: int = 61, + sG: int = 1, + spsi: int = 1, + root_options: RootSolveOptions = DEFAULT_ROOT_OPTIONS, + fold_tolerance: float = 1e-4, +) -> ContinuationResult: + """Trace a fixed-sign etabar branch through folds in iota(etabar).""" + + if not isinstance(num_points, int) or isinstance(num_points, bool) or num_points < 2: + raise ValueError("num_points must be an integer >= 2.") + if step_size is not None and step_size <= 0: + raise ValueError("step_size must be positive.") + if fold_tolerance < 0: + raise ValueError("fold_tolerance must be nonnegative.") + + etabar_start = jnp.asarray(etabar_start) + etabar_next = jnp.asarray(etabar_next) + if etabar_start.ndim != 0 or etabar_next.ndim != 0: + raise ValueError("etabar_start and etabar_next must be scalars.") + sign = jnp.where(etabar_start < 0, -1.0, 1.0) + if bool(etabar_start == 0) or bool(etabar_next == 0): + raise ValueError("Continuation seeds must have nonzero etabar.") + if bool(jnp.sign(etabar_start) != jnp.sign(etabar_next)): + raise ValueError("Continuation seeds must have the same etabar sign.") + + common = { + "axis": axis, + "B0": B0, + "sigma0": sigma0, + "I2": I2, + "nphi": nphi, + "order": "r1", + "sG": sG, + "spsi": spsi, + } + first = solve(etabar=etabar_start, **common) + second = solve(etabar=etabar_next, **common) + if not bool(first.root_report.converged) or not bool(second.root_report.converged): + raise ValueError("Both initial continuation points must converge.") + + solutions = [first, second] + pair0 = jnp.asarray((first.inputs.etabar, first.iota)) + pair1 = jnp.asarray((second.inputs.etabar, second.iota)) + secant = pair1 - pair0 + secant_norm = jnp.linalg.norm(secant) + if not bool(jnp.isfinite(secant_norm)) or bool(secant_norm == 0): + raise ValueError("Initial continuation points must be finite and distinct.") + tangent = secant / secant_norm + arclength_step = float(secant_norm) if step_size is None else step_size + tangents = [tangent, tangent] + responses = [ + parameter_response_derivative( + first.inputs, + first.geometry, + first.sigma, + first.iota, + parameter="etabar", + ), + parameter_response_derivative( + second.inputs, + second.geometry, + second.sigma, + second.iota, + parameter="etabar", + ), + ] + fold_flags = [ + jnp.abs(responses[0]) <= fold_tolerance, + (jnp.abs(responses[1]) <= fold_tolerance) | (responses[0] * responses[1] < 0), + ] + status = "complete" + + for _ in range(2, num_points): + previous = solutions[-1] + before_previous = solutions[-2] + previous_pair = jnp.asarray((previous.inputs.etabar, previous.iota)) + predicted_pair = previous_pair + arclength_step * tangent + predicted_etabar = predicted_pair[0] + if bool(sign * predicted_etabar <= 0): + status = "branch_zero_crossing" + break + + sigma_scale = arclength_step / jnp.linalg.norm( + previous_pair - jnp.asarray((before_previous.inputs.etabar, before_previous.iota)) + ) + predicted_sigma = previous.sigma + sigma_scale * (previous.sigma - before_previous.sigma) + initial_state = jnp.concatenate( + ( + jnp.log(jnp.abs(predicted_etabar))[None], + predicted_pair[1:], + predicted_sigma[1:], + ) + ) + base_inputs = previous.inputs + + def residual( + state, + base_inputs=base_inputs, + geometry=previous.geometry, + predicted_pair=predicted_pair, + tangent=tangent, + ): + etabar = sign * jnp.exp(state[0]) + iota = state[1] + sigma = jnp.concatenate((base_inputs.sigma0[None], state[2:])) + local_inputs = replace(base_inputs, etabar=etabar) + sigma_part = sigma_equation( + sigma, + iota, + inputs=local_inputs, + geometry=geometry, + ) + arclength_part = jnp.dot( + jnp.asarray((etabar, iota)) - predicted_pair, + tangent, + ) + return jnp.concatenate((sigma_part, arclength_part[None])) + + state, report = implicit_dense_root( + residual, + initial_state, + options=root_options, + ) + corrected_etabar = sign * jnp.exp(state[0]) + corrected_iota = state[1] + corrected_sigma = jnp.concatenate((base_inputs.sigma0[None], state[2:])) + corrected_inputs = replace(base_inputs, etabar=corrected_etabar) + corrected = first_order_solution( + corrected_inputs, + previous.geometry, + corrected_sigma, + corrected_iota, + report, + ) + if not bool(report.converged): + status = "solver_failure" + break + + corrected_pair = jnp.asarray((corrected_etabar, corrected_iota)) + new_tangent = corrected_pair - previous_pair + new_tangent = new_tangent / jnp.linalg.norm(new_tangent) + new_tangent = jnp.where(jnp.dot(new_tangent, tangent) < 0, -new_tangent, new_tangent) + response = parameter_response_derivative( + corrected.inputs, + corrected.geometry, + corrected.sigma, + corrected.iota, + parameter="etabar", + ) + fold = (jnp.abs(response) <= fold_tolerance) | (responses[-1] * response < 0) + solutions.append(corrected) + tangents.append(new_tangent) + responses.append(response) + fold_flags.append(fold) + tangent = new_tangent + + return ContinuationResult( + solutions=tuple(solutions), + tangents=jnp.stack(tuple(tangents)), + response_derivative=jnp.stack(tuple(responses)), + fold_detected=jnp.stack(tuple(fold_flags)), + requested_points=num_points, + status=status, + ) diff --git a/src/pyqsc_jax/criteria.py b/src/pyqsc_jax/criteria.py new file mode 100644 index 0000000..fb49a27 --- /dev/null +++ b/src/pyqsc_jax/criteria.py @@ -0,0 +1,200 @@ +"""Named, scalable stellarator design-criteria profiles.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import jax +import jax.numpy as jnp + +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.second_order import MU0 + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class CriterionEvaluation: + """One signed criterion margin and its pass/fail result.""" + + value: jax.Array + threshold: jax.Array + margin: jax.Array + passed: jax.Array + name: str = field(metadata={"static": True}) + sense: str = field(metadata={"static": True}) + units: str = field(metadata={"static": True}) + + +@dataclass(frozen=True) +class CriteriaReport: + """Evaluation of every item in one named criteria profile.""" + + profile_name: str + evaluations: tuple[CriterionEvaluation, ...] + + @property + def passed(self) -> bool: + """Whether every criterion passes.""" + + return all(bool(evaluation.passed) for evaluation in self.evaluations) + + @property + def margins(self) -> dict[str, jax.Array]: + """Signed raw margins keyed by criterion name.""" + + return {evaluation.name: evaluation.margin for evaluation in self.evaluations} + + @property + def values(self) -> dict[str, jax.Array]: + """Measured values keyed by criterion name.""" + + return {evaluation.name: evaluation.value for evaluation in self.evaluations} + + def __getitem__(self, name: str) -> CriterionEvaluation: + for evaluation in self.evaluations: + if evaluation.name == name: + return evaluation + raise KeyError(name) + + +@dataclass(frozen=True) +class Criteria: + """Configurable thresholds for a named near-axis design profile.""" + + minimum_axis_length: float + minimum_abs_iota: float + maximum_elongation: float + minimum_L_grad_B: float + minimum_axis_radius: float + minimum_singular_radius: float + minimum_L_grad_grad_B: float + maximum_B20_variation: float + minimum_beta: float + minimum_DMerc_times_r2: float + profile_name: str = "custom" + + @classmethod + def from_curvo_2025( + cls, + *, + major_radius: float = 1.0, + B0: float = 1.0, + **overrides: float, + ) -> Criteria: + """Return the scalable version of Curvo et al. (2025), Table 3.""" + + if major_radius <= 0: + raise ValueError("major_radius must be positive.") + if B0 <= 0: + raise ValueError("B0 must be positive.") + parameters = { + "minimum_axis_length": 0.0, + "minimum_abs_iota": 0.2, + "maximum_elongation": 10.0, + "minimum_L_grad_B": 0.1 * major_radius, + "minimum_axis_radius": 0.3 * major_radius, + "minimum_singular_radius": 0.05 * major_radius, + "minimum_L_grad_grad_B": 0.1 * major_radius, + "maximum_B20_variation": 5.0 * B0 / major_radius**2, + "minimum_beta": 1.0e-4, + "minimum_DMerc_times_r2": 0.0, + "profile_name": "curvo_2025", + } + unknown = set(overrides) - set(parameters) + if unknown: + names = ", ".join(sorted(unknown)) + raise ValueError(f"Unknown criteria override(s): {names}.") + parameters.update(overrides) + return cls(**parameters) + + def evaluate(self, solution: NearAxisSolution) -> CriteriaReport: + """Evaluate values, signed margins, and pass flags.""" + + if solution.second_order is None: + raise ValueError("The criteria profile requires a second-order solution.") + beta = -MU0 * solution.inputs.p2 * solution.r_singularity**2 / solution.inputs.B0**2 + values = ( + ( + "axis_length", + solution.axis_length, + self.minimum_axis_length, + "strict_min", + "m", + ), + ("abs_iota", jnp.abs(solution.iota), self.minimum_abs_iota, "min", "1"), + ( + "maximum_elongation", + jnp.max(solution.elongation), + self.maximum_elongation, + "max", + "1", + ), + ( + "minimum_L_grad_B", + jnp.min(solution.L_grad_B), + self.minimum_L_grad_B, + "min", + "m", + ), + ( + "minimum_axis_radius", + jnp.min(solution.R0), + self.minimum_axis_radius, + "min", + "m", + ), + ( + "singular_radius", + solution.r_singularity, + self.minimum_singular_radius, + "min", + "m", + ), + ( + "minimum_L_grad_grad_B", + jnp.min(solution.L_grad_grad_B), + self.minimum_L_grad_grad_B, + "min", + "m", + ), + ( + "B20_variation", + solution.B20_variation, + self.maximum_B20_variation, + "max", + "T/m^2", + ), + ("beta", beta, self.minimum_beta, "min", "1"), + ( + "DMerc_times_r2", + solution.DMerc_times_r2, + self.minimum_DMerc_times_r2, + "strict_min", + "1", + ), + ) + evaluations = [] + for name, value, threshold, sense, units in values: + value = jnp.asarray(value) + threshold = jnp.asarray(threshold, dtype=value.dtype) + if sense == "max": + margin = threshold - value + passed = value <= threshold + else: + margin = value - threshold + passed = value > threshold if sense == "strict_min" else value >= threshold + evaluations.append( + CriterionEvaluation( + value=value, + threshold=threshold, + margin=margin, + passed=passed, + name=name, + sense=sense, + units=units, + ) + ) + return CriteriaReport( + profile_name=self.profile_name, + evaluations=tuple(evaluations), + ) diff --git a/src/pyqsc_jax/diagnostics.py b/src/pyqsc_jax/diagnostics.py new file mode 100644 index 0000000..dd77ba0 --- /dev/null +++ b/src/pyqsc_jax/diagnostics.py @@ -0,0 +1,76 @@ +"""Near-axis equilibrium diagnostics.""" + +from __future__ import annotations + +import jax.numpy as jnp + +from pyqsc_jax.models import MercierDiagnostics, NearAxisSolution +from pyqsc_jax.second_order import MU0 + + +def mercier_diagnostics(solution: NearAxisSolution) -> MercierDiagnostics: + """Compute the leading magnetic-well and Mercier terms. + + The normalization and signs follow the standard pyQSC implementation. + """ + + r2 = solution.second_order + if r2 is None: + raise ValueError("A second-order solution is required for Mercier diagnostics.") + + inputs = solution.inputs + geometry = solution.geometry + etabar_squared = inputs.etabar**2 + curvature_squared = solution.curvature**2 + numerator = ( + etabar_squared**2 + + curvature_squared**2 * solution.sigma**2 + + etabar_squared * curvature_squared + ) + denominator = ( + etabar_squared**2 + + curvature_squared**2 * (1 + solution.sigma**2) + + 2 * etabar_squared * curvature_squared + ) + integrand = geometry.d_l_d_phi * numerator / denominator + weighted_integral = ( + jnp.sum(integrand) + * (2 * jnp.pi / (inputs.axis.nfp * inputs.nphi)) + * inputs.axis.nfp + * 2 + * jnp.pi + / geometry.axis_length + ) + DGeod_times_r2 = ( + -2 + * MU0**2 + * inputs.p2**2 + * solution.G0**4 + * etabar_squared + / (jnp.pi**3 * inputs.B0**10 * solution.iotaN**2) + * weighted_integral + ) + d2_volume_d_psi2 = ( + 4 + * jnp.pi**2 + * jnp.abs(solution.G0) + / inputs.B0**3 + * ( + 3 * etabar_squared + - 4 * r2.B20_mean / inputs.B0 + + 2 * (r2.G2 + solution.iota * inputs.I2) / solution.G0 + ) + ) + DWell_times_r2 = ( + MU0 + * inputs.p2 + * jnp.abs(solution.G0) + / (8 * jnp.pi**4 * inputs.B0**3) + * (d2_volume_d_psi2 - 8 * jnp.pi**2 * MU0 * inputs.p2 * jnp.abs(solution.G0) / inputs.B0**5) + ) + return MercierDiagnostics( + d2_volume_d_psi2=d2_volume_d_psi2, + DGeod_times_r2=DGeod_times_r2, + DWell_times_r2=DWell_times_r2, + DMerc_times_r2=DWell_times_r2 + DGeod_times_r2, + ) diff --git a/src/pyqsc_jax/field.py b/src/pyqsc_jax/field.py new file mode 100644 index 0000000..6ad22bd --- /dev/null +++ b/src/pyqsc_jax/field.py @@ -0,0 +1,232 @@ +"""Surface-free total magnetic field jets on the magnetic axis.""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp + +from pyqsc_jax.geometry import cylindrical_vector_to_cartesian +from pyqsc_jax.models import FieldJet, NearAxisSolution + + +def _differentiate_cylindrical_vector( + vector: jax.Array, + solution: NearAxisSolution, +) -> jax.Array: + """Differentiate a vector with respect to Boozer ``varphi``. + + Cylindrical components are periodic over one field period, while fixed + Cartesian components generally are not. The two connection terms below + account for rotation of the cylindrical basis. + """ + + derivative = solution.geometry.d_d_varphi @ vector + d_phi_d_varphi = 1 / solution.geometry.d_varphi_d_phi + derivative = derivative.at[:, 0].add(-d_phi_d_varphi * vector[:, 1]) + return derivative.at[:, 1].add(d_phi_d_varphi * vector[:, 0]) + + +def _to_cartesian(vector: jax.Array, solution: NearAxisSolution) -> jax.Array: + return cylindrical_vector_to_cartesian(vector, solution.phi) + + +def _regular_map_vectors( + solution: NearAxisSolution, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Return ``d1, d2, h11, h12, h22`` in periodic cylindrical components.""" + + r2 = solution.second_order + if r2 is None: + raise ValueError("A second-order solution is required for the regular coordinate map.") + geometry = solution.geometry + tangent = geometry.tangent_cylindrical + normal = geometry.normal_cylindrical + binormal = geometry.binormal_cylindrical + d1 = solution.X1c[:, None] * normal + solution.Y1c[:, None] * binormal + d2 = solution.Y1s[:, None] * binormal + h11 = ( + 2 * (r2.X20 + r2.X2c)[:, None] * normal + + 2 * (r2.Y20 + r2.Y2c)[:, None] * binormal + + 2 * (r2.Z20 + r2.Z2c)[:, None] * tangent + ) + h12 = ( + 2 * r2.X2s[:, None] * normal + + 2 * r2.Y2s[:, None] * binormal + + 2 * r2.Z2s[:, None] * tangent + ) + h22 = ( + 2 * (r2.X20 - r2.X2c)[:, None] * normal + + 2 * (r2.Y20 - r2.Y2c)[:, None] * binormal + + 2 * (r2.Z20 - r2.Z2c)[:, None] * tangent + ) + return d1, d2, h11, h12, h22 + + +def total_field_jet(solution: NearAxisSolution) -> FieldJet: + """Compute ``B``, ``grad(B)``, and ``grad(grad(B))`` on the axis. + + This is the regular-coordinate chain rule in equations (55)--(83) of the + surface-free plasma/coil derivation. It requires the complete second-order + near-axis solution but no finite-radius surface. + """ + + r2 = solution.second_order + if r2 is None: + raise ValueError("A second-order solution is required for the field Hessian.") + + geometry = solution.geometry + inputs = solution.inputs + length_scale = geometry.abs_G0_over_B0 + tangent = geometry.tangent_cylindrical + d1, d2, h11, h12, h22 = _regular_map_vectors(solution) + + derivative = lambda vector: _differentiate_cylindrical_vector(vector, solution) # noqa: E731 + V0 = length_scale * tangent + V1 = derivative(d1) + solution.iotaN * d2 + V2 = derivative(d2) - solution.iotaN * d1 + V11 = derivative(h11) + 2 * solution.iotaN * h12 + V12 = derivative(h12) + solution.iotaN * (h22 - h11) + V22 = derivative(h22) - 2 * solution.iotaN * h12 + + p0 = inputs.B0**2 / solution.G0 + p1 = 2 * inputs.B0**2 * inputs.etabar / solution.G0 + flux_term = inputs.B0**2 * (r2.G2 + solution.iota * inputs.I2) / solution.G0**2 + C11 = ( + inputs.B0**2 * inputs.etabar**2 + 2 * inputs.B0 * (r2.B20 + inputs.B2c) + ) / solution.G0 - flux_term + C22 = 2 * inputs.B0 * (r2.B20 - inputs.B2c) / solution.G0 - flux_term + + field_coordinate_gradient_cylindrical = jnp.stack( + ( + p0 * derivative(V0), + p1 * V0 + p0 * V1, + p0 * V2, + ), + axis=-1, + ) + field_coordinate_hessian_cylindrical = jnp.zeros( + (inputs.nphi, 3, 3, 3), + dtype=field_coordinate_gradient_cylindrical.dtype, + ) + coordinate_hessian_values = { + (0, 0): p0 * derivative(derivative(V0)), + (0, 1): p1 * derivative(V0) + p0 * derivative(V1), + (0, 2): p0 * derivative(V2), + (1, 1): 2 * C11[:, None] * V0 + 2 * p1 * V1 + p0 * V11, + (1, 2): p1 * V2 + p0 * V12, + (2, 2): 2 * C22[:, None] * V0 + p0 * V22, + } + for (first, second), value in coordinate_hessian_values.items(): + field_coordinate_hessian_cylindrical = field_coordinate_hessian_cylindrical.at[ + :, :, first, second + ].set(value) + field_coordinate_hessian_cylindrical = field_coordinate_hessian_cylindrical.at[ + :, :, second, first + ].set(value) + + field_coordinate_gradient = jnp.stack( + [ + _to_cartesian(field_coordinate_gradient_cylindrical[:, :, index], solution) + for index in range(3) + ], + axis=-1, + ) + field_coordinate_hessian = jnp.stack( + [ + jnp.stack( + [ + _to_cartesian( + field_coordinate_hessian_cylindrical[:, :, first, second], + solution, + ) + for second in range(3) + ], + axis=-1, + ) + for first in range(3) + ], + axis=-1, + ) + + coordinate_jacobian = jnp.stack( + ( + length_scale * geometry.tangent_cartesian, + _to_cartesian(d1, solution), + _to_cartesian(d2, solution), + ), + axis=-1, + ) + coordinate_hessian = jnp.zeros_like(field_coordinate_hessian) + map_hessian_values = { + (0, 0): length_scale**2 * solution.curvature[:, None] * geometry.normal_cartesian, + (0, 1): _to_cartesian(derivative(d1), solution), + (0, 2): _to_cartesian(derivative(d2), solution), + (1, 1): _to_cartesian(h11, solution), + (1, 2): _to_cartesian(h12, solution), + (2, 2): _to_cartesian(h22, solution), + } + for (first, second), value in map_hessian_values.items(): + coordinate_hessian = coordinate_hessian.at[:, :, first, second].set(value) + coordinate_hessian = coordinate_hessian.at[:, :, second, first].set(value) + + inverse_coordinate_jacobian = jnp.linalg.inv(coordinate_jacobian) + gradient = jnp.einsum( + "nia,naj->nij", + field_coordinate_gradient, + inverse_coordinate_jacobian, + ) + hessian = jnp.einsum( + "niab,naj,nbk->nijk", + field_coordinate_hessian, + inverse_coordinate_jacobian, + inverse_coordinate_jacobian, + ) - jnp.einsum( + "nia,nal,nlcb,ncj,nbk->nijk", + field_coordinate_gradient, + inverse_coordinate_jacobian, + coordinate_hessian, + inverse_coordinate_jacobian, + inverse_coordinate_jacobian, + ) + field = p0 * _to_cartesian(V0, solution) + + frenet_basis = jnp.stack( + ( + geometry.normal_cartesian, + geometry.binormal_cartesian, + geometry.tangent_cartesian, + ), + axis=-1, + ) + field_first_frenet = jnp.einsum( + "nia,nijk,njb,nkc->nabc", + frenet_basis, + hessian, + frenet_basis, + frenet_basis, + ) + hessian_frenet = jnp.transpose(field_first_frenet, (0, 2, 3, 1)) + + norm_squared = jnp.sum(hessian**2, axis=(1, 2, 3)) + inverse_scale_vs_varphi = jnp.sqrt(jnp.sqrt(norm_squared) / (4 * inputs.B0)) + coordinate_determinant = jnp.linalg.det(coordinate_jacobian) + divergence = jnp.trace(gradient, axis1=1, axis2=2) + divergence_gradient = jnp.einsum("niik->nk", hessian) + return FieldJet( + field=field, + gradient=gradient, + hessian=hessian, + hessian_frenet=hessian_frenet, + coordinate_jacobian=coordinate_jacobian, + inverse_coordinate_jacobian=inverse_coordinate_jacobian, + coordinate_hessian=coordinate_hessian, + minimum_absolute_coordinate_jacobian=jnp.min(jnp.abs(coordinate_determinant)), + maximum_field_error=jnp.max(jnp.abs(field - solution.B_axis)), + maximum_gradient_error=jnp.max(jnp.abs(gradient - solution.grad_B_axis)), + maximum_divergence=jnp.max(jnp.abs(divergence)), + maximum_derivative_asymmetry=jnp.max(jnp.abs(hessian - jnp.swapaxes(hessian, 2, 3))), + maximum_divergence_gradient=jnp.max(jnp.abs(divergence_gradient)), + grad_grad_B_inverse_scale_length_vs_varphi=inverse_scale_vs_varphi, + L_grad_grad_B=1 / inverse_scale_vs_varphi, + grad_grad_B_inverse_scale_length=jnp.max(inverse_scale_vs_varphi), + ) diff --git a/src/pyqsc_jax/first_order.py b/src/pyqsc_jax/first_order.py new file mode 100644 index 0000000..5443bf9 --- /dev/null +++ b/src/pyqsc_jax/first_order.py @@ -0,0 +1,325 @@ +"""First-order quasisymmetric near-axis construction.""" + +from dataclasses import replace +from typing import Any + +import jax +import jax.numpy as jnp + +from pyqsc_jax.axis import Axis +from pyqsc_jax.geometry import AxisGeometry, compute_axis_geometry +from pyqsc_jax.models import NearAxisInputs, NearAxisSolution, RootSolveReport +from pyqsc_jax.solvers import DEFAULT_ROOT_OPTIONS, RootSolveOptions, implicit_dense_root + +ArrayLike = Any + + +def sigma_from_state(state: ArrayLike, sigma0: ArrayLike) -> jax.Array: + """Replace the state-vector iota slot by the prescribed ``sigma(0)``.""" + + return jnp.asarray(state).at[0].set(jnp.asarray(sigma0)) + + +def sigma_residual( + state: ArrayLike, + *, + inputs: NearAxisInputs, + geometry: AxisGeometry, +) -> jax.Array: + """Periodic first-order sigma-equation residual.""" + + state = jnp.asarray(state) + sigma = sigma_from_state(state, inputs.sigma0) + iota = state[0] + return sigma_equation( + sigma, + iota, + inputs=inputs, + geometry=geometry, + ) + + +def sigma_equation( + sigma: ArrayLike, + iota: ArrayLike, + *, + inputs: NearAxisInputs, + geometry: AxisGeometry, +) -> jax.Array: + """Periodic sigma equation for explicit sigma, iota, and inputs.""" + + sigma = jnp.asarray(sigma) + iota = jnp.asarray(iota) + helicity = geometry.frame_helicity * inputs.spsi * inputs.sG + iotaN = iota + helicity * inputs.axis.nfp + eta_over_curvature_squared = inputs.etabar**2 / geometry.curvature**2 + G0_over_B0 = inputs.sG * geometry.abs_G0_over_B0 + return ( + geometry.d_d_varphi @ sigma + + iotaN * (eta_over_curvature_squared**2 + 1 + sigma**2) + - 2 + * eta_over_curvature_squared + * (-inputs.spsi * geometry.torsion + inputs.I2 / inputs.B0) + * G0_over_B0 + ) + + +def solve_sigma( + inputs: NearAxisInputs, + geometry: AxisGeometry, + *, + root_options: RootSolveOptions = DEFAULT_ROOT_OPTIONS, +) -> tuple[jax.Array, jax.Array, RootSolveReport]: + """Solve for periodic sigma and rotational transform.""" + + initial_state = jnp.full((inputs.nphi,), inputs.sigma0) + initial_state = initial_state.at[0].set(0.0) + residual_function = lambda state: sigma_residual( # noqa: E731 + state, + inputs=inputs, + geometry=geometry, + ) + state, report = implicit_dense_root( + residual_function, + initial_state, + options=root_options, + ) + return sigma_from_state(state, inputs.sigma0), state[0], report + + +def _assemble_gradient( + tangent: jax.Array, + normal: jax.Array, + binormal: jax.Array, + *, + nn: jax.Array, + bn: jax.Array, + nb: jax.Array, + bb: jax.Array, + tn: jax.Array, + nt: jax.Array, + tt: jax.Array, +) -> jax.Array: + outer = lambda left, right: jnp.einsum("ni,nj->nij", left, right) # noqa: E731 + pyqsc_component_expression = ( + nn[:, None, None] * outer(normal, normal) + + bn[:, None, None] * outer(binormal, normal) + + nb[:, None, None] * outer(normal, binormal) + + bb[:, None, None] * outer(binormal, binormal) + + tn[:, None, None] * outer(tangent, normal) + + nt[:, None, None] * outer(normal, tangent) + + tt[:, None, None] * outer(tangent, tangent) + ) + return jnp.swapaxes(pyqsc_component_expression, -1, -2) + + +def first_order_solution( + inputs: NearAxisInputs, + geometry: AxisGeometry, + sigma: jax.Array, + iota: jax.Array, + root_report: RootSolveReport, +) -> NearAxisSolution: + """Assemble first-order shape coefficients, fields, and diagnostics.""" + + helicity = geometry.frame_helicity * inputs.spsi * inputs.sG + iotaN = iota + helicity * inputs.axis.nfp + G0 = inputs.sG * geometry.abs_G0_over_B0 * inputs.B0 + X1s = jnp.zeros_like(geometry.curvature) + X1c = inputs.etabar / geometry.curvature + Y1s = inputs.sG * inputs.spsi * geometry.curvature / inputs.etabar + Y1c = inputs.sG * inputs.spsi * geometry.curvature * sigma / inputs.etabar + + untwisting_angle = -helicity * inputs.axis.nfp * geometry.varphi + sine = jnp.sin(untwisting_angle) + cosine = jnp.cos(untwisting_angle) + X1s_untwisted = X1s * cosine + X1c * sine + X1c_untwisted = -X1s * sine + X1c * cosine + Y1s_untwisted = Y1s * cosine + Y1c * sine + Y1c_untwisted = -Y1s * sine + Y1c * cosine + + p = X1s**2 + X1c**2 + Y1s**2 + Y1c**2 + q = X1s * Y1c - X1c * Y1s + elongation = (p + jnp.sqrt(p**2 - 4 * q**2)) / (2 * jnp.abs(q)) + mean_elongation = jnp.sum(elongation * geometry.d_l_d_phi) / jnp.sum(geometry.d_l_d_phi) + + d_X1c_d_varphi = geometry.d_d_varphi @ X1c + d_Y1s_d_varphi = geometry.d_d_varphi @ Y1s + d_Y1c_d_varphi = geometry.d_d_varphi @ Y1c + factor = inputs.spsi * inputs.B0 / geometry.abs_G0_over_B0 + tn = inputs.sG * inputs.B0 * geometry.curvature + nt = tn + bb = factor * (X1c * d_Y1s_d_varphi - iotaN * X1c * Y1c) + nn = factor * (d_X1c_d_varphi * Y1s + iotaN * X1c * Y1c) + bn = factor * ( + -inputs.sG * inputs.spsi * geometry.abs_G0_over_B0 * geometry.torsion - iotaN * X1c**2 + ) + nb = factor * ( + d_Y1c_d_varphi * Y1s + - d_Y1s_d_varphi * Y1c + + inputs.sG * inputs.spsi * geometry.abs_G0_over_B0 * geometry.torsion + + iotaN * (Y1s**2 + Y1c**2) + ) + tt = jnp.zeros_like(tn) + grad_B_axis_cylindrical = _assemble_gradient( + geometry.tangent_cylindrical, + geometry.normal_cylindrical, + geometry.binormal_cylindrical, + nn=nn, + bn=bn, + nb=nb, + bb=bb, + tn=tn, + nt=nt, + tt=tt, + ) + grad_B_axis = _assemble_gradient( + geometry.tangent_cartesian, + geometry.normal_cartesian, + geometry.binormal_cartesian, + nn=nn, + bn=bn, + nb=nb, + bb=bb, + tn=tn, + nt=nt, + tt=tt, + ) + grad_B_frobenius_squared = jnp.sum(grad_B_axis**2, axis=(-2, -1)) + L_grad_B = inputs.B0 * jnp.sqrt(2 / grad_B_frobenius_squared) + B_axis_cylindrical = inputs.sG * inputs.B0 * geometry.tangent_cylindrical + B_axis = inputs.sG * inputs.B0 * geometry.tangent_cartesian + + return NearAxisSolution( + inputs=inputs, + geometry=geometry, + root_report=root_report, + sigma=sigma, + iota=iota, + iotaN=iotaN, + helicity=helicity, + G0=G0, + X1s=X1s, + X1c=X1c, + Y1s=Y1s, + Y1c=Y1c, + X1s_untwisted=X1s_untwisted, + X1c_untwisted=X1c_untwisted, + Y1s_untwisted=Y1s_untwisted, + Y1c_untwisted=Y1c_untwisted, + elongation=elongation, + mean_elongation=mean_elongation, + B_axis_cylindrical=B_axis_cylindrical, + B_axis=B_axis, + grad_B_axis_cylindrical=grad_B_axis_cylindrical, + grad_B_axis=grad_B_axis, + L_grad_B=L_grad_B, + ) + + +def _normalize_order(order: int | str) -> int: + if order in (1, "r1"): + return 1 + if order in (2, "r2"): + return 2 + if order in (3, "r3"): + return 3 + raise ValueError("order must be 1, 2, 3, 'r1', 'r2', or 'r3'.") + + +def solve( + *, + axis: Axis, + etabar: ArrayLike | None = None, + iota: ArrayLike | None = None, + B0: ArrayLike = 1.0, + sigma0: ArrayLike = 0.0, + I2: ArrayLike = 0.0, + p2: ArrayLike = 0.0, + B2c: ArrayLike = 0.0, + B2s: ArrayLike = 0.0, + nphi: int = 61, + order: int | str = 1, + sG: int = 1, + spsi: int = 1, + solve_for: str = "iota", + root_options: RootSolveOptions = DEFAULT_ROOT_OPTIONS, + fold_tolerance: float = 1e-8, +) -> NearAxisSolution: + """Construct an immutable forward or target-transform solution.""" + + normalized_order = _normalize_order(order) + if solve_for not in ("iota", "etabar", "I2"): + raise ValueError("solve_for must be 'iota', 'etabar', or 'I2'.") + if fold_tolerance < 0: + raise ValueError("fold_tolerance must be nonnegative.") + if solve_for == "iota": + if etabar is None: + raise ValueError("etabar is required when solve_for='iota'.") + if iota is not None: + raise ValueError("iota is prescribed only when solve_for is 'etabar' or 'I2'.") + else: + if iota is None: + raise ValueError(f"iota is required when solve_for={solve_for!r}.") + if solve_for == "I2" and etabar is None: + raise ValueError("etabar is required when solve_for='I2'.") + if etabar is None: + etabar = -1.0 + inputs = NearAxisInputs( + axis=axis, + etabar=etabar, + B0=B0, + sigma0=sigma0, + I2=I2, + p2=p2, + B2c=B2c, + B2s=B2s, + nphi=nphi, + order=normalized_order, + sG=sG, + spsi=spsi, + solve_for=solve_for, + ) + geometry = compute_axis_geometry(axis, nphi=nphi) + if solve_for == "iota": + sigma, solved_iota, root_report = solve_sigma( + inputs, + geometry, + root_options=root_options, + ) + inverse = None + else: + from pyqsc_jax.inverse import solve_target_iota + + inputs, sigma, solved_iota, root_report, inverse = solve_target_iota( + inputs, + geometry, + target_iota=iota, + root_options=root_options, + fold_tolerance=fold_tolerance, + ) + solution = first_order_solution(inputs, geometry, sigma, solved_iota, root_report) + solution = replace(solution, inverse=inverse) + if normalized_order >= 2: + from pyqsc_jax.second_order import solve_second_order + + solution = solve_second_order(solution) + if normalized_order == 3: + from pyqsc_jax.third_order import solve_third_order + + solution = solve_third_order(solution) + return solution + + +def Qsc( + rc: ArrayLike, + zs: ArrayLike, + *, + rs: ArrayLike = (), + zc: ArrayLike = (), + nfp: int = 1, + **kwargs: Any, +) -> NearAxisSolution: + """pyQSC-familiar convenience constructor for the immutable solution.""" + + return solve(axis=Axis(rc=rc, rs=rs, zc=zc, zs=zs, nfp=nfp), **kwargs) diff --git a/src/pyqsc_jax/geometry.py b/src/pyqsc_jax/geometry.py new file mode 100644 index 0000000..a3833c8 --- /dev/null +++ b/src/pyqsc_jax/geometry.py @@ -0,0 +1,205 @@ +"""Sampled magnetic-axis geometry and Frenet frames.""" + +from dataclasses import dataclass +from typing import Any + +import jax +import jax.numpy as jnp + +from pyqsc_jax.axis import Axis, AxisSamples, evaluate_axis +from pyqsc_jax.models import GeometryDiagnostics +from pyqsc_jax.spectral import differentiation_matrix, periodic_grid + +ArrayLike = Any + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class AxisGeometry: + """Immutable sampled geometry over one magnetic field period. + + Cylindrical vector components use the final-axis order ``(R, phi, Z)``. + Cartesian vector components use ``(x, y, z)``. + """ + + axis: Axis + samples: AxisSamples + position_cartesian: jax.Array + d_r_d_phi_cylindrical: jax.Array + d2_r_d_phi2_cylindrical: jax.Array + d3_r_d_phi3_cylindrical: jax.Array + d_l_d_phi: jax.Array + axis_length: jax.Array + tangent_cylindrical: jax.Array + normal_cylindrical: jax.Array + binormal_cylindrical: jax.Array + tangent_cartesian: jax.Array + normal_cartesian: jax.Array + binormal_cartesian: jax.Array + curvature: jax.Array + torsion: jax.Array + frame_helicity: jax.Array + varphi: jax.Array + d_varphi_d_phi: jax.Array + d_d_phi: jax.Array + d_d_varphi: jax.Array + abs_G0_over_B0: jax.Array + diagnostics: GeometryDiagnostics + + +def cylindrical_vector_to_cartesian(vector: ArrayLike, phi: ArrayLike) -> jax.Array: + """Convert vectors from cylindrical ``(R, phi, Z)`` to Cartesian basis.""" + + vector = jnp.asarray(vector) + phi = jnp.asarray(phi) + cosine = jnp.cos(phi) + sine = jnp.sin(phi) + return jnp.stack( + ( + vector[..., 0] * cosine - vector[..., 1] * sine, + vector[..., 0] * sine + vector[..., 1] * cosine, + vector[..., 2], + ), + axis=-1, + ) + + +def _frame_helicity(normal_cylindrical: jax.Array) -> jax.Array: + x_positive = normal_cylindrical[:, 0] >= 0 + z_positive = normal_cylindrical[:, 2] >= 0 + quadrant = jnp.where( + x_positive & z_positive, + 1, + jnp.where(~x_positive & z_positive, 2, jnp.where(~x_positive & ~z_positive, 3, 4)), + ) + next_quadrant = jnp.roll(quadrant, -1) + increment = jnp.where( + (quadrant == 4) & (next_quadrant == 1), + 1, + jnp.where( + (quadrant == 1) & (next_quadrant == 4), + -1, + next_quadrant - quadrant, + ), + ) + return jnp.rint(jnp.sum(increment) / 4).astype(jnp.int32) + + +def compute_axis_geometry( + axis: Axis, + *, + nphi: int = 61, + speed_tolerance: ArrayLike = 1e-12, + curvature_tolerance: ArrayLike = 1e-10, + radius_tolerance: ArrayLike = 1e-10, +) -> AxisGeometry: + """Evaluate axis geometry and Frenet validity over one field period.""" + + if not isinstance(nphi, int) or isinstance(nphi, bool) or nphi < 3: + raise ValueError("nphi must be an integer >= 3.") + + period = 2 * jnp.pi / axis.nfp + phi = periodic_grid(nphi, period=period) + samples = evaluate_axis(axis, phi) + cosine = jnp.cos(phi) + sine = jnp.sin(phi) + position_cartesian = jnp.stack((samples.R * cosine, samples.R * sine, samples.Z), axis=-1) + + d_r = jnp.stack((samples.d_R_d_phi, samples.R, samples.d_Z_d_phi), axis=-1) + d2_r = jnp.stack( + ( + samples.d2_R_d_phi2 - samples.R, + 2 * samples.d_R_d_phi, + samples.d2_Z_d_phi2, + ), + axis=-1, + ) + d3_r = jnp.stack( + ( + samples.d3_R_d_phi3 - 3 * samples.d_R_d_phi, + 3 * samples.d2_R_d_phi2 - samples.R, + samples.d3_Z_d_phi3, + ), + axis=-1, + ) + + d_l_d_phi = jnp.linalg.norm(d_r, axis=-1) + d2_l_d_phi2 = jnp.sum(d_r * d2_r, axis=-1) / d_l_d_phi + tangent = d_r / d_l_d_phi[:, None] + d_tangent_d_l = (-d_r * d2_l_d_phi2[:, None] / d_l_d_phi[:, None] + d2_r) / d_l_d_phi[ + :, None + ] ** 2 + curvature = jnp.linalg.norm(d_tangent_d_l, axis=-1) + pointwise_frenet_valid = (d_l_d_phi > speed_tolerance) & (curvature > curvature_tolerance) + normal = jnp.where( + pointwise_frenet_valid[:, None], + d_tangent_d_l / curvature[:, None], + jnp.nan, + ) + binormal = jnp.cross(tangent, normal) + + cross_first_second = jnp.cross(d_r, d2_r) + torsion = jnp.where( + pointwise_frenet_valid, + jnp.sum(d_r * jnp.cross(d2_r, d3_r), axis=-1) / jnp.sum(cross_first_second**2, axis=-1), + jnp.nan, + ) + + d_phi = period / nphi + axis_length = jnp.sum(d_l_d_phi) * d_phi * axis.nfp + B0_over_abs_G0 = nphi / jnp.sum(d_l_d_phi) + abs_G0_over_B0 = 1 / B0_over_abs_G0 + d_varphi_d_phi = B0_over_abs_G0 * d_l_d_phi + d_d_phi = differentiation_matrix(nphi, period=period) + d_d_varphi = d_d_phi / d_varphi_d_phi[:, None] + varphi = jnp.concatenate( + (jnp.zeros(1, dtype=d_l_d_phi.dtype), jnp.cumsum(d_l_d_phi[:-1] + d_l_d_phi[1:])) + ) + varphi = varphi * (0.5 * d_phi * 2 * jnp.pi / axis_length) + + tangent_cartesian = cylindrical_vector_to_cartesian(tangent, phi) + normal_cartesian = cylindrical_vector_to_cartesian(normal, phi) + binormal_cartesian = cylindrical_vector_to_cartesian(binormal, phi) + frame = jnp.stack((tangent_cartesian, normal_cartesian, binormal_cartesian), axis=-2) + gram = frame @ jnp.swapaxes(frame, -1, -2) + identity = jnp.eye(3, dtype=frame.dtype) + orthogonality_error = jnp.max(jnp.abs(gram - identity)) + determinant = jnp.linalg.det(frame) + minimum_speed = jnp.min(d_l_d_phi) + minimum_curvature = jnp.min(curvature) + minimum_radius = jnp.min(samples.R) + diagnostics = GeometryDiagnostics( + minimum_speed=minimum_speed, + minimum_curvature=minimum_curvature, + minimum_cylindrical_radius=minimum_radius, + maximum_frame_orthogonality_error=orthogonality_error, + minimum_frame_determinant=jnp.min(determinant), + frenet_valid=jnp.all(pointwise_frenet_valid) & jnp.isfinite(orthogonality_error), + cylindrical_coordinates_valid=minimum_radius > radius_tolerance, + ) + + return AxisGeometry( + axis=axis, + samples=samples, + position_cartesian=position_cartesian, + d_r_d_phi_cylindrical=d_r, + d2_r_d_phi2_cylindrical=d2_r, + d3_r_d_phi3_cylindrical=d3_r, + d_l_d_phi=d_l_d_phi, + axis_length=axis_length, + tangent_cylindrical=tangent, + normal_cylindrical=normal, + binormal_cylindrical=binormal, + tangent_cartesian=tangent_cartesian, + normal_cartesian=normal_cartesian, + binormal_cartesian=binormal_cartesian, + curvature=curvature, + torsion=torsion, + frame_helicity=_frame_helicity(normal), + varphi=varphi, + d_varphi_d_phi=d_varphi_d_phi, + d_d_phi=d_d_phi, + d_d_varphi=d_d_varphi, + abs_G0_over_B0=abs_G0_over_B0, + diagnostics=diagnostics, + ) diff --git a/src/pyqsc_jax/inverse.py b/src/pyqsc_jax/inverse.py new file mode 100644 index 0000000..e9c60b1 --- /dev/null +++ b/src/pyqsc_jax/inverse.py @@ -0,0 +1,134 @@ +"""Branch-local inverse solves for a prescribed rotational transform.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Any + +import jax +import jax.numpy as jnp + +from pyqsc_jax.first_order import sigma_equation, sigma_residual +from pyqsc_jax.geometry import AxisGeometry +from pyqsc_jax.models import InverseSolveDiagnostics, NearAxisInputs, RootSolveReport +from pyqsc_jax.solvers import DEFAULT_ROOT_OPTIONS, RootSolveOptions, implicit_dense_root + +ArrayLike = Any + + +def parameter_response_derivative( + inputs: NearAxisInputs, + geometry: AxisGeometry, + sigma: jax.Array, + iota: jax.Array, + *, + parameter: str, +) -> jax.Array: + """Return the local forward derivative d(iota)/d(parameter).""" + + if parameter not in ("etabar", "I2"): + raise ValueError("parameter must be 'etabar' or 'I2'.") + parameter_value = getattr(inputs, parameter) + forward_state = jnp.asarray(sigma).at[0].set(iota) + forward_jacobian = jax.jacfwd( + lambda candidate: sigma_residual( + candidate, + inputs=inputs, + geometry=geometry, + ) + )(forward_state) + parameter_derivative = jax.jacfwd( + lambda value: sigma_residual( + forward_state, + inputs=replace(inputs, **{parameter: value}), + geometry=geometry, + ) + )(parameter_value) + return jnp.linalg.solve(forward_jacobian, -parameter_derivative)[0] + + +def solve_target_iota( + inputs: NearAxisInputs, + geometry: AxisGeometry, + *, + target_iota: ArrayLike, + root_options: RootSolveOptions = DEFAULT_ROOT_OPTIONS, + fold_tolerance: float = 1e-8, +) -> tuple[ + NearAxisInputs, + jax.Array, + jax.Array, + RootSolveReport, + InverseSolveDiagnostics, +]: + """Solve the sigma equation at fixed iota for etabar or I2.""" + + target_iota = jnp.asarray(target_iota) + if target_iota.ndim != 0: + raise ValueError("iota must be a scalar.") + if inputs.solve_for not in ("etabar", "I2"): + raise ValueError("Target-iota mode requires solve_for='etabar' or solve_for='I2'.") + + initial_state = jnp.full((inputs.nphi,), inputs.sigma0) + if inputs.solve_for == "etabar": + parameter_sign = jnp.where(inputs.etabar < 0, -1.0, 1.0) + safe_magnitude = jnp.maximum(jnp.abs(inputs.etabar), jnp.finfo(inputs.etabar.dtype).tiny) + initial_state = initial_state.at[0].set(jnp.log(safe_magnitude)) + + def parameter_from_state(state): + return parameter_sign * jnp.exp(state[0]) + + def inputs_from_parameter(parameter): + return replace(inputs, etabar=parameter) + + else: + parameter_sign = jnp.sign(inputs.I2) + initial_state = initial_state.at[0].set(inputs.I2) + + def parameter_from_state(state): + return state[0] + + def inputs_from_parameter(parameter): + return replace(inputs, I2=parameter) + + def residual(state): + parameter = parameter_from_state(state) + local_inputs = inputs_from_parameter(parameter) + sigma = state.at[0].set(inputs.sigma0) + return sigma_equation( + sigma, + target_iota, + inputs=local_inputs, + geometry=geometry, + ) + + state, report = implicit_dense_root( + residual, + initial_state, + options=root_options, + ) + solved_parameter = parameter_from_state(state) + solved_inputs = inputs_from_parameter(solved_parameter) + sigma = state.at[0].set(inputs.sigma0) + + response_derivative = parameter_response_derivative( + solved_inputs, + geometry, + sigma, + target_iota, + parameter=inputs.solve_for, + ) + absolute_response = jnp.abs(response_derivative) + branch_fold = ~jnp.isfinite(response_derivative) | (absolute_response <= fold_tolerance) + diagnostics = InverseSolveDiagnostics( + target_iota=target_iota, + achieved_iota=target_iota, + solved_value=solved_parameter, + response_derivative=response_derivative, + absolute_response_derivative=absolute_response, + fold_tolerance=jnp.asarray(fold_tolerance), + branch_fold=branch_fold, + parameter_sign=jnp.sign(solved_parameter), + parameter=inputs.solve_for, + ) + return solved_inputs, sigma, target_iota, report, diagnostics diff --git a/src/pyqsc_jax/models.py b/src/pyqsc_jax/models.py new file mode 100644 index 0000000..e255869 --- /dev/null +++ b/src/pyqsc_jax/models.py @@ -0,0 +1,504 @@ +"""Small immutable result and diagnostic models.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, ClassVar + +import jax +import jax.numpy as jnp + +from pyqsc_jax.axis import Axis + +if TYPE_CHECKING: + from pyqsc_jax.geometry import AxisGeometry + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class GeometryDiagnostics: + """Validity and conditioning information for sampled axis geometry.""" + + minimum_speed: jax.Array + minimum_curvature: jax.Array + minimum_cylindrical_radius: jax.Array + maximum_frame_orthogonality_error: jax.Array + minimum_frame_determinant: jax.Array + frenet_valid: jax.Array + cylindrical_coordinates_valid: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class RootSolveReport: + """Convergence evidence for a nonlinear root solve.""" + + initial_residual_norm: jax.Array + residual_norm: jax.Array + tolerance: jax.Array + step_norm: jax.Array + iterations: jax.Array + backtracking_steps: jax.Array + jacobian_condition_number: jax.Array + converged: jax.Array + finite: jax.Array + stagnated: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class InverseSolveDiagnostics: + """Local branch and fold diagnostics for a target-transform solve.""" + + target_iota: jax.Array + achieved_iota: jax.Array + solved_value: jax.Array + response_derivative: jax.Array + absolute_response_derivative: jax.Array + fold_tolerance: jax.Array + branch_fold: jax.Array + parameter_sign: jax.Array + parameter: str = field(metadata={"static": True}) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class LinearSolveReport: + """Residual and conditioning evidence for a dense linear solve.""" + + residual_norm: jax.Array + relative_residual_norm: jax.Array + matrix_condition_number: jax.Array + finite: jax.Array + converged: jax.Array + well_conditioned: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class MercierDiagnostics: + """Leading near-axis magnetic-well and Mercier contributions.""" + + d2_volume_d_psi2: jax.Array + DGeod_times_r2: jax.Array + DWell_times_r2: jax.Array + DMerc_times_r2: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class FieldJet: + """Total on-axis magnetic field through two Cartesian derivatives. + + ``gradient`` and ``hessian`` use field-component-first ordering: + ``gradient[n, i, j] = d B_i / d x_j`` and + ``hessian[n, i, j, k] = d² B_i / (d x_j d x_k)``. + """ + + field: jax.Array + gradient: jax.Array + hessian: jax.Array + hessian_frenet: jax.Array + coordinate_jacobian: jax.Array + inverse_coordinate_jacobian: jax.Array + coordinate_hessian: jax.Array + minimum_absolute_coordinate_jacobian: jax.Array + maximum_field_error: jax.Array + maximum_gradient_error: jax.Array + maximum_divergence: jax.Array + maximum_derivative_asymmetry: jax.Array + maximum_divergence_gradient: jax.Array + grad_grad_B_inverse_scale_length_vs_varphi: jax.Array + L_grad_grad_B: jax.Array + grad_grad_B_inverse_scale_length: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class SingularityDiagnostics: + """First loss of regularity in the quadratic near-axis coordinate map.""" + + r_singularity: jax.Array + r_singularity_vs_varphi: jax.Array + inv_r_singularity_vs_varphi: jax.Array + theta_singularity_vs_varphi: jax.Array + residual_norm_vs_varphi: jax.Array + maximum_residual_norm: jax.Array + g0: jax.Array + g1c: jax.Array + g1s: jax.Array + g20: jax.Array + g2s: jax.Array + g2c: jax.Array + angular_resolution: int = field(metadata={"static": True}) + newton_iterations: int = field(metadata={"static": True}) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class NearAxisInputs: + """Normalized immutable inputs to a near-axis solve.""" + + axis: Axis + etabar: jax.Array + B0: jax.Array = 1.0 + sigma0: jax.Array = 0.0 + I2: jax.Array = 0.0 + p2: jax.Array = 0.0 + B2c: jax.Array = 0.0 + B2s: jax.Array = 0.0 + nphi: int = field(default=61, metadata={"static": True}) + order: int = field(default=1, metadata={"static": True}) + sG: int = field(default=1, metadata={"static": True}) + spsi: int = field(default=1, metadata={"static": True}) + solve_for: str = field(default="iota", metadata={"static": True}) + + def __post_init__(self) -> None: + if not isinstance(self.nphi, int) or isinstance(self.nphi, bool) or self.nphi < 3: + raise ValueError("nphi must be an integer >= 3.") + if self.order not in (1, 2, 3): + raise ValueError("order must be 1, 2, or 3.") + if self.sG not in (-1, 1): + raise ValueError("sG must be +1 or -1.") + if self.spsi not in (-1, 1): + raise ValueError("spsi must be +1 or -1.") + if self.solve_for not in ("iota", "etabar", "I2"): + raise ValueError("solve_for must be 'iota', 'etabar', or 'I2'.") + for name in ("etabar", "B0", "sigma0", "I2", "p2", "B2c", "B2s"): + value = jnp.asarray(getattr(self, name)) + if value.ndim: + raise ValueError(f"{name} must be a scalar.") + object.__setattr__(self, name, value) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class SecondOrderData: + """Complete second-order coefficient solution and direct diagnostics.""" + + linear_report: LinearSolveReport + V1: jax.Array + V2: jax.Array + V3: jax.Array + X20: jax.Array + X2s: jax.Array + X2c: jax.Array + Y20: jax.Array + Y2s: jax.Array + Y2c: jax.Array + Z20: jax.Array + Z2s: jax.Array + Z2c: jax.Array + beta_1s: jax.Array + B20: jax.Array + B20_mean: jax.Array + B20_anomaly: jax.Array + B20_residual: jax.Array + B20_variation: jax.Array + G2: jax.Array + N_helicity: jax.Array + d_curvature_d_varphi: jax.Array + d_torsion_d_varphi: jax.Array + d_X20_d_varphi: jax.Array + d_X2s_d_varphi: jax.Array + d_X2c_d_varphi: jax.Array + d_Y20_d_varphi: jax.Array + d_Y2s_d_varphi: jax.Array + d_Y2c_d_varphi: jax.Array + d_Z20_d_varphi: jax.Array + d_Z2s_d_varphi: jax.Array + d_Z2c_d_varphi: jax.Array + d2_X1c_d_varphi2: jax.Array + d2_Y1c_d_varphi2: jax.Array + d2_Y1s_d_varphi2: jax.Array + X20_untwisted: jax.Array + X2s_untwisted: jax.Array + X2c_untwisted: jax.Array + Y20_untwisted: jax.Array + Y2s_untwisted: jax.Array + Y2c_untwisted: jax.Array + Z20_untwisted: jax.Array + Z2s_untwisted: jax.Array + Z2c_untwisted: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class ThirdOrderData: + """Third-order flux-constraint surface corrections.""" + + flux_constraint_coefficient: jax.Array + B0_order_a_squared_to_cancel: jax.Array + flux_constraint_residual: jax.Array + consistency_error: jax.Array + X3s1: jax.Array + X3c1: jax.Array + Y3s1: jax.Array + Y3c1: jax.Array + Z3s1: jax.Array + Z3c1: jax.Array + X3s3: jax.Array + X3c3: jax.Array + Y3s3: jax.Array + Y3c3: jax.Array + Z3s3: jax.Array + Z3c3: jax.Array + d_X3c1_d_varphi: jax.Array + d_Y3s1_d_varphi: jax.Array + d_Y3c1_d_varphi: jax.Array + X3s1_untwisted: jax.Array + X3c1_untwisted: jax.Array + Y3s1_untwisted: jax.Array + Y3c1_untwisted: jax.Array + Z3s1_untwisted: jax.Array + Z3c1_untwisted: jax.Array + X3s3_untwisted: jax.Array + X3c3_untwisted: jax.Array + Y3s3_untwisted: jax.Array + Y3c3_untwisted: jax.Array + Z3s3_untwisted: jax.Array + Z3c3_untwisted: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class ShearData: + """Order-r-squared rotational-transform correction and intermediates.""" + + B31c: jax.Array + iota2: jax.Array + numerator: jax.Array + denominator: jax.Array + Lambda_tilde: jax.Array + integrating_factor: jax.Array + sigma_average: jax.Array + Z31c: jax.Array + Z31s: jax.Array + X31c: jax.Array + X31s: jax.Array + Y31s: jax.Array + stellarator_symmetric: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class NearAxisSolution: + """Canonical immutable near-axis solution. + + Sampled vector arrays use a leading ``nphi`` axis. Vector components are + Cartesian unless a field name explicitly contains ``cylindrical``. + """ + + inputs: NearAxisInputs + geometry: AxisGeometry + root_report: RootSolveReport + sigma: jax.Array + iota: jax.Array + iotaN: jax.Array + helicity: jax.Array + G0: jax.Array + X1s: jax.Array + X1c: jax.Array + Y1s: jax.Array + Y1c: jax.Array + X1s_untwisted: jax.Array + X1c_untwisted: jax.Array + Y1s_untwisted: jax.Array + Y1c_untwisted: jax.Array + elongation: jax.Array + mean_elongation: jax.Array + B_axis_cylindrical: jax.Array + B_axis: jax.Array + grad_B_axis_cylindrical: jax.Array + grad_B_axis: jax.Array + L_grad_B: jax.Array + second_order: SecondOrderData | None = None + mercier: MercierDiagnostics | None = None + field_jet: FieldJet | None = None + singularity: SingularityDiagnostics | None = None + third_order: ThirdOrderData | None = None + shear: ShearData | None = None + inverse: InverseSolveDiagnostics | None = None + + _SECOND_ORDER_NAMES: ClassVar[frozenset[str]] = frozenset( + field.name for field in SecondOrderData.__dataclass_fields__.values() + ) + _THIRD_ORDER_NAMES: ClassVar[frozenset[str]] = frozenset( + field.name for field in ThirdOrderData.__dataclass_fields__.values() + ) + _SHEAR_NAMES: ClassVar[frozenset[str]] = frozenset( + field.name for field in ShearData.__dataclass_fields__.values() + ) + _INVERSE_NAMES: ClassVar[frozenset[str]] = frozenset( + field.name for field in InverseSolveDiagnostics.__dataclass_fields__.values() + ) + _MERCIER_NAMES: ClassVar[frozenset[str]] = frozenset( + field.name for field in MercierDiagnostics.__dataclass_fields__.values() + ) + _FIELD_JET_NAMES: ClassVar[frozenset[str]] = frozenset( + { + "grad_grad_B_axis", + "grad_grad_B", + "L_grad_grad_B", + "grad_grad_B_inverse_scale_length_vs_varphi", + "grad_grad_B_inverse_scale_length", + } + ) + _SINGULARITY_NAMES: ClassVar[frozenset[str]] = frozenset( + { + "r_singularity", + "r_singularity_vs_varphi", + "inv_r_singularity_vs_varphi", + "r_singularity_basic_vs_varphi", + "r_singularity_theta_vs_varphi", + "r_singularity_residual_sqnorm", + } + ) + + def __getattr__(self, name: str): + if name in self._SECOND_ORDER_NAMES: + second_order = object.__getattribute__(self, "second_order") + if second_order is None: + raise AttributeError(f"First-order solution has no {name!r} quantity.") + return getattr(second_order, name) + if name in self._THIRD_ORDER_NAMES: + third_order = object.__getattribute__(self, "third_order") + if third_order is None: + raise AttributeError(f"Lower-order solution has no {name!r} quantity.") + return getattr(third_order, name) + if name in self._SHEAR_NAMES: + shear = object.__getattribute__(self, "shear") + if shear is None: + raise AttributeError( + f"Magnetic shear has not been calculated; no {name!r} quantity." + ) + return getattr(shear, name) + if name in self._INVERSE_NAMES: + inverse = object.__getattribute__(self, "inverse") + if inverse is None: + raise AttributeError(f"Forward solution has no inverse diagnostic {name!r}.") + return getattr(inverse, name) + if name in self._MERCIER_NAMES: + mercier = object.__getattribute__(self, "mercier") + if mercier is None: + raise AttributeError("First-order solution has no Mercier diagnostics.") + return getattr(mercier, name) + if name in self._FIELD_JET_NAMES: + raise AttributeError("First-order solution has no second-derivative field jet.") + if name in self._SINGULARITY_NAMES: + raise AttributeError("First-order solution has no singular-radius diagnostics.") + raise AttributeError(f"{type(self).__name__!s} has no attribute {name!r}.") + + def _require_mercier(self) -> MercierDiagnostics: + mercier = object.__getattribute__(self, "mercier") + if mercier is None: + raise AttributeError("First-order solution has no Mercier diagnostics.") + return mercier + + def _require_field_jet(self) -> FieldJet: + field_jet = object.__getattribute__(self, "field_jet") + if field_jet is None: + raise AttributeError("First-order solution has no second-derivative field jet.") + return field_jet + + def _require_singularity(self) -> SingularityDiagnostics: + singularity = object.__getattribute__(self, "singularity") + if singularity is None: + raise AttributeError("First-order solution has no singular-radius diagnostics.") + return singularity + + @property + def axis(self) -> Axis: + return self.inputs.axis + + @property + def phi(self) -> jax.Array: + return self.geometry.samples.phi + + @property + def varphi(self) -> jax.Array: + return self.geometry.varphi + + @property + def R0(self) -> jax.Array: + return self.geometry.samples.R + + @property + def Z0(self) -> jax.Array: + return self.geometry.samples.Z + + @property + def curvature(self) -> jax.Array: + return self.geometry.curvature + + @property + def torsion(self) -> jax.Array: + return self.geometry.torsion + + @property + def axis_length(self) -> jax.Array: + return self.geometry.axis_length + + @property + def d2_volume_d_psi2(self) -> jax.Array: + return self._require_mercier().d2_volume_d_psi2 + + @property + def DGeod_times_r2(self) -> jax.Array: + return self._require_mercier().DGeod_times_r2 + + @property + def DWell_times_r2(self) -> jax.Array: + return self._require_mercier().DWell_times_r2 + + @property + def DMerc_times_r2(self) -> jax.Array: + return self._require_mercier().DMerc_times_r2 + + @property + def grad_grad_B_axis(self) -> jax.Array: + """Cartesian Hessian in ``(sample, field, derivative, derivative)`` order.""" + + return self._require_field_jet().hessian + + @property + def grad_grad_B(self) -> jax.Array: + """pyQSC-compatible Frenet Hessian in ``(sample, d, d, field)`` order.""" + + return self._require_field_jet().hessian_frenet + + @property + def L_grad_grad_B(self) -> jax.Array: + return self._require_field_jet().L_grad_grad_B + + @property + def grad_grad_B_inverse_scale_length_vs_varphi(self) -> jax.Array: + return self._require_field_jet().grad_grad_B_inverse_scale_length_vs_varphi + + @property + def grad_grad_B_inverse_scale_length(self) -> jax.Array: + return self._require_field_jet().grad_grad_B_inverse_scale_length + + @property + def r_singularity(self) -> jax.Array: + return self._require_singularity().r_singularity + + @property + def r_singularity_vs_varphi(self) -> jax.Array: + return self._require_singularity().r_singularity_vs_varphi + + @property + def inv_r_singularity_vs_varphi(self) -> jax.Array: + return self._require_singularity().inv_r_singularity_vs_varphi + + @property + def r_singularity_basic_vs_varphi(self) -> jax.Array: + return self._require_singularity().r_singularity_vs_varphi + + @property + def r_singularity_theta_vs_varphi(self) -> jax.Array: + return self._require_singularity().theta_singularity_vs_varphi + + @property + def r_singularity_residual_sqnorm(self) -> jax.Array: + return self._require_singularity().residual_norm_vs_varphi ** 2 diff --git a/src/pyqsc_jax/near_axis.py b/src/pyqsc_jax/near_axis.py new file mode 100644 index 0000000..f71a0a8 --- /dev/null +++ b/src/pyqsc_jax/near_axis.py @@ -0,0 +1,703 @@ +"""ESSOS-compatible mutable facade over the immutable near-axis core.""" + +from __future__ import annotations + +from typing import Any + +import jax +import jax.numpy as jnp + +from pyqsc_jax.axis import Axis +from pyqsc_jax.first_order import solve +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.shear import solve_magnetic_shear +from pyqsc_jax.solvers import implicit_dense_root +from pyqsc_jax.vmec import VmecExport +from pyqsc_jax.vmec import to_vmec as export_to_vmec + +ArrayLike = Any + + +class near_axis: # noqa: N801 + """Compatibility adapter for the historical ``near_axis`` interface. + + The canonical API is :func:`pyqsc_jax.solve`. This class intentionally + keeps ESSOS's mutable ``x``/``dofs`` facade and historical component-axis + ordering while delegating all near-axis physics to the immutable core. + """ + + def __init__( + self, + rc: ArrayLike = (1.0, 0.1), + zs: ArrayLike = (0.0, 0.1), + etabar: ArrayLike = 1.0, + B0: ArrayLike = 1.0, + sigma0: ArrayLike = 0.0, + I2: ArrayLike = 0.0, + nphi: int = 31, + spsi: int = 1, + sG: int = 1, + nfp: int = 2, + order: int | str = "r1", + B2c: ArrayLike = 0.0, + p2: ArrayLike = 0.0, + B2s: ArrayLike = 0.0, + ) -> None: + if not isinstance(nphi, int) or isinstance(nphi, bool) or nphi < 3 or nphi % 2 == 0: + raise ValueError("The compatibility API requires odd integer nphi >= 3.") + if isinstance(order, bool) or order not in (1, 2, 3, "r1", "r2", "r3"): + raise ValueError("order must be one of 1, 2, 3, 'r1', 'r2', or 'r3'.") + + self.rc = jnp.asarray(rc) + self.zs = jnp.asarray(zs) + if self.rc.ndim != 1 or self.zs.ndim != 1 or self.rc.size != self.zs.size: + raise ValueError("rc and zs must be one-dimensional arrays of equal length.") + self.etabar = jnp.asarray(etabar) + self.B0 = jnp.asarray(B0) + self.sigma0 = jnp.asarray(sigma0) + self.I2 = jnp.asarray(I2) + self.p2 = jnp.asarray(p2) + self.B2c = jnp.asarray(B2c) + self.B2s = jnp.asarray(B2s) + self.nphi = nphi + self.spsi = spsi + self.sG = sG + self.nfp = nfp + self.order = order + self.nfourier = self.rc.size + self._dofs = jnp.concatenate((self.rc, self.zs, self.etabar[None])) + self._refresh() + + def _canonical_solution( + self, + rc: ArrayLike, + zs: ArrayLike, + etabar: ArrayLike, + ) -> NearAxisSolution: + canonical_order = ( + "r1" if self.order in (1, "r1") else ("r2" if self.order in (2, "r2") else "r3") + ) + return solve( + axis=Axis.stellarator_symmetric(rc=rc, zs=zs, nfp=self.nfp), + etabar=etabar, + B0=self.B0, + sigma0=self.sigma0, + I2=self.I2, + p2=self.p2, + B2c=self.B2c, + B2s=self.B2s, + nphi=self.nphi, + order=canonical_order, + sG=self.sG, + spsi=self.spsi, + ) + + @staticmethod + def _legacy_tuple(solution: NearAxisSolution) -> tuple[jax.Array, ...]: + geometry = solution.geometry + normal = geometry.normal_cylindrical + binormal = geometry.binormal_cylindrical + return ( + solution.R0, + solution.Z0, + solution.sigma, + solution.elongation, + solution.B_axis.T, + jnp.moveaxis(solution.grad_B_axis, 0, -1), + solution.axis_length, + solution.iota, + solution.iotaN, + solution.G0, + solution.helicity, + solution.X1c_untwisted, + solution.X1s_untwisted, + solution.Y1s_untwisted, + solution.Y1c_untwisted, + normal[:, 0], + normal[:, 1], + normal[:, 2], + binormal[:, 0], + binormal[:, 1], + binormal[:, 2], + solution.L_grad_B, + 1 / solution.L_grad_B, + solution.torsion, + solution.curvature, + solution.varphi, + geometry.samples.d_R_d_phi, + geometry.samples.d_Z_d_phi, + ) + + def _refresh(self) -> None: + for name in NearAxisSolution._SHEAR_NAMES: + self.__dict__.pop(name, None) + solution = self._canonical_solution(self.rc, self.zs, self.etabar) + self.solution = solution + self.phi = solution.phi + ( + self.R0, + self.Z0, + self.sigma, + self.elongation, + self.B_axis, + self.grad_B_axis, + self.axis_length, + self.iota, + self.iotaN, + self.G0, + self.helicity, + self.X1c_untwisted, + self.X1s_untwisted, + self.Y1s_untwisted, + self.Y1c_untwisted, + self.normal_R, + self.normal_phi, + self.normal_z, + self.binormal_R, + self.binormal_phi, + self.binormal_z, + self.L_grad_B, + self.inv_L_grad_B, + self.torsion, + self.curvature, + self.varphi, + self.R0p, + self.Z0p, + ) = self._legacy_tuple(solution) + if solution.second_order is not None: + for name in solution._SECOND_ORDER_NAMES: + setattr(self, name, getattr(solution, name)) + self.d2_volume_d_psi2 = solution.d2_volume_d_psi2 + self.DGeod_times_r2 = solution.DGeod_times_r2 + self.DWell_times_r2 = solution.DWell_times_r2 + self.DMerc_times_r2 = solution.DMerc_times_r2 + self.grad_grad_B = solution.grad_grad_B + self.grad_grad_B_axis = jnp.moveaxis(solution.grad_grad_B_axis, 0, -1) + self.L_grad_grad_B = solution.L_grad_grad_B + self.grad_grad_B_inverse_scale_length_vs_varphi = ( + solution.grad_grad_B_inverse_scale_length_vs_varphi + ) + self.grad_grad_B_inverse_scale_length = solution.grad_grad_B_inverse_scale_length + self.r_singularity = solution.r_singularity + self.r_singularity_vs_varphi = solution.r_singularity_vs_varphi + self.inv_r_singularity_vs_varphi = solution.inv_r_singularity_vs_varphi + self.r_singularity_basic_vs_varphi = solution.r_singularity_basic_vs_varphi + self.r_singularity_theta_vs_varphi = solution.r_singularity_theta_vs_varphi + self.r_singularity_residual_sqnorm = solution.r_singularity_residual_sqnorm + if solution.third_order is not None: + for name in solution._THIRD_ORDER_NAMES: + setattr(self, name, getattr(solution, name)) + + @property + def dofs(self) -> jax.Array: + """Mutable legacy degrees of freedom ordered ``rc, zs, etabar``.""" + + return self._dofs + + @dofs.setter + def dofs(self, new_dofs: ArrayLike) -> None: + new_dofs = jnp.asarray(new_dofs) + if new_dofs.ndim != 1 or new_dofs.size != 2 * self.nfourier + 1: + raise ValueError(f"dofs must have shape ({2 * self.nfourier + 1},).") + self._dofs = new_dofs + self.rc = new_dofs[: self.nfourier] + self.zs = new_dofs[self.nfourier : 2 * self.nfourier] + self.etabar = new_dofs[-1] + self._refresh() + + @property + def x(self) -> jax.Array: + """Alias for :attr:`dofs`, retained for ESSOS optimizers.""" + + return self.dofs + + @x.setter + def x(self, new_x: ArrayLike) -> None: + self.dofs = new_x + + def _tree_flatten(self): + children = ( + self.rc, + self.zs, + self.etabar, + self.B0, + self.sigma0, + self.I2, + self.B2c, + self.p2, + self.B2s, + ) + auxiliary = { + "nphi": self.nphi, + "spsi": self.spsi, + "sG": self.sG, + "nfp": self.nfp, + "order": self.order, + } + return children, auxiliary + + @classmethod + def _tree_unflatten(cls, auxiliary, children): + rc, zs, etabar, B0, sigma0, I2, B2c, p2, B2s = children + return cls( + rc=rc, + zs=zs, + etabar=etabar, + B0=B0, + sigma0=sigma0, + I2=I2, + B2c=B2c, + p2=p2, + B2s=B2s, + **auxiliary, + ) + + def calculate(self, rc: ArrayLike, zs: ArrayLike, etabar: ArrayLike): + """Return the historical tuple, evaluated by the canonical core.""" + + return self._legacy_tuple(self._canonical_solution(rc, zs, etabar)) + + def calculate_shear(self, B31c: ArrayLike = 0.0) -> None: + """Populate the historical order-r-squared transform correction.""" + + self.solution = solve_magnetic_shear(self.solution, B31c=B31c) + for name in self.solution._SHEAR_NAMES: + setattr(self, name, getattr(self.solution, name)) + + def B_covariant(self, points: ArrayLike) -> jax.Array: + """First-order covariant Boozer components ``(B_r, B_theta, B_phi)``.""" + + r, _, _ = jnp.asarray(points) + return jnp.asarray((0.0, r * r * self.I2, self.G0)) + + def B_contravariant(self, points: ArrayLike) -> jax.Array: + """First-order contravariant Boozer components.""" + + r, _, _ = jnp.asarray(points) + Bphi = r * self.AbsB(points) / self.jacobian(points) + return jnp.asarray((0.0, self.iotaN * Bphi, Bphi)) + + def AbsB(self, points: ArrayLike) -> jax.Array: + """First-order field strength in near-axis coordinates.""" + + r, theta, _ = jnp.asarray(points) + return self.B0 * (1 + r * self.etabar * jnp.cos(theta)) + + def jacobian(self, points: ArrayLike) -> jax.Array: + """First-order coordinate Jacobian.""" + + r, _, _ = jnp.asarray(points) + field_strength = self.AbsB(points) + return r * self.B0 * (self.G0 + self.iota * self.I2) / field_strength**2 + + def interpolated_array_at_point(self, array: ArrayLike, point: ArrayLike) -> jax.Array: + """Periodically interpolate a sampled one-field-period array.""" + + period = 2 * jnp.pi / self.nfp + array = jnp.asarray(array) + return jnp.interp( + jnp.asarray(point), + jnp.append(self.phi, period), + jnp.append(array, array[0]), + period=period, + ) + + def Frenet_to_cylindrical_1_point( + self, + phi0: ArrayLike, + X_at_this_theta: ArrayLike, + Y_at_this_theta: ArrayLike, + Z_at_this_theta: ArrayLike | None = None, + ) -> tuple[jax.Array, jax.Array, jax.Array]: + """Map one displaced Frenet point to cylindrical coordinates.""" + + sine = jnp.sin(phi0) + cosine = jnp.cos(phi0) + R0 = self.interpolated_array_at_point(self.R0, phi0) + Z0 = self.interpolated_array_at_point(self.Z0, phi0) + X = self.interpolated_array_at_point(X_at_this_theta, phi0) + Y = self.interpolated_array_at_point(Y_at_this_theta, phi0) + if Z_at_this_theta is None: + Z_at_this_theta = jnp.zeros_like(X_at_this_theta) + Z = self.interpolated_array_at_point(Z_at_this_theta, phi0) + normal_R = self.interpolated_array_at_point(self.normal_R, phi0) + normal_phi = self.interpolated_array_at_point(self.normal_phi, phi0) + normal_z = self.interpolated_array_at_point(self.normal_z, phi0) + binormal_R = self.interpolated_array_at_point(self.binormal_R, phi0) + binormal_phi = self.interpolated_array_at_point(self.binormal_phi, phi0) + binormal_z = self.interpolated_array_at_point(self.binormal_z, phi0) + tangent = self.solution.geometry.tangent_cylindrical + tangent_R = self.interpolated_array_at_point(tangent[:, 0], phi0) + tangent_phi = self.interpolated_array_at_point(tangent[:, 1], phi0) + tangent_z = self.interpolated_array_at_point(tangent[:, 2], phi0) + + normal_x = normal_R * cosine - normal_phi * sine + normal_y = normal_R * sine + normal_phi * cosine + binormal_x = binormal_R * cosine - binormal_phi * sine + binormal_y = binormal_R * sine + binormal_phi * cosine + tangent_x = tangent_R * cosine - tangent_phi * sine + tangent_y = tangent_R * sine + tangent_phi * cosine + x = R0 * cosine + X * normal_x + Y * binormal_x + Z * tangent_x + y = R0 * sine + X * normal_y + Y * binormal_y + Z * tangent_y + z = Z0 + X * normal_z + Y * binormal_z + Z * tangent_z + return jnp.hypot(x, y), z, jnp.arctan2(y, x) + + def Frenet_to_cylindrical_residual_func( + self, + phi0: ArrayLike, + phi_target: ArrayLike, + X_at_this_theta: ArrayLike, + Y_at_this_theta: ArrayLike, + Z_at_this_theta: ArrayLike | None = None, + ) -> jax.Array: + """Wrapped cylindrical-angle residual for a Frenet point.""" + + _, _, phi = self.Frenet_to_cylindrical_1_point( + phi0, + X_at_this_theta, + Y_at_this_theta, + Z_at_this_theta, + ) + difference = phi - phi_target + return jnp.arctan2(jnp.sin(difference), jnp.cos(difference)) + + def residual_phi0_of_theta_varphi_func( + self, + phi0: ArrayLike, + r: ArrayLike, + theta: ArrayLike, + varphi: ArrayLike, + ) -> jax.Array: + """Residual for inversion at fixed Boozer toroidal angle.""" + + X, Y, Z = self._frenet_displacements(r, theta) + _, _, phi = self.Frenet_to_cylindrical_1_point(phi0, X, Y, Z) + nu0 = self.interpolated_array_at_point(self.varphi - self.phi, phi0) + X1c = self.interpolated_array_at_point(self.X1c_untwisted, phi0) + X1s = self.interpolated_array_at_point(self.X1s_untwisted, phi0) + Y1c = self.interpolated_array_at_point(self.Y1c_untwisted, phi0) + Y1s = self.interpolated_array_at_point(self.Y1s_untwisted, phi0) + bR = self.interpolated_array_at_point(self.binormal_R, phi0) + bZ = self.interpolated_array_at_point(self.binormal_z, phi0) + nR = self.interpolated_array_at_point(self.normal_R, phi0) + nZ = self.interpolated_array_at_point(self.normal_z, phi0) + R0 = self.interpolated_array_at_point(self.R0, phi0) + R0p = self.interpolated_array_at_point(self.R0p, phi0) + Z0p = self.interpolated_array_at_point(self.Z0p, phi0) + nu1c = X1c * (bR * Z0p - bZ * R0p) / R0 + Y1c * (nZ * R0p - nR * Z0p) / R0 + nu1s = X1s * (bR * Z0p - bZ * R0p) / R0 + Y1s * (nZ * R0p - nR * Z0p) / R0 + nu = nu0 + r * (nu1c * jnp.cos(theta) + nu1s * jnp.sin(theta)) + return phi + nu - varphi + + def _frenet_displacements( + self, + r: ArrayLike, + theta: ArrayLike, + ) -> tuple[jax.Array, jax.Array, jax.Array]: + """Assemble all available radial-order Frenet displacements.""" + + cosine = jnp.cos(theta) + sine = jnp.sin(theta) + X = r * (self.X1c_untwisted * cosine + self.X1s_untwisted * sine) + Y = r * (self.Y1c_untwisted * cosine + self.Y1s_untwisted * sine) + Z = jnp.zeros_like(X) + if self.solution.second_order is not None: + cosine2 = jnp.cos(2 * theta) + sine2 = jnp.sin(2 * theta) + X = X + r**2 * ( + self.X20_untwisted + self.X2c_untwisted * cosine2 + self.X2s_untwisted * sine2 + ) + Y = Y + r**2 * ( + self.Y20_untwisted + self.Y2c_untwisted * cosine2 + self.Y2s_untwisted * sine2 + ) + Z = Z + r**2 * ( + self.Z20_untwisted + self.Z2c_untwisted * cosine2 + self.Z2s_untwisted * sine2 + ) + if self.solution.third_order is not None: + cosine3 = jnp.cos(3 * theta) + sine3 = jnp.sin(3 * theta) + X = X + r**3 * ( + self.X3c1_untwisted * cosine + + self.X3s1_untwisted * sine + + self.X3c3_untwisted * cosine3 + + self.X3s3_untwisted * sine3 + ) + Y = Y + r**3 * ( + self.Y3c1_untwisted * cosine + + self.Y3s1_untwisted * sine + + self.Y3c3_untwisted * cosine3 + + self.Y3s3_untwisted * sine3 + ) + Z = Z + r**3 * ( + self.Z3c1_untwisted * cosine + + self.Z3s1_untwisted * sine + + self.Z3c3_untwisted * cosine3 + + self.Z3s3_untwisted * sine3 + ) + return X, Y, Z + + def phi_of_theta_varphi( + self, + r: ArrayLike, + theta: ArrayLike, + varphi: ArrayLike, + ) -> jax.Array: + """Invert the regular-coordinate map for cylindrical toroidal angle.""" + + residual = lambda phi0: self.residual_phi0_of_theta_varphi_func( # noqa: E731 + phi0, + r, + theta, + varphi, + ) + phi_on_axis, _ = implicit_dense_root(residual, jnp.asarray(varphi)) + X, Y, Z = self._frenet_displacements(r, theta) + _, _, phi = self.Frenet_to_cylindrical_1_point(phi_on_axis, X, Y, Z) + return phi + + def Frenet_to_cylindrical( + self, + r: ArrayLike, + ntheta: int = 20, + phi_is_varphi: bool = False, + ) -> tuple[jax.Array, jax.Array, jax.Array]: + """Map the available-order surface over one field period.""" + + theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False) + toroidal_grid = self.phi + + def for_theta(theta_value): + X, Y, Z = self._frenet_displacements(r, theta_value) + + def for_toroidal_angle(target): + if phi_is_varphi: + residual = lambda phi0: self.residual_phi0_of_theta_varphi_func( # noqa: E731 + phi0, + r, + theta_value, + target, + ) + else: + residual = lambda phi0: self.Frenet_to_cylindrical_residual_func( # noqa: E731 + phi0, + target, + X, + Y, + Z, + ) + phi0, _ = implicit_dense_root(residual, target) + R, cylindrical_Z, _ = self.Frenet_to_cylindrical_1_point(phi0, X, Y, Z) + return R, cylindrical_Z, phi0 + + return jax.vmap(for_toroidal_angle)(toroidal_grid) + + return jax.vmap(for_theta)(theta) + + def to_Fourier( # noqa: N802 + self, + R_2D: ArrayLike, + Z_2D: ArrayLike, + nfp: int, + mpol: int, + ntor: int, + ) -> tuple[jax.Array, jax.Array]: + """Convert a sampled stellarator-symmetric surface to Fourier data.""" + + R_2D = jnp.asarray(R_2D) + Z_2D = jnp.asarray(Z_2D) + ntheta, nphi = R_2D.shape + theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False) + phi = jnp.linspace(0, 2 * jnp.pi / nfp, nphi, endpoint=False) + phi2d, theta2d = jnp.meshgrid(phi, theta, indexing="xy") + m = jnp.arange(mpol + 1) + n = jnp.arange(-ntor, ntor + 1) + angle = ( + m[None, :, None, None] * theta2d[None, None, :, :] + - n[:, None, None, None] * nfp * phi2d[None, None, :, :] + ) + factor = 2 / (ntheta * nphi) + RBC = factor * jnp.sum(R_2D[None, None, :, :] * jnp.cos(angle), axis=(-2, -1)) + ZBS = factor * jnp.sum(Z_2D[None, None, :, :] * jnp.sin(angle), axis=(-2, -1)) + RBC = RBC.at[ntor, 0].set(jnp.mean(R_2D)) + RBC = RBC.at[:ntor, 0].set(0) + ZBS = ZBS.at[:ntor, 0].set(0) + return RBC, ZBS + + def get_boundary( + self, + r: ArrayLike = 0.1, + ntheta: int = 30, + nphi: int = 120, + ntheta_fourier: int = 20, + mpol: int = 5, + ntor: int = 5, + phi_is_varphi: bool = False, + phi_offset: ArrayLike = 0.0, + ) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + """Return a full-torus available-order surface in Cartesian coordinates.""" + + R_period, Z_period, _ = self.Frenet_to_cylindrical( + r, + ntheta=ntheta_fourier, + phi_is_varphi=phi_is_varphi, + ) + RBC, ZBS = self.to_Fourier(R_period, Z_period, self.nfp, mpol, ntor) + theta = jnp.linspace(0, 2 * jnp.pi, ntheta) + original_phi = jnp.linspace(0, 2 * jnp.pi, nphi) + phi_offset + phi2d, theta2d = jnp.meshgrid(original_phi, theta, indexing="xy") + + if phi_is_varphi: + phi2d = jax.vmap( + lambda theta_row, varphi_row: jax.vmap( + lambda theta_value, varphi_value: self.phi_of_theta_varphi( + r, + theta_value, + varphi_value, + ) + )(theta_row, varphi_row) + )(theta2d, phi2d) + + m = jnp.arange(mpol + 1) + n = jnp.arange(-ntor, ntor + 1) + angle = ( + m[None, :, None, None] * theta2d[None, None, :, :] + - n[:, None, None, None] * self.nfp * original_phi[None, None, None, :] + ) + R = jnp.sum(RBC[:, :, None, None] * jnp.cos(angle), axis=(0, 1)) + Z = jnp.sum(ZBS[:, :, None, None] * jnp.sin(angle), axis=(0, 1)) + return R * jnp.cos(phi2d), R * jnp.sin(phi2d), Z, R + + def to_vmec( + self, + filename, + r: float = 0.1, + params: dict[str, Any] | None = None, + ntheta: int = 40, + ntorMax: int = 14, # noqa: N803 + ) -> VmecExport: + """Write a fast, diagnosed VMEC input while preserving the pyQSC call form.""" + + parameters = dict(params or {}) + mpol = int(parameters.pop("mpol", min(ntheta // 2 - 1, 12))) + ntor = int(parameters.pop("ntor", min((self.nphi - 1) // 2, ntorMax))) + result = export_to_vmec( + self.solution, + filename, + r=r, + parameters=parameters, + ntheta=ntheta, + mpol=mpol, + ntor=ntor, + ntor_max=ntorMax, + ) + self.RBC = result.boundary.RBC.T + self.RBS = result.boundary.RBS.T + self.ZBC = result.boundary.ZBC.T + self.ZBS = result.boundary.ZBS.T + return result + + def B_mag(self, r: ArrayLike, theta: ArrayLike, phi: ArrayLike) -> jax.Array: + """Available-order field strength using the legacy angle convention.""" + + thetaN = theta - (self.iota - self.iotaN) * phi + field_strength = self.B0 * (1 + r * self.etabar * jnp.cos(thetaN)) + if self.solution.second_order is not None: + B20 = self.interpolated_array_at_point(self.B20, phi) + field_strength = field_strength + r**2 * ( + B20 + + self.B2c * jnp.cos(2 * thetaN) + + self.solution.inputs.B2s * jnp.sin(2 * thetaN) + ) + return field_strength + + def plot( + self, + r: float = 0.1, + ntheta: int = 40, + nphi: int = 120, + ntheta_fourier: int = 20, + ax=None, + show: bool = True, + close: bool = False, + axis_equal: bool = True, + **kwargs, + ): + """Plot the available-order boundary without importing ESSOS.""" + + import matplotlib.pyplot as plt + import numpy as np + from matplotlib import cm + from matplotlib.colors import LightSource, Normalize + + created_axes = ax is None or getattr(ax, "name", None) != "3d" + if created_axes: + figure = plt.figure() + ax = figure.add_subplot(projection="3d") + else: + figure = ax.figure + + x, y, z, _ = self.get_boundary( + r=r, + ntheta=ntheta, + nphi=nphi, + ntheta_fourier=ntheta_fourier, + ) + theta = jnp.linspace(0, 2 * jnp.pi, ntheta) + phi = jnp.linspace(0, 2 * jnp.pi, nphi) + phi2d, theta2d = jnp.meshgrid(phi, theta) + field_strength = np.asarray(self.B_mag(r, theta2d, phi2d)) + normalization = Normalize(vmin=field_strength.min(), vmax=field_strength.max()) + colormap = cm.viridis + facecolors = LightSource(azdeg=0, altdeg=10).shade( + field_strength, + colormap, + norm=normalization, + ) + kwargs.setdefault("alpha", 1) + ax.plot_surface( + np.asarray(x), + np.asarray(y), + np.asarray(z), + facecolors=facecolors, + rstride=1, + cstride=1, + antialiased=False, + linewidth=0, + shade=False, + **kwargs, + ) + if created_axes: + colorbar = figure.colorbar( + cm.ScalarMappable(cmap=colormap, norm=normalization), + ax=ax, + shrink=0.7, + ) + colorbar.ax.set_title(r"$|B|$ [T]") + ax.grid(False) + if axis_equal: + ranges = ( + np.ptp(np.asarray(x)), + np.ptp(np.asarray(y)), + np.ptp(np.asarray(z)), + ) + radius = max(ranges) / 2 + centers = ( + np.mean(np.asarray(x)), + np.mean(np.asarray(y)), + np.mean(np.asarray(z)), + ) + ax.set_xlim(centers[0] - radius, centers[0] + radius) + ax.set_ylim(centers[1] - radius, centers[1] + radius) + ax.set_zlim(centers[2] - radius, centers[2] + radius) + if show: + plt.show() + if close: + plt.close(figure) + return figure, ax + + +jax.tree_util.register_pytree_node( + near_axis, + near_axis._tree_flatten, + near_axis._tree_unflatten, +) diff --git a/src/pyqsc_jax/optimize.py b/src/pyqsc_jax/optimize.py new file mode 100644 index 0000000..90feaa0 --- /dev/null +++ b/src/pyqsc_jax/optimize.py @@ -0,0 +1,276 @@ +"""Dense B20 diagnostics and exact elimination of the affine B2c input.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +import jax +import jax.numpy as jnp + +from pyqsc_jax.first_order import solve +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.second_order import solve_second_order + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class B20Diagnostics: + """Dense-grid and toroidal-spectrum measures of nonconstant B20.""" + + weighted_mean: jax.Array + anomaly: jax.Array + weighted_l2: jax.Array + smooth_maximum: jax.Array + grid_maximum: jax.Array + peak_to_peak: jax.Array + fourier_modes: jax.Array + fourier_coefficients: jax.Array + nonzero_fourier_norm: jax.Array + nonzero_fourier_l1: jax.Array + fourier_tail_ratio: jax.Array + smooth_maximum_power: int + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class B2cOptimizationResult: + """Exact weighted-L2 optimum for the affine B2c dependence.""" + + solution: NearAxisSolution + diagnostics: B20Diagnostics + B2c_optimal: jax.Array + affine_intercept: jax.Array + affine_response: jax.Array + projected_response_norm_squared: jax.Array + affine_reconstruction_error: jax.Array + degenerate: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class B20ResolutionVerification: + """Independent fixed-parameter B20 checks on successively finer grids.""" + + resolutions: jax.Array + weighted_l2: jax.Array + smooth_maximum: jax.Array + grid_maximum: jax.Array + peak_to_peak: jax.Array + nonzero_fourier_l1: jax.Array + fourier_tail_ratio: jax.Array + relative_weighted_l2_change: jax.Array + relative_grid_maximum_change: jax.Array + + +def b20_diagnostics( + solution: NearAxisSolution, + *, + smooth_maximum_power: int = 16, +) -> B20Diagnostics: + """Evaluate normalized dense-grid and nonzero-mode B20 diagnostics.""" + + if solution.second_order is None: + raise ValueError("A second-order solution is required for B20 diagnostics.") + if ( + not isinstance(smooth_maximum_power, int) + or isinstance(smooth_maximum_power, bool) + or smooth_maximum_power < 2 + ): + raise ValueError("smooth_maximum_power must be an integer >= 2.") + weights = solution.geometry.d_l_d_phi + weight_sum = jnp.sum(weights) + B20 = solution.B20 + B0 = solution.inputs.B0 + weighted_mean = jnp.sum(weights * B20) / weight_sum + anomaly = B20 - weighted_mean + normalized = anomaly / B0 + weighted_l2 = jnp.sqrt(jnp.sum(weights * normalized**2) / weight_sum) + smooth_maximum = ( + jnp.sum(weights * jnp.abs(normalized) ** smooth_maximum_power) / weight_sum + ) ** (1 / smooth_maximum_power) + grid_maximum = jnp.max(jnp.abs(normalized)) + peak_to_peak = (jnp.max(B20) - jnp.min(B20)) / B0 + + maximum_mode = (solution.inputs.nphi - 1) // 2 + modes = jnp.arange(1, maximum_mode + 1) + phase = jnp.exp(-1j * modes[:, None] * solution.inputs.axis.nfp * solution.varphi[None, :]) + coefficients = jnp.sum(weights[None, :] * anomaly[None, :] * phase, axis=1) / weight_sum + coefficient_power = jnp.abs(coefficients / B0) ** 2 + nonzero_fourier_norm = jnp.sqrt(2 * jnp.sum(coefficient_power)) + nonzero_fourier_l1 = 2 * jnp.sum(jnp.abs(coefficients / B0)) + tail_start = maximum_mode * 3 // 4 + tiny = jnp.finfo(B20.dtype).tiny + fourier_tail_ratio = jnp.sqrt( + jnp.sum(coefficient_power[tail_start:]) / jnp.maximum(jnp.sum(coefficient_power), tiny) + ) + return B20Diagnostics( + weighted_mean=weighted_mean, + anomaly=anomaly, + weighted_l2=weighted_l2, + smooth_maximum=smooth_maximum, + grid_maximum=grid_maximum, + peak_to_peak=peak_to_peak, + fourier_modes=modes, + fourier_coefficients=coefficients, + nonzero_fourier_norm=nonzero_fourier_norm, + nonzero_fourier_l1=nonzero_fourier_l1, + fourier_tail_ratio=fourier_tail_ratio, + smooth_maximum_power=smooth_maximum_power, + ) + + +def _first_order_with_B2c(solution: NearAxisSolution, B2c: jax.Array) -> NearAxisSolution: + inputs = replace(solution.inputs, B2c=B2c) + return replace( + solution, + inputs=inputs, + second_order=None, + mercier=None, + field_jet=None, + singularity=None, + third_order=None, + shear=None, + ) + + +def _affine_B20(solution: NearAxisSolution) -> tuple[jax.Array, jax.Array]: + if solution.second_order is None: + raise ValueError("A second-order solution is required to eliminate B2c.") + zero = solve_second_order( + _first_order_with_B2c(solution, jnp.zeros_like(solution.inputs.B2c)), + attach_diagnostics=False, + ) + one = solve_second_order( + _first_order_with_B2c(solution, jnp.ones_like(solution.inputs.B2c)), + attach_diagnostics=False, + ) + return zero.B20, one.B20 - zero.B20 + + +def optimal_B2c_value( + solution: NearAxisSolution, + *, + degeneracy_tolerance: float = 1e-24, +) -> jax.Array: + """Return the exact B2c minimizing the nonconstant weighted-L2 B20.""" + + if degeneracy_tolerance < 0: + raise ValueError("degeneracy_tolerance must be nonnegative.") + intercept, response = _affine_B20(solution) + weights = solution.geometry.d_l_d_phi + weight_sum = jnp.sum(weights) + projected_intercept = intercept - jnp.sum(weights * intercept) / weight_sum + projected_response = response - jnp.sum(weights * response) / weight_sum + denominator = jnp.sum(weights * projected_response**2) + numerator = jnp.sum(weights * projected_intercept * projected_response) + return jnp.where( + denominator > degeneracy_tolerance, + -numerator / denominator, + solution.inputs.B2c, + ) + + +def optimize_B2c( + solution: NearAxisSolution, + *, + smooth_maximum_power: int = 16, + degeneracy_tolerance: float = 1e-24, +) -> B2cOptimizationResult: + """Eliminate B2c analytically and return the fully recomputed solution.""" + + intercept, response = _affine_B20(solution) + weights = solution.geometry.d_l_d_phi + weight_sum = jnp.sum(weights) + projected_intercept = intercept - jnp.sum(weights * intercept) / weight_sum + projected_response = response - jnp.sum(weights * response) / weight_sum + denominator = jnp.sum(weights * projected_response**2) + numerator = jnp.sum(weights * projected_intercept * projected_response) + degenerate = denominator <= degeneracy_tolerance + optimum = jnp.where(degenerate, solution.inputs.B2c, -numerator / denominator) + optimal = solve_second_order(_first_order_with_B2c(solution, optimum)) + if solution.inputs.order == 3: + from pyqsc_jax.third_order import solve_third_order + + optimal = solve_third_order(optimal) + if solution.shear is not None: + from pyqsc_jax.shear import solve_magnetic_shear + + optimal = solve_magnetic_shear(optimal, B31c=solution.shear.B31c) + reconstruction_error = jnp.max(jnp.abs(optimal.B20 - (intercept + optimum * response))) + return B2cOptimizationResult( + solution=optimal, + diagnostics=b20_diagnostics( + optimal, + smooth_maximum_power=smooth_maximum_power, + ), + B2c_optimal=optimum, + affine_intercept=intercept, + affine_response=response, + projected_response_norm_squared=denominator, + affine_reconstruction_error=reconstruction_error, + degenerate=degenerate, + ) + + +def verify_B20_resolution( + solution: NearAxisSolution, + *, + multipliers: tuple[int, ...] = (1, 2, 4), + smooth_maximum_power: int = 16, +) -> B20ResolutionVerification: + """Recompute one fixed physical candidate at independent resolutions.""" + + if not multipliers or any( + not isinstance(multiplier, int) or isinstance(multiplier, bool) or multiplier < 1 + for multiplier in multipliers + ): + raise ValueError("multipliers must be a nonempty tuple of positive integers.") + inputs = solution.inputs + resolutions = tuple(multiplier * (inputs.nphi - 1) + 1 for multiplier in multipliers) + diagnostics = [] + for nphi in resolutions: + candidate = solve( + axis=inputs.axis, + etabar=inputs.etabar, + B0=inputs.B0, + sigma0=inputs.sigma0, + I2=inputs.I2, + p2=inputs.p2, + B2c=inputs.B2c, + B2s=inputs.B2s, + nphi=nphi, + order="r2", + sG=inputs.sG, + spsi=inputs.spsi, + ) + diagnostics.append(b20_diagnostics(candidate, smooth_maximum_power=smooth_maximum_power)) + weighted_l2 = jnp.stack(tuple(item.weighted_l2 for item in diagnostics)) + smooth_maximum = jnp.stack(tuple(item.smooth_maximum for item in diagnostics)) + grid_maximum = jnp.stack(tuple(item.grid_maximum for item in diagnostics)) + peak_to_peak = jnp.stack(tuple(item.peak_to_peak for item in diagnostics)) + nonzero_fourier_l1 = jnp.stack(tuple(item.nonzero_fourier_l1 for item in diagnostics)) + tail_ratio = jnp.stack(tuple(item.fourier_tail_ratio for item in diagnostics)) + tiny = jnp.finfo(weighted_l2.dtype).tiny + relative_l2_change = jnp.concatenate( + ( + jnp.zeros(1, dtype=weighted_l2.dtype), + jnp.abs(jnp.diff(weighted_l2)) / jnp.maximum(weighted_l2[:-1], tiny), + ) + ) + relative_maximum_change = jnp.concatenate( + ( + jnp.zeros(1, dtype=grid_maximum.dtype), + jnp.abs(jnp.diff(grid_maximum)) / jnp.maximum(grid_maximum[:-1], tiny), + ) + ) + return B20ResolutionVerification( + resolutions=jnp.asarray(resolutions), + weighted_l2=weighted_l2, + smooth_maximum=smooth_maximum, + grid_maximum=grid_maximum, + peak_to_peak=peak_to_peak, + nonzero_fourier_l1=nonzero_fourier_l1, + fourier_tail_ratio=tail_ratio, + relative_weighted_l2_change=relative_l2_change, + relative_grid_maximum_change=relative_maximum_change, + ) diff --git a/src/pyqsc_jax/plasma.py b/src/pyqsc_jax/plasma.py new file mode 100644 index 0000000..ea79b13 --- /dev/null +++ b/src/pyqsc_jax/plasma.py @@ -0,0 +1,1106 @@ +"""Surface-free plasma-current source and free-space field jets.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import jax +import jax.numpy as jnp + +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.second_order import MU0 + +ArrayLike = Any + + +def _positive_scalar(value: ArrayLike, *, name: str) -> jax.Array: + array = jnp.asarray(value) + if array.ndim: + raise ValueError(f"{name} must be a scalar.") + try: + if bool(array <= 0): + raise ValueError(f"{name} must be positive.") + except jax.errors.TracerBoolConversionError: + pass + return array + + +def enclosed_current_from_covariant( + I2: ArrayLike, + *, + formal_radius: ArrayLike, + chi: int, +) -> jax.Array: + """Convert covariant ``I2`` to enclosed toroidal current in amperes.""" + + radius = _positive_scalar(formal_radius, name="formal_radius") + if chi not in (-1, 1): + raise ValueError("chi must be +1 or -1.") + return 2 * jnp.pi * chi * jnp.asarray(I2) * radius**2 / MU0 + + +def covariant_current_from_enclosed( + enclosed_current: ArrayLike, + *, + formal_radius: ArrayLike, + chi: int, +) -> jax.Array: + """Convert enclosed toroidal current in amperes to covariant ``I2``.""" + + radius = _positive_scalar(formal_radius, name="formal_radius") + if chi not in (-1, 1): + raise ValueError("chi must be +1 or -1.") + return MU0 * jnp.asarray(enclosed_current) / (2 * jnp.pi * chi * radius**2) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class PlasmaCurrentSource: + """Positive-volume current measure through quadratic radial order. + + The vector coefficients follow + + ``W / L = r * w1 + r**2 * (w2_cosine*cos(theta) + w2_sine*sin(theta))``. + """ + + formal_radius: jax.Array + parallel_current_mu0: jax.Array + enclosed_toroidal_current: jax.Array + C2: jax.Array + beta_1s: jax.Array + w1: jax.Array + wstar2_cosine: jax.Array + wstar2_sine: jax.Array + w2_cosine: jax.Array + w2_sine: jax.Array + axis_length_per_radian: jax.Array + chi: int = field(metadata={"static": True}) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class PlasmaFieldData: + """Matched on-axis free-space plasma field and error metadata.""" + + field: jax.Array + regularized_axis_integral: jax.Array + matched_axis_and_core: jax.Array + second_order_shape_correction: jax.Array + core_binormal_constant: jax.Array + core_normal_constant: jax.Array + maximum_matching_scale_error: jax.Array + formal_radius_to_curvature_radius: jax.Array + estimated_field_remainder: jax.Array + current_source: PlasmaCurrentSource + angular_resolution: int = field(metadata={"static": True}) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class PlasmaGradientData: + """Local plasma gradient and external vacuum value/gradient target.""" + + field: PlasmaFieldData + gradient: jax.Array + gradient_frenet: jax.Array + external_field: jax.Array + external_gradient: jax.Array + external_gradient_frenet: jax.Array + external_gradient_stf: jax.Array + external_gradient_independent: jax.Array + maximum_divergence: jax.Array + maximum_ampere_error: jax.Array + maximum_external_asymmetry: jax.Array + maximum_external_trace: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class PlasmaHessianData: + """Complete local plasma Hessian and external 3+5+7 vacuum target.""" + + field: PlasmaGradientData + hessian: jax.Array + hessian_frenet: jax.Array + external_hessian: jax.Array + external_hessian_stf: jax.Array + external_hessian_independent: jax.Array + maximum_derivative_asymmetry: jax.Array + maximum_external_symmetry_error: jax.Array + maximum_external_trace: jax.Array + estimated_hessian_remainder: jax.Array + + +def plasma_current_source( + solution: NearAxisSolution, + *, + formal_radius: ArrayLike, +) -> PlasmaCurrentSource: + """Construct the exact positive-volume weighted source through ``O(r²)``.""" + + if solution.second_order is None: + raise ValueError("The plasma current source requires a second-order solution.") + radius = _positive_scalar(formal_radius, name="formal_radius") + inputs = solution.inputs + geometry = solution.geometry + chi = inputs.sG * inputs.spsi + axis_length_per_radian = geometry.abs_G0_over_B0 + x = solution.X1c + y = solution.Y1s + y_sigma = solution.Y1c + derivative = geometry.d_d_varphi + current_density_mu0 = 2 * chi * inputs.I2 + C2 = solution.G2 + solution.N_helicity * inputs.I2 + + cosine_tangent = -chi * inputs.spsi * inputs.B0 * solution.beta_1s + cosine_normal = current_density_mu0 * ( + (derivative @ x) / axis_length_per_radian - geometry.torsion * y_sigma + ) + cosine_binormal = ( + current_density_mu0 + * ((derivative @ y_sigma) / axis_length_per_radian + geometry.torsion * x) + - 2 * chi * C2 * y / axis_length_per_radian + ) + sine_normal = ( + -current_density_mu0 * geometry.torsion * y + 2 * chi * C2 * x / axis_length_per_radian + ) + sine_binormal = ( + current_density_mu0 * (derivative @ y) / axis_length_per_radian + + 2 * chi * C2 * y_sigma / axis_length_per_radian + ) + + tangent = geometry.tangent_cartesian + normal = geometry.normal_cartesian + binormal = geometry.binormal_cartesian + w1 = current_density_mu0 * tangent + wstar2_cosine = ( + cosine_tangent * tangent + + cosine_normal[:, None] * normal + + cosine_binormal[:, None] * binormal + ) + wstar2_sine = sine_normal[:, None] * normal + sine_binormal[:, None] * binormal + w2_cosine = wstar2_cosine - (current_density_mu0 * geometry.curvature * x)[:, None] * tangent + return PlasmaCurrentSource( + formal_radius=radius, + parallel_current_mu0=current_density_mu0, + enclosed_toroidal_current=enclosed_current_from_covariant( + inputs.I2, + formal_radius=radius, + chi=chi, + ), + C2=C2, + beta_1s=solution.beta_1s, + w1=w1, + wstar2_cosine=wstar2_cosine, + wstar2_sine=wstar2_sine, + w2_cosine=w2_cosine, + w2_sine=wstar2_sine, + axis_length_per_radian=axis_length_per_radian, + chi=chi, + ) + + +def evaluate_weighted_current( + source: PlasmaCurrentSource, + radial_coordinate: ArrayLike, + theta: ArrayLike, +) -> jax.Array: + """Evaluate ``W = chi * mu0 * J_r * J`` at all toroidal samples.""" + + radial_coordinate, theta = jnp.broadcast_arrays( + jnp.asarray(radial_coordinate), + jnp.asarray(theta), + ) + radial = radial_coordinate[..., None, None] + cosine = jnp.cos(theta)[..., None, None] + sine = jnp.sin(theta)[..., None, None] + return source.axis_length_per_radian * ( + radial * source.w1 + radial**2 * (cosine * source.w2_cosine + sine * source.w2_sine) + ) + + +def _full_torus_axis_samples( + solution: NearAxisSolution, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + geometry = solution.geometry + nfp = solution.inputs.axis.nfp + cylindrical_period = 2 * jnp.pi / nfp + boozer_period = 2 * jnp.pi / nfp + period_index = jnp.arange(nfp) + full_phi = (geometry.samples.phi[None, :] + period_index[:, None] * cylindrical_period).reshape( + -1 + ) + full_varphi = (geometry.varphi[None, :] + period_index[:, None] * boozer_period).reshape(-1) + radius = jnp.tile(geometry.samples.R, nfp) + height = jnp.tile(geometry.samples.Z, nfp) + position = jnp.stack( + ( + radius * jnp.cos(full_phi), + radius * jnp.sin(full_phi), + height, + ), + axis=-1, + ) + tangent_cylindrical = jnp.tile( + geometry.tangent_cylindrical, + (nfp, 1), + ) + tangent = jnp.stack( + ( + tangent_cylindrical[:, 0] * jnp.cos(full_phi) + - tangent_cylindrical[:, 1] * jnp.sin(full_phi), + tangent_cylindrical[:, 0] * jnp.sin(full_phi) + + tangent_cylindrical[:, 1] * jnp.cos(full_phi), + tangent_cylindrical[:, 2], + ), + axis=-1, + ) + d_phi = cylindrical_period / solution.inputs.nphi + weights = jnp.tile(geometry.d_varphi_d_phi * d_phi, nfp) + return full_varphi, position, tangent, weights + + +def regularized_axis_integral(solution: NearAxisSolution) -> jax.Array: + """Evaluate the full-torus periodic finite-part integral in Eq. (136).""" + + source_varphi, source_position, source_tangent, weights = _full_torus_axis_samples(solution) + observation_varphi = solution.varphi + observation_position = solution.geometry.position_cartesian + displacement = observation_position[:, None, :] - source_position[None, :, :] + distance_squared = jnp.sum(displacement**2, axis=-1) + angle = source_varphi[None, :] - observation_varphi[:, None] + sine_half = jnp.sin(0.5 * angle) + coincident = jnp.abs(sine_half) < 16 * jnp.finfo(angle.dtype).eps + safe_distance_squared = jnp.where(coincident, 1.0, distance_squared) + safe_sine = jnp.where(coincident, 1.0, jnp.abs(sine_half)) + filament = ( + solution.geometry.abs_G0_over_B0 + * jnp.cross( + source_tangent[None, :, :], + displacement, + ) + / safe_distance_squared[..., None] ** 1.5 + ) + singular_model = ( + solution.geometry.curvature[:, None, None] + * solution.geometry.binormal_cartesian[:, None, :] + / (4 * safe_sine[..., None]) + ) + integrand = jnp.where( + coincident[..., None], + 0.0, + filament - singular_model, + ) + return jnp.sum(weights[None, :, None] * integrand, axis=1) + + +def _second_order_shape_correction( + solution: NearAxisSolution, + source: PlasmaCurrentSource, + *, + angular_resolution: int, +) -> jax.Array: + if ( + not isinstance(angular_resolution, int) + or isinstance(angular_resolution, bool) + or angular_resolution < 8 + ): + raise ValueError("angular_resolution must be an integer >= 8.") + theta = 2 * jnp.pi * jnp.arange(angular_resolution) / angular_resolution + cosine = jnp.cos(theta)[None, :, None] + sine = jnp.sin(theta)[None, :, None] + cosine2 = jnp.cos(2 * theta)[None, :, None] + sine2 = jnp.sin(2 * theta)[None, :, None] + geometry = solution.geometry + tangent = geometry.tangent_cartesian[:, None, :] + normal = geometry.normal_cartesian[:, None, :] + binormal = geometry.binormal_cartesian[:, None, :] + e = ( + solution.X1c[:, None, None] * cosine * normal + + (solution.Y1s[:, None, None] * sine + solution.Y1c[:, None, None] * cosine) * binormal + ) + X2 = ( + solution.X20[:, None, None] + + solution.X2c[:, None, None] * cosine2 + + solution.X2s[:, None, None] * sine2 + ) + Y2 = ( + solution.Y20[:, None, None] + + solution.Y2c[:, None, None] * cosine2 + + solution.Y2s[:, None, None] * sine2 + ) + Z2 = ( + solution.Z20[:, None, None] + + solution.Z2c[:, None, None] * cosine2 + + solution.Z2s[:, None, None] * sine2 + ) + xi2 = X2 * normal + Y2 * binormal + Z2 * tangent + w1 = source.w1[:, None, :] + wstar2 = source.wstar2_cosine[:, None, :] * cosine + source.wstar2_sine[:, None, :] * sine + E2 = jnp.sum(e**2, axis=-1, keepdims=True) + integrand = (jnp.cross(w1, xi2) + jnp.cross(wstar2, e)) / E2 - 2 * jnp.sum( + e * xi2, axis=-1, keepdims=True + ) * jnp.cross(w1, e) / E2**2 + return -0.5 * source.formal_radius**2 * jnp.mean(integrand, axis=1) + + +def matched_plasma_field_kernel( + solution: NearAxisSolution, + source: PlasmaCurrentSource, + regularized_integral: jax.Array, + *, + reference_length: ArrayLike, +) -> jax.Array: + """Combine finite-part and local-core terms at an arbitrary matching length.""" + + reference_length = _positive_scalar( + reference_length, + name="reference_length", + ) + geometry = solution.geometry + axis_scale = geometry.abs_G0_over_B0 + x = solution.X1c + sigma = solution.Y1c / solution.Y1s + trace_Q = x**2 + (1 + sigma**2) / x**2 + core_binormal = -0.5 - 0.5 * jnp.log((trace_Q + 2) / 4) + (x**2 + 1) / (trace_Q + 2) + core_normal = -source.chi * sigma / (trace_Q + 2) + curvature_binormal = geometry.curvature[:, None] * geometry.binormal_cartesian + finite_part = regularized_integral - curvature_binormal * jnp.log( + reference_length / (4 * axis_scale) + ) + return ( + finite_part + + curvature_binormal + * (jnp.log(2 * reference_length / source.formal_radius) + core_binormal[:, None]) + + (geometry.curvature * core_normal)[:, None] * geometry.normal_cartesian + ) + + +def plasma_field_on_axis( + solution: NearAxisSolution, + *, + formal_radius: ArrayLike, + angular_resolution: int = 128, +) -> PlasmaFieldData: + """Evaluate the matched leading on-axis free-space plasma field.""" + + source = plasma_current_source( + solution, + formal_radius=formal_radius, + ) + regularized_integral = regularized_axis_integral(solution) + axis_scale = solution.geometry.abs_G0_over_B0 + matched = matched_plasma_field_kernel( + solution, + source, + regularized_integral, + reference_length=4 * axis_scale, + ) + independent_matching_scale = matched_plasma_field_kernel( + solution, + source, + regularized_integral, + reference_length=7 * axis_scale, + ) + shape_correction = _second_order_shape_correction( + solution, + source, + angular_resolution=angular_resolution, + ) + current_prefactor = source.parallel_current_mu0 * source.formal_radius**2 / 4 + field_value = current_prefactor * matched + shape_correction + x = solution.X1c + sigma = solution.Y1c / solution.Y1s + trace_Q = x**2 + (1 + sigma**2) / x**2 + core_binormal = -0.5 - 0.5 * jnp.log((trace_Q + 2) / 4) + (x**2 + 1) / (trace_Q + 2) + core_normal = -source.chi * sigma / (trace_Q + 2) + radius_to_curvature = source.formal_radius * jnp.max(solution.geometry.curvature) + logarithm = jnp.abs(jnp.log(source.formal_radius / solution.geometry.abs_G0_over_B0)) + estimated_remainder = ( + jnp.abs(source.parallel_current_mu0) + * source.formal_radius**4 + / solution.geometry.abs_G0_over_B0**3 + * (1 + logarithm) + ) + return PlasmaFieldData( + field=field_value, + regularized_axis_integral=regularized_integral, + matched_axis_and_core=matched, + second_order_shape_correction=shape_correction, + core_binormal_constant=core_binormal, + core_normal_constant=core_normal, + maximum_matching_scale_error=jnp.max(jnp.abs(matched - independent_matching_scale)), + formal_radius_to_curvature_radius=radius_to_curvature, + estimated_field_remainder=estimated_remainder, + current_source=source, + angular_resolution=angular_resolution, + ) + + +def project_symmetric_trace_free_rank2(tensor: ArrayLike) -> jax.Array: + """Project final two axes onto symmetric trace-free rank-two tensors.""" + + tensor = jnp.asarray(tensor) + if tensor.shape[-2:] != (3, 3): + raise ValueError("A rank-two Cartesian tensor must end in shape (3, 3).") + symmetric = 0.5 * (tensor + jnp.swapaxes(tensor, -1, -2)) + trace = jnp.trace(symmetric, axis1=-2, axis2=-1) + return ( + symmetric + - trace[..., None, None] + * jnp.eye( + 3, + dtype=tensor.dtype, + ) + / 3 + ) + + +def pack_symmetric_trace_free_rank2(tensor: ArrayLike) -> jax.Array: + """Pack an STF matrix as ``(xx, yy, xy, xz, yz)``.""" + + tensor = project_symmetric_trace_free_rank2(tensor) + return jnp.stack( + ( + tensor[..., 0, 0], + tensor[..., 1, 1], + tensor[..., 0, 1], + tensor[..., 0, 2], + tensor[..., 1, 2], + ), + axis=-1, + ) + + +def unpack_symmetric_trace_free_rank2(components: ArrayLike) -> jax.Array: + """Unpack ``(xx, yy, xy, xz, yz)`` into an STF matrix.""" + + components = jnp.asarray(components) + if components.shape[-1:] != (5,): + raise ValueError("STF rank-two components must end in length 5.") + xx, yy, xy, xz, yz = jnp.moveaxis(components, -1, 0) + return jnp.stack( + ( + jnp.stack((xx, xy, xz), axis=-1), + jnp.stack((xy, yy, yz), axis=-1), + jnp.stack((xz, yz, -xx - yy), axis=-1), + ), + axis=-2, + ) + + +def elliptical_channel_gradient( + x: ArrayLike, + sigma: ArrayLike, + *, + parallel_current_mu0: ArrayLike, + chi: int, + frame: ArrayLike | None = None, +) -> jax.Array: + """Return the leading field-component-first gradient of an elliptical channel.""" + + if chi not in (-1, 1): + raise ValueError("chi must be +1 or -1.") + x = jnp.asarray(x) + sigma = jnp.asarray(sigma) + parallel_current_mu0 = jnp.asarray(parallel_current_mu0) + trace_Q = x**2 + (1 + sigma**2) / x**2 + derivative_first = ( + parallel_current_mu0[..., None, None] + / (trace_Q + 2)[..., None, None] + * jnp.stack( + ( + jnp.stack( + ( + jnp.zeros_like(x), + jnp.zeros_like(x), + jnp.zeros_like(x), + ), + axis=-1, + ), + jnp.stack( + ( + jnp.zeros_like(x), + chi * sigma, + 1 + (1 + sigma**2) / x**2, + ), + axis=-1, + ), + jnp.stack( + ( + jnp.zeros_like(x), + -(1 + x**2), + -chi * sigma, + ), + axis=-1, + ), + ), + axis=-2, + ) + ) + field_first_frenet = jnp.swapaxes(derivative_first, -1, -2) + if frame is None: + return field_first_frenet + frame = jnp.asarray(frame) + if frame.shape[-2:] != (3, 3): + raise ValueError("frame must end in shape (3, 3).") + return jnp.einsum( + "...ai,...ab,...bj->...ij", + frame, + field_first_frenet, + frame, + ) + + +def plasma_gradient_on_axis( + solution: NearAxisSolution, + *, + formal_radius: ArrayLike, + angular_resolution: int = 128, +) -> PlasmaGradientData: + """Evaluate the local plasma gradient and subtract it from the total jet.""" + + plasma_field = plasma_field_on_axis( + solution, + formal_radius=formal_radius, + angular_resolution=angular_resolution, + ) + x = solution.X1c + sigma = solution.Y1c / solution.Y1s + frame = jnp.stack( + ( + solution.geometry.tangent_cartesian, + solution.geometry.normal_cartesian, + solution.geometry.binormal_cartesian, + ), + axis=-2, + ) + gradient_frenet = elliptical_channel_gradient( + x, + sigma, + parallel_current_mu0=plasma_field.current_source.parallel_current_mu0, + chi=plasma_field.current_source.chi, + ) + gradient = elliptical_channel_gradient( + x, + sigma, + parallel_current_mu0=plasma_field.current_source.parallel_current_mu0, + chi=plasma_field.current_source.chi, + frame=frame, + ) + external_field = solution.B_axis - plasma_field.field + external_gradient = solution.grad_B_axis - gradient + external_gradient_frenet = jnp.einsum( + "...ai,...ij,...bj->...ab", + frame, + external_gradient, + frame, + ) + external_gradient_stf = project_symmetric_trace_free_rank2( + external_gradient, + ) + plasma_divergence = jnp.trace(gradient, axis1=-2, axis2=-1) + ampere = ( + gradient_frenet[:, 2, 1] + - gradient_frenet[:, 1, 2] + - plasma_field.current_source.parallel_current_mu0 + ) + external_asymmetry = external_gradient - jnp.swapaxes( + external_gradient, + -1, + -2, + ) + external_trace = jnp.trace( + external_gradient, + axis1=-2, + axis2=-1, + ) + return PlasmaGradientData( + field=plasma_field, + gradient=gradient, + gradient_frenet=gradient_frenet, + external_field=external_field, + external_gradient=external_gradient, + external_gradient_frenet=external_gradient_frenet, + external_gradient_stf=external_gradient_stf, + external_gradient_independent=pack_symmetric_trace_free_rank2( + external_gradient, + ), + maximum_divergence=jnp.max(jnp.abs(plasma_divergence)), + maximum_ampere_error=jnp.max(jnp.abs(ampere)), + maximum_external_asymmetry=jnp.max(jnp.abs(external_asymmetry)), + maximum_external_trace=jnp.max(jnp.abs(external_trace)), + ) + + +def _symmetrize_rank3(tensor: ArrayLike) -> jax.Array: + """Symmetrize the final three axes of a Cartesian tensor.""" + + tensor = jnp.asarray(tensor) + if tensor.shape[-3:] != (3, 3, 3): + raise ValueError("A rank-three Cartesian tensor must end in shape (3, 3, 3).") + return ( + tensor + + jnp.transpose(tensor, (*range(tensor.ndim - 3), -3, -1, -2)) + + jnp.transpose(tensor, (*range(tensor.ndim - 3), -2, -3, -1)) + + jnp.transpose(tensor, (*range(tensor.ndim - 3), -2, -1, -3)) + + jnp.transpose(tensor, (*range(tensor.ndim - 3), -1, -3, -2)) + + jnp.transpose(tensor, (*range(tensor.ndim - 3), -1, -2, -3)) + ) / 6 + + +def project_symmetric_trace_free_rank3(tensor: ArrayLike) -> jax.Array: + """Project final three axes onto fully symmetric trace-free rank three.""" + + tensor = jnp.asarray(tensor) + symmetric = _symmetrize_rank3(tensor) + trace = jnp.einsum("...iik->...k", symmetric) + identity = jnp.eye(3, dtype=tensor.dtype) + correction = ( + jnp.einsum("ij,...k->...ijk", identity, trace) + + jnp.einsum("ik,...j->...ijk", identity, trace) + + jnp.einsum("jk,...i->...ijk", identity, trace) + ) / 5 + return symmetric - correction + + +def pack_symmetric_trace_free_rank3(tensor: ArrayLike) -> jax.Array: + """Pack an STF rank-three tensor as seven Cartesian components.""" + + tensor = project_symmetric_trace_free_rank3(tensor) + return jnp.stack( + ( + tensor[..., 0, 0, 0], + tensor[..., 0, 0, 1], + tensor[..., 0, 0, 2], + tensor[..., 0, 1, 1], + tensor[..., 0, 1, 2], + tensor[..., 1, 1, 1], + tensor[..., 1, 1, 2], + ), + axis=-1, + ) + + +def _rank3_stf_basis(dtype: jnp.dtype) -> jax.Array: + basis = jnp.zeros((7, 3, 3, 3), dtype=dtype) + + def set_symmetric(array, basis_index, indices, value): + from itertools import permutations + + for permutation in set(permutations(indices)): + array = array.at[ + basis_index, + permutation[0], + permutation[1], + permutation[2], + ].set(value) + return array + + basis = set_symmetric(basis, 0, (0, 0, 0), 1) + basis = set_symmetric(basis, 0, (0, 2, 2), -1) + basis = set_symmetric(basis, 1, (0, 0, 1), 1) + basis = set_symmetric(basis, 1, (1, 2, 2), -1) + basis = set_symmetric(basis, 2, (0, 0, 2), 1) + basis = set_symmetric(basis, 2, (2, 2, 2), -1) + basis = set_symmetric(basis, 3, (0, 1, 1), 1) + basis = set_symmetric(basis, 3, (0, 2, 2), -1) + basis = set_symmetric(basis, 4, (0, 1, 2), 1) + basis = set_symmetric(basis, 5, (1, 1, 1), 1) + basis = set_symmetric(basis, 5, (1, 2, 2), -1) + basis = set_symmetric(basis, 6, (1, 1, 2), 1) + basis = set_symmetric(basis, 6, (2, 2, 2), -1) + return basis + + +def unpack_symmetric_trace_free_rank3(components: ArrayLike) -> jax.Array: + """Unpack ``(xxx,xxy,xxz,xyy,xyz,yyy,yyz)`` into an STF tensor.""" + + components = jnp.asarray(components) + if components.shape[-1:] != (7,): + raise ValueError("STF rank-three components must end in length 7.") + return jnp.einsum( + "...a,aijk->...ijk", + components, + _rank3_stf_basis(components.dtype), + ) + + +def _affine_potential_third_derivative( + coefficients: jax.Array, + elongation: jax.Array, +) -> jax.Array: + """Equation (204)-(205) for scalar or vector affine coefficients.""" + + inverse = 1 / elongation + denominator = (elongation + inverse) ** 2 + alpha_plus = coefficients[..., 0] + alpha_minus = coefficients[..., 1] + prefix_shape = coefficients.shape[:-1] + result = jnp.zeros( + (*prefix_shape, 2, 2, 2), + dtype=coefficients.dtype, + ) + broadcast_shape = (elongation.shape[0],) + (1,) * (alpha_plus.ndim - 1) + elongation = elongation.reshape(broadcast_shape) + inverse = inverse.reshape(broadcast_shape) + denominator = denominator.reshape(broadcast_shape) + plus_plus_plus = 2 * jnp.pi / denominator * (2 + inverse**2) * alpha_plus + plus_plus_minus = 2 * jnp.pi / denominator * inverse**2 * alpha_minus + plus_minus_minus = 2 * jnp.pi / denominator * elongation**2 * alpha_plus + minus_minus_minus = 2 * jnp.pi / denominator * (elongation**2 + 2) * alpha_minus + result = result.at[..., 0, 0, 0].set(plus_plus_plus) + result = result.at[..., 1, 1, 1].set(minus_minus_minus) + for indices in ((0, 0, 1), (0, 1, 0), (1, 0, 0)): + result = result.at[..., indices[0], indices[1], indices[2]].set(plus_plus_minus) + for indices in ((0, 1, 1), (1, 0, 1), (1, 1, 0)): + result = result.at[..., indices[0], indices[1], indices[2]].set(plus_minus_minus) + return result + + +def _oriented_ellipse_svd( + solution: NearAxisSolution, +) -> tuple[jax.Array, jax.Array, jax.Array]: + x = solution.X1c + y = solution.Y1s + sigma = solution.Y1c / y + ellipse = jnp.stack( + ( + jnp.stack((x, jnp.zeros_like(x)), axis=-1), + jnp.stack((y * sigma, y), axis=-1), + ), + axis=-2, + ) + physical_rotation, singular_values, parameter_rotation_transpose = jnp.linalg.svd( + ellipse, full_matrices=False + ) + orientation = jnp.linalg.det(physical_rotation) + column_correction = jnp.stack( + (jnp.ones_like(orientation), orientation), + axis=-1, + ) + physical_rotation = physical_rotation * column_correction[:, None, :] + parameter_rotation_transpose = parameter_rotation_transpose * column_correction[:, :, None] + return ( + physical_rotation, + singular_values[:, 0], + jnp.swapaxes(parameter_rotation_transpose, -1, -2), + ) + + +def _principal_transverse_plasma_hessian( + solution: NearAxisSolution, + source: PlasmaCurrentSource, +) -> tuple[jax.Array, jax.Array]: + """Equations (161)-(214) in the oriented ellipse principal frame.""" + + physical_rotation, elongation, parameter_rotation = _oriented_ellipse_svd(solution) + inverse = 1 / elongation + d_lambda = elongation + inverse + tangent = solution.geometry.tangent_cartesian + normal = solution.geometry.normal_cartesian + binormal = solution.geometry.binormal_cartesian + e_plus = physical_rotation[:, 0, 0, None] * normal + physical_rotation[:, 1, 0, None] * binormal + e_minus = ( + physical_rotation[:, 0, 1, None] * normal + physical_rotation[:, 1, 1, None] * binormal + ) + principal_frame = jnp.stack((tangent, e_plus, e_minus), axis=-2) + current_coefficients = jnp.stack( + (source.wstar2_cosine, source.wstar2_sine), + axis=-1, + ) + current_principal = jnp.einsum( + "...gi,...ij->...gj", + jnp.einsum( + "...gi,...ij->...gj", + principal_frame, + current_coefficients, + ), + parameter_rotation, + ) + c_star = jnp.stack( + ( + current_principal[..., 0] / elongation[:, None], + current_principal[..., 1] * elongation[:, None], + ), + axis=-1, + ) + + H_normal = jnp.stack( + ( + jnp.stack( + (solution.X20 + solution.X2c, solution.X2s), + axis=-1, + ), + jnp.stack( + (solution.X2s, solution.X20 - solution.X2c), + axis=-1, + ), + ), + axis=-2, + ) + H_binormal = jnp.stack( + ( + jnp.stack( + (solution.Y20 + solution.Y2c, solution.Y2s), + axis=-1, + ), + jnp.stack( + (solution.Y2s, solution.Y20 - solution.Y2c), + axis=-1, + ), + ), + axis=-2, + ) + transformed_normal = jnp.einsum( + "...ia,...ij,...jb->...ab", + parameter_rotation, + H_normal, + parameter_rotation, + ) + transformed_binormal = jnp.einsum( + "...ia,...ij,...jb->...ab", + parameter_rotation, + H_binormal, + parameter_rotation, + ) + G_plus = ( + physical_rotation[:, 0, 0, None, None] * transformed_normal + + physical_rotation[:, 1, 0, None, None] * transformed_binormal + ) + G_minus = ( + physical_rotation[:, 0, 1, None, None] * transformed_normal + + physical_rotation[:, 1, 1, None, None] * transformed_binormal + ) + d1 = 2 * G_plus[:, 0, 0] / elongation + 2 * elongation * G_minus[:, 0, 1] + d2 = 2 * G_plus[:, 0, 1] / elongation + 2 * elongation * G_minus[:, 1, 1] + delta = jnp.stack((d1 / elongation, elongation * d2), axis=-1) + p3_cosine = (G_plus[:, 0, 0] - G_plus[:, 1, 1]) / (4 * elongation) - elongation * G_minus[ + :, 0, 1 + ] / 2 + p3_sine = ( + G_plus[:, 0, 1] / (2 * elongation) + elongation * (G_minus[:, 0, 0] - G_minus[:, 1, 1]) / 4 + ) + ellipse_parameter = (elongation - inverse) / d_lambda + ell_cosine = ( + -4 + * jnp.pi + * (1 + ellipse_parameter**3) + * p3_cosine + / (3 * elongation * (elongation**2 + 3 * inverse**2)) + ) + ell_sine = ( + -4 + * jnp.pi + * (1 - ellipse_parameter**3) + * p3_sine + / (3 * inverse * (3 * elongation**2 + inverse**2)) + ) + boundary_third = jnp.zeros( + (solution.inputs.nphi, 2, 2, 2), + dtype=elongation.dtype, + ) + boundary_third = boundary_third.at[:, 0, 0, 0].set(6 * ell_cosine) + boundary_third = boundary_third.at[:, 1, 1, 1].set(-6 * ell_sine) + for indices in ((0, 0, 1), (0, 1, 0), (1, 0, 0)): + boundary_third = boundary_third.at[ + :, + indices[0], + indices[1], + indices[2], + ].set(6 * ell_sine) + for indices in ((0, 1, 1), (1, 0, 1), (1, 1, 0)): + boundary_third = boundary_third.at[ + :, + indices[0], + indices[1], + indices[2], + ].set(-6 * ell_cosine) + + n_components = physical_rotation[:, 0, :] + K = jnp.zeros((solution.inputs.nphi, 2, 2), dtype=elongation.dtype) + K = K.at[:, 0, 0].set(2 * jnp.pi * inverse / d_lambda) + K = K.at[:, 1, 1].set(2 * jnp.pi * elongation / d_lambda) + M = -source.parallel_current_mu0 * K / (2 * jnp.pi) + product_rule = ( + jnp.einsum("...a,...bc->...abc", n_components, K) + + jnp.einsum("...b,...ac->...abc", n_components, K) + + jnp.einsum("...c,...ab->...abc", n_components, K) + ) + potential_third = -_affine_potential_third_derivative(c_star, elongation) / (2 * jnp.pi) + tangent_contribution = ( + source.parallel_current_mu0 * solution.geometry.curvature / (4 * jnp.pi) + )[:, None, None, None] * ( + _affine_potential_third_derivative( + n_components, + elongation, + ) + - product_rule + ) + source.parallel_current_mu0 / (2 * jnp.pi) * ( + _affine_potential_third_derivative(delta, elongation) - boundary_third + ) + potential_third = potential_third.at[:, 0, :, :, :].add(tangent_contribution) + + transverse = jnp.zeros( + (solution.inputs.nphi, 2, 2, 3), + dtype=elongation.dtype, + ) + transverse = transverse.at[:, :, :, 0].set( + potential_third[:, 2, :, :, 0] - potential_third[:, 1, :, :, 1] + ) + transverse = transverse.at[:, :, :, 1].set( + potential_third[:, 0, :, :, 1] + - (solution.geometry.curvature * n_components[:, 1])[:, None, None] * M + ) + transverse = transverse.at[:, :, :, 2].set( + (solution.geometry.curvature * n_components[:, 0])[:, None, None] * M + - potential_third[:, 0, :, :, 0] + ) + return transverse, principal_frame + + +def _plasma_hessian_frenet( + solution: NearAxisSolution, + gradient_frenet: jax.Array, + source: PlasmaCurrentSource, +) -> jax.Array: + transverse, principal_frame = _principal_transverse_plasma_hessian( + solution, + source, + ) + frenet_frame = jnp.stack( + ( + solution.geometry.tangent_cartesian, + solution.geometry.normal_cartesian, + solution.geometry.binormal_cartesian, + ), + axis=-2, + ) + principal_from_frenet = jnp.einsum( + "...ai,...bi->...ab", + principal_frame, + frenet_frame, + ) + principal_derivative_first = jnp.zeros( + (solution.inputs.nphi, 3, 3, 3), + dtype=transverse.dtype, + ) + principal_derivative_first = principal_derivative_first.at[ + :, + 1:, + 1:, + :, + ].set(transverse) + frenet_derivative_first = jnp.einsum( + "...pa,...qb,...rg,...pqr->...abg", + principal_from_frenet, + principal_from_frenet, + principal_from_frenet, + principal_derivative_first, + ) + + derivative_first_gradient = jnp.swapaxes(gradient_frenet, -1, -2) + derivative_along_axis = ( + jnp.einsum( + "nm,mab->nab", + solution.geometry.d_d_varphi, + derivative_first_gradient, + ) + / solution.geometry.abs_G0_over_B0 + ) + zeros = jnp.zeros_like(solution.geometry.curvature) + connection = jnp.stack( + ( + jnp.stack( + (zeros, solution.geometry.curvature, zeros), + axis=-1, + ), + jnp.stack( + ( + -solution.geometry.curvature, + zeros, + solution.geometry.torsion, + ), + axis=-1, + ), + jnp.stack( + (zeros, -solution.geometry.torsion, zeros), + axis=-1, + ), + ), + axis=-2, + ) + tangential = ( + derivative_along_axis + - jnp.einsum( + "...ad,...dg->...ag", + connection, + derivative_first_gradient, + ) + - jnp.einsum( + "...ad,...gd->...ag", + derivative_first_gradient, + connection, + ) + ) + frenet_derivative_first = frenet_derivative_first.at[:, 0, :, :].set(tangential) + frenet_derivative_first = frenet_derivative_first.at[:, :, 0, :].set(tangential) + return jnp.transpose(frenet_derivative_first, (0, 3, 1, 2)) + + +def plasma_hessian_on_axis( + solution: NearAxisSolution, + *, + formal_radius: ArrayLike, + angular_resolution: int = 128, +) -> PlasmaHessianData: + """Evaluate the complete local plasma Hessian and external vacuum jet.""" + + gradient = plasma_gradient_on_axis( + solution, + formal_radius=formal_radius, + angular_resolution=angular_resolution, + ) + hessian_frenet = _plasma_hessian_frenet( + solution, + gradient.gradient_frenet, + gradient.field.current_source, + ) + frame = jnp.stack( + ( + solution.geometry.tangent_cartesian, + solution.geometry.normal_cartesian, + solution.geometry.binormal_cartesian, + ), + axis=-2, + ) + hessian = jnp.einsum( + "...gi,...gab,...aj,...bk->...ijk", + frame, + hessian_frenet, + frame, + frame, + ) + external_hessian = solution.grad_grad_B_axis - hessian + external_symmetric = _symmetrize_rank3(external_hessian) + external_stf = project_symmetric_trace_free_rank3(external_hessian) + external_trace = jnp.einsum("...iik->...k", external_hessian) + external_symmetry = external_hessian - external_symmetric + derivative_asymmetry = hessian - jnp.swapaxes(hessian, -1, -2) + radius = gradient.field.current_source.formal_radius + axis_scale = solution.geometry.abs_G0_over_B0 + logarithm = jnp.abs(jnp.log(radius / axis_scale)) + remainder = jnp.max(jnp.abs(hessian)) * (radius / axis_scale) ** 2 * (1 + logarithm) + return PlasmaHessianData( + field=gradient, + hessian=hessian, + hessian_frenet=hessian_frenet, + external_hessian=external_hessian, + external_hessian_stf=external_stf, + external_hessian_independent=pack_symmetric_trace_free_rank3(external_hessian), + maximum_derivative_asymmetry=jnp.max(jnp.abs(derivative_asymmetry)), + maximum_external_symmetry_error=jnp.max(jnp.abs(external_symmetry)), + maximum_external_trace=jnp.max(jnp.abs(external_trace)), + estimated_hessian_remainder=remainder, + ) diff --git a/src/pyqsc_jax/plotting.py b/src/pyqsc_jax/plotting.py new file mode 100644 index 0000000..b1e65cd --- /dev/null +++ b/src/pyqsc_jax/plotting.py @@ -0,0 +1,320 @@ +"""Optional Matplotlib plotting helpers that return their figure objects.""" + +from __future__ import annotations + +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np + +from pyqsc_jax.axis import Axis, evaluate_axis +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.plasma import PlasmaHessianData +from pyqsc_jax.vmec import frenet_displacements + + +def _matplotlib(): + try: + import matplotlib.pyplot as plt + except ImportError as error: + raise ImportError( + "Plotting requires Matplotlib. Install pyqsc-jax with the 'plot' extra." + ) from error + return plt + + +def plot_axis( + solution_or_axis: NearAxisSolution | Axis, + *, + ax: Any = None, + samples: int = 361, + label: str | None = None, + **plot_kwargs: Any, +): + """Plot a full-torus magnetic axis and return ``(figure, axes)``.""" + + if not isinstance(samples, int) or isinstance(samples, bool) or samples < 4: + raise ValueError("samples must be an integer >= 4.") + axis = ( + solution_or_axis.inputs.axis + if isinstance(solution_or_axis, NearAxisSolution) + else solution_or_axis + ) + if not isinstance(axis, Axis): + raise TypeError("solution_or_axis must be an Axis or NearAxisSolution.") + phi = jnp.linspace(0, 2 * jnp.pi, samples) + axis_samples = evaluate_axis(axis, phi) + x = axis_samples.R * jnp.cos(phi) + y = axis_samples.R * jnp.sin(phi) + z = axis_samples.Z + + plt = _matplotlib() + if ax is None: + figure = plt.figure(figsize=(5.5, 4.5)) + ax = figure.add_subplot(111, projection="3d") + else: + figure = ax.figure + defaults = {"linewidth": 2.0} + defaults.update(plot_kwargs) + ax.plot(np.asarray(x), np.asarray(y), np.asarray(z), label=label, **defaults) + ax.set_xlabel("x [m]") + ax.set_ylabel("y [m]") + ax.set_zlabel("z [m]") + ax.set_box_aspect((1, 1, 1)) + if label is not None: + ax.legend() + return figure, ax + + +def surface_coordinates( + solution: NearAxisSolution, + *, + radius: float = 0.05, + ntheta: int = 32, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Return full-torus Cartesian coordinates of a near-axis surface.""" + + if radius <= 0: + raise ValueError("radius must be positive.") + if not isinstance(ntheta, int) or isinstance(ntheta, bool) or ntheta < 4: + raise ValueError("ntheta must be an integer >= 4.") + theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False)[:, None] + phi0 = solution.phi[None, :] + X, Y, Z = frenet_displacements( + solution, + jnp.asarray(radius), + theta, + phi0, + ) + axis_position = jnp.stack( + ( + solution.R0 * jnp.cos(solution.phi), + solution.R0 * jnp.sin(solution.phi), + solution.Z0, + ), + axis=-1, + ) + period_position = ( + axis_position[None, :, :] + + X[..., None] * solution.geometry.normal_cartesian[None, :, :] + + Y[..., None] * solution.geometry.binormal_cartesian[None, :, :] + + Z[..., None] * solution.geometry.tangent_cartesian[None, :, :] + ) + period = 2 * jnp.pi / solution.inputs.axis.nfp + period_angles = jnp.arange(solution.inputs.axis.nfp) * period + cosine = jnp.cos(period_angles)[None, :, None] + sine = jnp.sin(period_angles)[None, :, None] + x_period = period_position[..., 0][:, None, :] + y_period = period_position[..., 1][:, None, :] + x = x_period * cosine - y_period * sine + y = x_period * sine + y_period * cosine + z = jnp.broadcast_to( + period_position[..., 2][:, None, :], + x.shape, + ) + x = x.reshape(ntheta, -1) + y = y.reshape(ntheta, -1) + z = z.reshape(ntheta, -1) + return ( + jnp.concatenate((x, x[:, :1]), axis=1), + jnp.concatenate((y, y[:, :1]), axis=1), + jnp.concatenate((z, z[:, :1]), axis=1), + ) + + +def plot_surface_3d( + solution: NearAxisSolution, + *, + radius: float = 0.05, + ntheta: int = 32, + ax: Any = None, + cmap: str = "viridis", + alpha: float = 0.9, + plot_axis_line: bool = True, + **surface_kwargs: Any, +): + """Plot a full-torus near-axis surface and return ``(figure, axes)``.""" + + x, y, z = surface_coordinates(solution, radius=radius, ntheta=ntheta) + plt = _matplotlib() + if ax is None: + figure = plt.figure(figsize=(6.0, 4.8)) + ax = figure.add_subplot(111, projection="3d") + else: + figure = ax.figure + defaults = { + "cmap": cmap, + "linewidth": 0, + "antialiased": True, + "alpha": alpha, + } + defaults.update(surface_kwargs) + ax.plot_surface(np.asarray(x), np.asarray(y), np.asarray(z), **defaults) + if plot_axis_line: + phi = jnp.linspace(0, 2 * jnp.pi, 361) + axis_samples = evaluate_axis(solution.inputs.axis, phi) + ax.plot( + np.asarray(axis_samples.R * jnp.cos(phi)), + np.asarray(axis_samples.R * jnp.sin(phi)), + np.asarray(axis_samples.Z), + color="black", + linewidth=1.8, + ) + ax.set_xlabel("x [m]") + ax.set_ylabel("y [m]") + ax.set_zlabel("z [m]") + ax.set_box_aspect((1, 1, 0.5)) + return figure, ax + + +def plot_b20( + solution: NearAxisSolution, + *, + ax: Any = None, + label: str | None = None, + **plot_kwargs: Any, +): + """Plot the nonconstant second-order field-strength coefficient.""" + + if solution.second_order is None: + raise ValueError("B20 plotting requires an r2 or r3 solution.") + plt = _matplotlib() + if ax is None: + figure, ax = plt.subplots(figsize=(6.0, 3.6)) + else: + figure = ax.figure + normalized_angle = np.asarray(solution.varphi * solution.inputs.axis.nfp / (2 * jnp.pi)) + defaults = {"linewidth": 2.0} + defaults.update(plot_kwargs) + ax.plot( + normalized_angle, + np.asarray(solution.B20_anomaly), + label=label, + **defaults, + ) + ax.set_xlabel("Boozer angle / field period") + ax.set_ylabel(r"$B_{20}-\langle B_{20}\rangle$ [T/m$^2$]") + if label is not None: + ax.legend() + return figure, ax + + +def plot_field_jet_norms( + result: PlasmaHessianData, + *, + axes: Any = None, +): + """Plot total, plasma, and external field-jet Frobenius norms.""" + + plt = _matplotlib() + if axes is None: + figure, axes = plt.subplots(3, 1, figsize=(7.0, 7.5), sharex=True) + else: + axes = np.asarray(axes) + if axes.shape != (3,): + raise ValueError("axes must contain exactly three Matplotlib axes.") + figure = axes[0].figure + + plasma_field = result.field.field.field + external_field = result.field.external_field + total_field = plasma_field + external_field + plasma_gradient = result.field.gradient + external_gradient = result.field.external_gradient + total_gradient = plasma_gradient + external_gradient + plasma_hessian = result.hessian + external_hessian = result.external_hessian + total_hessian = plasma_hessian + external_hessian + samples = np.arange(total_field.shape[0]) / total_field.shape[0] + + tensors = ( + (total_field, plasma_field, external_field, r"$|B|$ [T]"), + ( + total_gradient, + plasma_gradient, + external_gradient, + r"$|\nabla B|_F$ [T/m]", + ), + ( + total_hessian, + plasma_hessian, + external_hessian, + r"$|\nabla\nabla B|_F$ [T/m$^2$]", + ), + ) + for axis, (total, plasma, external, ylabel) in zip(axes, tensors, strict=True): + component_axes = tuple(range(1, total.ndim)) + axis.plot( + samples, + np.sqrt(np.sum(np.square(np.asarray(total)), axis=component_axes)), + label="total", + ) + axis.plot( + samples, + np.sqrt(np.sum(np.square(np.asarray(plasma)), axis=component_axes)), + label="plasma", + ) + axis.plot( + samples, + np.sqrt(np.sum(np.square(np.asarray(external)), axis=component_axes)), + label="external", + ) + axis.set_ylabel(ylabel) + axes[-1].set_xlabel("axis sample / field period") + axes[0].legend(ncol=3) + figure.tight_layout() + return figure, axes + + +def field_split_frenet_components( + result: PlasmaHessianData, + solution: NearAxisSolution, +) -> jax.Array: + """Return total/plasma/external field components in the Frenet frame. + + The result has shape ``(3, 3, nphi)``. The first index orders + ``(total, plasma, external)`` and the second orders + ``(tangent, normal, binormal)``. + """ + + plasma = result.field.field.field + external = result.field.external_field + fields = jnp.stack((plasma + external, plasma, external)) + frames = jnp.stack( + ( + solution.geometry.tangent_cartesian, + solution.geometry.normal_cartesian, + solution.geometry.binormal_cartesian, + ) + ) + return jnp.einsum("fpi,cpi->fcp", fields, frames) + + +def plot_field_split_components( + result: PlasmaHessianData, + solution: NearAxisSolution, + *, + axes: Any = None, +): + """Plot angle-dependent total/plasma/external Frenet field components.""" + + plt = _matplotlib() + if axes is None: + figure, axes = plt.subplots(3, 1, figsize=(7.0, 7.2), sharex=True) + else: + axes = np.asarray(axes) + if axes.shape != (3,): + raise ValueError("axes must contain exactly three Matplotlib axes.") + figure = axes[0].figure + components = np.asarray(field_split_frenet_components(result, solution)) + angle = np.asarray(solution.varphi * solution.inputs.axis.nfp / (2 * jnp.pi)) + contributions = ("total", "plasma", "external") + component_labels = (r"$B_t$ [T]", r"$B_n$ [T]", r"$B_b$ [T]") + for component_index, (axis, ylabel) in enumerate(zip(axes, component_labels, strict=True)): + for field_index, label in enumerate(contributions): + axis.plot(angle, components[field_index, component_index], label=label) + axis.set_ylabel(ylabel) + axes[-1].set_xlabel("Boozer angle / field period") + axes[0].legend(ncol=3) + figure.tight_layout() + return figure, axes diff --git a/src/pyqsc_jax/second_order.py b/src/pyqsc_jax/second_order.py new file mode 100644 index 0000000..bc4cb34 --- /dev/null +++ b/src/pyqsc_jax/second_order.py @@ -0,0 +1,473 @@ +"""Complete second-order quasisymmetric near-axis coefficient solve.""" + +from dataclasses import dataclass, replace + +import jax +import jax.numpy as jnp + +from pyqsc_jax.models import NearAxisSolution, SecondOrderData +from pyqsc_jax.solvers import implicit_dense_linear_solve + +MU0 = 4e-7 * jnp.pi + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class SecondOrderResiduals: + """Independent residuals of the four coupled second-order equations.""" + + force_balance_1: jax.Array + force_balance_2: jax.Array + area_constraint_1: jax.Array + area_constraint_2: jax.Array + + @property + def maximum_absolute(self) -> jax.Array: + """Maximum absolute collocation residual across all four equations.""" + + return jnp.max( + jnp.abs( + jnp.stack( + ( + self.force_balance_1, + self.force_balance_2, + self.area_constraint_1, + self.area_constraint_2, + ) + ) + ) + ) + + +def _assemble_periodic_system( + first_order: NearAxisSolution, + *, + X2s: jax.Array, + X2c: jax.Array, + Z20: jax.Array, + Z2s: jax.Array, + Z2c: jax.Array, + beta_1s: jax.Array, +) -> tuple[ + jax.Array, + jax.Array, + jax.Array, + jax.Array, + jax.Array, + jax.Array, +]: + """Assemble the coupled periodic system for ``X20`` and ``Y20``.""" + + inputs = first_order.inputs + geometry = first_order.geometry + D = geometry.d_d_varphi + lp = geometry.abs_G0_over_B0 + X1c = first_order.X1c + Y1s = first_order.Y1s + Y1c = first_order.Y1c + curvature = first_order.curvature + torsion = first_order.torsion + iotaN = first_order.iotaN + sign_product = inputs.sG * inputs.spsi + current_factor = inputs.spsi * inputs.I2 / inputs.B0 + + Y2s_from_X20 = -sign_product * curvature**2 / inputs.etabar**2 + Y2s_inhomogeneous = sign_product * ( + -curvature / 2 + curvature**2 / inputs.etabar**2 * (-X2c + X2s * first_order.sigma) + ) + Y2c_from_X20 = -sign_product * curvature**2 * first_order.sigma / inputs.etabar**2 + Y2c_inhomogeneous = ( + sign_product * curvature**2 / inputs.etabar**2 * (X2s + X2c * first_order.sigma) + ) + + fX0_from_X20 = -4 * sign_product * lp * (Y2c_from_X20 * Z2s - Y2s_from_X20 * Z2c) + fX0_from_Y20 = -torsion * lp - 4 * sign_product * lp * Z2s + 2 * current_factor * lp + fX0_inhomogeneous = ( + curvature * lp * Z20 + - 4 * sign_product * lp * (Y2c_inhomogeneous * Z2s - Y2s_inhomogeneous * Z2c) + - current_factor * (curvature * sign_product / 2) * lp + + beta_1s * lp * Y1c / 2 + ) + + fXs_from_X20 = ( + -torsion * lp * Y2s_from_X20 + - 4 * sign_product * lp * Y2c_from_X20 * Z20 + + 2 * current_factor * lp * Y2s_from_X20 + ) + fXs_from_Y20 = -4 * sign_product * lp * (-Z2c + Z20) + fXs_inhomogeneous = ( + D @ X2s + - 2 * iotaN * X2c + - torsion * lp * Y2s_inhomogeneous + + curvature * lp * Z2s + - 4 * sign_product * lp * Y2c_inhomogeneous * Z20 + - current_factor * (curvature * sign_product / 2 - 2 * Y2s_inhomogeneous) * lp + - lp * beta_1s * Y1s / 2 + ) + + fXc_from_X20 = ( + -torsion * lp * Y2c_from_X20 + + 4 * sign_product * lp * Y2s_from_X20 * Z20 + + 2 * current_factor * lp * Y2c_from_X20 + ) + fXc_from_Y20 = -torsion * lp - 4 * sign_product * lp * Z2s + 2 * current_factor * lp + fXc_inhomogeneous = ( + D @ X2c + + 2 * iotaN * X2s + - torsion * lp * Y2c_inhomogeneous + + curvature * lp * Z2c + + 4 * sign_product * lp * Y2s_inhomogeneous * Z20 + - current_factor * (curvature * sign_product / 2 - 2 * Y2c_inhomogeneous) * lp + - lp * beta_1s * Y1c / 2 + ) + + fY0_from_X20 = torsion * lp - 2 * current_factor * lp + fY0_from_Y20 = jnp.zeros_like(curvature) + fY0_inhomogeneous = ( + -4 * sign_product * lp * (X2s * Z2c - X2c * Z2s) + + current_factor * curvature * X1c**2 * lp / 2 + - lp * beta_1s * X1c / 2 + ) + + fYs_from_X20 = -2 * iotaN * Y2c_from_X20 - 4 * sign_product * lp * Z2c + fYs_from_Y20 = jnp.full_like(curvature, -2 * iotaN) + fYs_inhomogeneous = ( + D @ Y2s_inhomogeneous + - 2 * iotaN * Y2c_inhomogeneous + + torsion * lp * X2s + + 4 * sign_product * lp * X2c * Z20 + - 2 * current_factor * X2s * lp + ) + + fYc_from_X20 = 2 * iotaN * Y2s_from_X20 + 4 * sign_product * lp * Z2s + fYc_from_Y20 = jnp.zeros_like(curvature) + fYc_inhomogeneous = ( + D @ Y2c_inhomogeneous + + 2 * iotaN * Y2s_inhomogeneous + + torsion * lp * X2c + - 4 * sign_product * lp * X2s * Z20 + - current_factor * (-curvature * X1c**2 / 2 + 2 * X2c) * lp + + lp * beta_1s * X1c / 2 + ) + + block_00 = Y1c[:, None] * D * Y2s_from_X20[None, :] - Y1s[:, None] * D * Y2c_from_X20[None, :] + block_01 = -2 * Y1s[:, None] * D + block_10 = ( + -X1c[:, None] * D + + Y1s[:, None] * D * Y2s_from_X20[None, :] + + Y1c[:, None] * D * Y2c_from_X20[None, :] + ) + block_11 = jnp.zeros_like(D) + + block_00 = block_00 + jnp.diag( + X1c * fXs_from_X20 - Y1s * fY0_from_X20 + Y1c * fYs_from_X20 - Y1s * fYc_from_X20 + ) + block_01 = block_01 + jnp.diag( + X1c * fXs_from_Y20 - Y1s * fY0_from_Y20 + Y1c * fYs_from_Y20 - Y1s * fYc_from_Y20 + ) + block_10 = block_10 + jnp.diag( + -X1c * fX0_from_X20 + + X1c * fXc_from_X20 + - Y1c * fY0_from_X20 + + Y1s * fYs_from_X20 + + Y1c * fYc_from_X20 + ) + block_11 = block_11 + jnp.diag( + -X1c * fX0_from_Y20 + + X1c * fXc_from_Y20 + - Y1c * fY0_from_Y20 + + Y1s * fYs_from_Y20 + + Y1c * fYc_from_Y20 + ) + matrix = jnp.concatenate( + ( + jnp.concatenate((block_00, block_01), axis=1), + jnp.concatenate((block_10, block_11), axis=1), + ), + axis=0, + ) + + right_hand_side_1 = -( + X1c * fXs_inhomogeneous + - Y1s * fY0_inhomogeneous + + Y1c * fYs_inhomogeneous + - Y1s * fYc_inhomogeneous + ) + right_hand_side_2 = -( + -X1c * fX0_inhomogeneous + + X1c * fXc_inhomogeneous + - Y1c * fY0_inhomogeneous + + Y1s * fYs_inhomogeneous + + Y1c * fYc_inhomogeneous + ) + right_hand_side = jnp.concatenate((right_hand_side_1, right_hand_side_2)) + return ( + matrix, + right_hand_side, + Y2s_from_X20, + Y2s_inhomogeneous, + Y2c_from_X20, + Y2c_inhomogeneous, + ) + + +def solve_second_order( + first_order: NearAxisSolution, + *, + attach_diagnostics: bool = True, +) -> NearAxisSolution: + """Add the complete finite-pressure/current second-order solution.""" + + inputs = first_order.inputs + geometry = first_order.geometry + D = geometry.d_d_varphi + B0_over_abs_G0 = 1 / geometry.abs_G0_over_B0 + lp = geometry.abs_G0_over_B0 + X1c = first_order.X1c + Y1s = first_order.Y1s + Y1c = first_order.Y1c + curvature = first_order.curvature + torsion = first_order.torsion + iotaN = first_order.iotaN + sign_product = inputs.sG * inputs.spsi + + V1 = X1c**2 + Y1c**2 + Y1s**2 + V2 = 2 * Y1s * Y1c + V3 = X1c**2 + Y1c**2 - Y1s**2 + factor = -B0_over_abs_G0 / 8 + Z20 = factor * (D @ V1) + Z2s = factor * (D @ V2 - 2 * iotaN * V3) + Z2c = factor * (D @ V3 + 2 * iotaN * V2) + + qs = -iotaN * X1c - Y1s * torsion * lp + qc = D @ X1c - Y1c * torsion * lp + rs = D @ Y1s - iotaN * Y1c + rc = D @ Y1c + iotaN * Y1s + X1c * torsion * lp + X2s = ( + B0_over_abs_G0 + * ( + D @ Z2s + - 2 * iotaN * Z2c + + B0_over_abs_G0 * (lp**2 * inputs.B2s / inputs.B0 + (qc * qs + rc * rs) / 2) + ) + / curvature + ) + X2c = ( + B0_over_abs_G0 + * ( + D @ Z2c + + 2 * iotaN * Z2s + - B0_over_abs_G0 + * ( + -(lp**2) * inputs.B2c / inputs.B0 + + lp**2 * inputs.etabar**2 / 2 + - (qc**2 - qs**2 + rc**2 - rs**2) / 4 + ) + ) + / curvature + ) + beta_1s = -4 * sign_product * MU0 * inputs.p2 * inputs.etabar * lp / (iotaN * inputs.B0**2) + + ( + matrix, + right_hand_side, + Y2s_from_X20, + Y2s_inhomogeneous, + Y2c_from_X20, + Y2c_inhomogeneous, + ) = _assemble_periodic_system( + first_order, + X2s=X2s, + X2c=X2c, + Z20=Z20, + Z2s=Z2s, + Z2c=Z2c, + beta_1s=beta_1s, + ) + solution, linear_report = implicit_dense_linear_solve(matrix, right_hand_side) + X20, Y20 = jnp.split(solution, 2) + Y2s = Y2s_inhomogeneous + Y2s_from_X20 * X20 + Y2c = Y2c_inhomogeneous + Y2c_from_X20 * X20 + Y20 + + B20 = inputs.B0 * ( + curvature * X20 + - B0_over_abs_G0 * (D @ Z20) + + inputs.etabar**2 / 2 + - MU0 * inputs.p2 / inputs.B0**2 + - B0_over_abs_G0**2 * (qc**2 + qs**2 + rc**2 + rs**2) / 4 + ) + weights = geometry.d_l_d_phi + B20_mean = jnp.sum(B20 * weights) / jnp.sum(weights) + B20_anomaly = B20 - B20_mean + B20_residual = jnp.sqrt(jnp.sum(B20_anomaly**2 * weights) / jnp.sum(weights)) / inputs.B0 + B20_variation = jnp.max(B20) - jnp.min(B20) + G2 = -MU0 * inputs.p2 * first_order.G0 / inputs.B0**2 - first_order.iota * inputs.I2 + N_helicity = -first_order.helicity * inputs.axis.nfp + + d_X1c_d_varphi = D @ X1c + d_Y1c_d_varphi = D @ Y1c + d_Y1s_d_varphi = D @ Y1s + untwisting_angle = -first_order.helicity * inputs.axis.nfp * first_order.varphi + sine = jnp.sin(2 * untwisting_angle) + cosine = jnp.cos(2 * untwisting_angle) + + def untwist(sine_coefficient, cosine_coefficient): + return ( + sine_coefficient * cosine + cosine_coefficient * sine, + -sine_coefficient * sine + cosine_coefficient * cosine, + ) + + X2s_untwisted, X2c_untwisted = untwist(X2s, X2c) + Y2s_untwisted, Y2c_untwisted = untwist(Y2s, Y2c) + Z2s_untwisted, Z2c_untwisted = untwist(Z2s, Z2c) + second_order = SecondOrderData( + linear_report=linear_report, + V1=V1, + V2=V2, + V3=V3, + X20=X20, + X2s=X2s, + X2c=X2c, + Y20=Y20, + Y2s=Y2s, + Y2c=Y2c, + Z20=Z20, + Z2s=Z2s, + Z2c=Z2c, + beta_1s=beta_1s, + B20=B20, + B20_mean=B20_mean, + B20_anomaly=B20_anomaly, + B20_residual=B20_residual, + B20_variation=B20_variation, + G2=G2, + N_helicity=N_helicity, + d_curvature_d_varphi=D @ curvature, + d_torsion_d_varphi=D @ torsion, + d_X20_d_varphi=D @ X20, + d_X2s_d_varphi=D @ X2s, + d_X2c_d_varphi=D @ X2c, + d_Y20_d_varphi=D @ Y20, + d_Y2s_d_varphi=D @ Y2s, + d_Y2c_d_varphi=D @ Y2c, + d_Z20_d_varphi=D @ Z20, + d_Z2s_d_varphi=D @ Z2s, + d_Z2c_d_varphi=D @ Z2c, + d2_X1c_d_varphi2=D @ d_X1c_d_varphi, + d2_Y1c_d_varphi2=D @ d_Y1c_d_varphi, + d2_Y1s_d_varphi2=D @ d_Y1s_d_varphi, + X20_untwisted=X20, + X2s_untwisted=X2s_untwisted, + X2c_untwisted=X2c_untwisted, + Y20_untwisted=Y20, + Y2s_untwisted=Y2s_untwisted, + Y2c_untwisted=Y2c_untwisted, + Z20_untwisted=Z20, + Z2s_untwisted=Z2s_untwisted, + Z2c_untwisted=Z2c_untwisted, + ) + solution = replace(first_order, second_order=second_order) + if not attach_diagnostics: + return solution + from pyqsc_jax.diagnostics import mercier_diagnostics + from pyqsc_jax.field import total_field_jet + from pyqsc_jax.singularity import singularity_diagnostics + + return replace( + solution, + mercier=mercier_diagnostics(solution), + field_jet=total_field_jet(solution), + singularity=singularity_diagnostics(solution), + ) + + +def second_order_residuals(solution: NearAxisSolution) -> SecondOrderResiduals: + """Evaluate all four r2 equations independently of matrix assembly.""" + + r2 = solution.second_order + if r2 is None: + raise ValueError("A second-order solution is required.") + inputs = solution.inputs + D = solution.geometry.d_d_varphi + lp = solution.geometry.abs_G0_over_B0 + sign_product = inputs.sG * inputs.spsi + current_factor = inputs.spsi * inputs.I2 / inputs.B0 + + fX0 = ( + D @ r2.X20 + - solution.torsion * lp * r2.Y20 + + solution.curvature * lp * r2.Z20 + - 4 * sign_product * lp * (r2.Y2c * r2.Z2s - r2.Y2s * r2.Z2c) + - current_factor * (solution.curvature * solution.X1c * solution.Y1c / 2 - 2 * r2.Y20) * lp + + lp * r2.beta_1s * solution.Y1c / 2 + ) + fXs = ( + D @ r2.X2s + - 2 * solution.iotaN * r2.X2c + - solution.torsion * lp * r2.Y2s + + solution.curvature * lp * r2.Z2s + - 4 * sign_product * lp * (-r2.Y20 * r2.Z2c + r2.Y2c * r2.Z20) + - current_factor * (solution.curvature * solution.X1c * solution.Y1s / 2 - 2 * r2.Y2s) * lp + - lp * r2.beta_1s * solution.Y1s / 2 + ) + fXc = ( + D @ r2.X2c + + 2 * solution.iotaN * r2.X2s + - solution.torsion * lp * r2.Y2c + + solution.curvature * lp * r2.Z2c + - 4 * sign_product * lp * (r2.Y20 * r2.Z2s - r2.Y2s * r2.Z20) + - current_factor * (solution.curvature * solution.X1c * solution.Y1c / 2 - 2 * r2.Y2c) * lp + - lp * r2.beta_1s * solution.Y1c / 2 + ) + fY0 = ( + D @ r2.Y20 + + solution.torsion * lp * r2.X20 + - 4 * sign_product * lp * (r2.X2s * r2.Z2c - r2.X2c * r2.Z2s) + - current_factor * (-solution.curvature * solution.X1c**2 / 2 + 2 * r2.X20) * lp + - lp * r2.beta_1s * solution.X1c / 2 + ) + fYs = ( + D @ r2.Y2s + - 2 * solution.iotaN * r2.Y2c + + solution.torsion * lp * r2.X2s + - 4 * sign_product * lp * (r2.X20 * r2.Z2c - r2.X2c * r2.Z20) + - 2 * current_factor * r2.X2s * lp + ) + fYc = ( + D @ r2.Y2c + + 2 * solution.iotaN * r2.Y2s + + solution.torsion * lp * r2.X2c + - 4 * sign_product * lp * (r2.X2s * r2.Z20 - r2.X20 * r2.Z2s) + - current_factor * (-solution.curvature * solution.X1c**2 / 2 + 2 * r2.X2c) * lp + + lp * r2.beta_1s * solution.X1c / 2 + ) + force_balance_1 = ( + solution.X1c * fXs - solution.Y1s * fY0 + solution.Y1c * fYs - solution.Y1s * fYc + ) + force_balance_2 = ( + -solution.X1c * fX0 + + solution.X1c * fXc + - solution.Y1c * fY0 + + solution.Y1s * fYs + + solution.Y1c * fYc + ) + area_constraint_1 = ( + -solution.X1c * r2.Y2c + + solution.X1c * r2.Y20 + + r2.X2s * solution.Y1s + + r2.X2c * solution.Y1c + - r2.X20 * solution.Y1c + ) + area_constraint_2 = ( + solution.X1c * r2.Y2s + + r2.X2c * solution.Y1s + - r2.X2s * solution.Y1c + + r2.X20 * solution.Y1s + + sign_product * solution.X1c * solution.curvature / 2 + ) + return SecondOrderResiduals( + force_balance_1=force_balance_1, + force_balance_2=force_balance_2, + area_constraint_1=area_constraint_1, + area_constraint_2=area_constraint_2, + ) diff --git a/src/pyqsc_jax/shear.py b/src/pyqsc_jax/shear.py new file mode 100644 index 0000000..4dda027 --- /dev/null +++ b/src/pyqsc_jax/shear.py @@ -0,0 +1,372 @@ +"""Magnetic shear from the order-r-cubed generalized sigma equation. + +The equations are adapted from ``landreman/pyQSC:qsc/calculate_r3.py`` at the +audited BSD-2-Clause upstream commit recorded in the refactor baseline. Their +source derivation is Rodríguez et al., Physics of Plasmas 29, 012507 (2022). +""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Any + +import jax.numpy as jnp + +from pyqsc_jax.models import NearAxisSolution, ShearData + +ArrayLike = Any + + +def solve_magnetic_shear( + solution: NearAxisSolution, + *, + B31c: ArrayLike = 0.0, +) -> NearAxisSolution: + """Attach the standard-MHS order-r-squared transform correction. + + ``B31c`` uses the inverse-field-squared convention of the generalized + sigma equation. Current variation at this order and ``B31s`` are zero. + """ + + r2 = solution.second_order + if r2 is None: + raise ValueError("A second-order solution is required to calculate magnetic shear.") + inputs = solution.inputs + if inputs.sG != 1 or inputs.spsi != 1: + raise NotImplementedError( + "Magnetic shear currently requires sG=spsi=1; the upstream signs are not generalized." + ) + + B31c = jnp.asarray(B31c) + if B31c.ndim != 0: + raise ValueError("B31c must be a scalar.") + + D = solution.geometry.d_d_varphi + epsilon_scale = jnp.sqrt(2 / inputs.B0) + scale2 = epsilon_scale**2 + G2 = r2.G2 * scale2 + G0 = solution.G0 + I2 = inputs.I2 * scale2 + X1c = solution.X1c * epsilon_scale + Y1c = solution.Y1c * epsilon_scale + Y1s = solution.Y1s * epsilon_scale + X20 = r2.X20 * scale2 + X2s = r2.X2s * scale2 + X2c = r2.X2c * scale2 + Y20 = r2.Y20 * scale2 + Y2s = r2.Y2s * scale2 + Y2c = r2.Y2c * scale2 + Z20 = r2.Z20 * scale2 + Z2s = r2.Z2s * scale2 + Z2c = r2.Z2c * scale2 + torsion = -solution.torsion + curvature = solution.curvature + iota = solution.iotaN + dldp = solution.geometry.abs_G0_over_B0 + dX1c = D @ X1c + dY1c = D @ Y1c + dY1s = D @ Y1s + dZ20 = D @ Z20 + dZ2c = D @ Z2c + dZ2s = D @ Z2s + dX20 = D @ X20 + dX2c = D @ X2c + dX2s = D @ X2s + dY20 = D @ Y20 + dY2c = D @ Y2c + dY2s = D @ Y2s + + inverse_B0_squared = 1 / inputs.B0**2 + Ba0 = G0 + Ba1 = G2 + solution.iotaN * I2 + eta = inputs.etabar * jnp.sqrt(2) * inverse_B0_squared**0.25 + B1c = -2 * inverse_B0_squared * eta + B20 = ( + (0.75 * inputs.etabar**2 / jnp.sqrt(inverse_B0_squared) - r2.B20) + * 4 + * inverse_B0_squared**2 + ) + B31s = jnp.asarray(0, dtype=B31c.dtype) + I4 = jnp.asarray(0, dtype=B31c.dtype) + + Z31c = ( + -1 + / (3 * Ba0 * X1c * Y1s) + * ( + 2 * iota * (X1c * X2s - Y2c * Y1s + Y1c * Y2s) + - 2 * Ba0 * X2s * Y1c * Z20 + + 2 * Ba0 * X2c * Y1s * Z20 + + 2 * Ba0 * X1c * Y2s * Z20 + - 4 * Ba0 * X2s * Y1c * Z2c + - 2 * Ba0 * X20 * Y1s * Z2c + + 4 * Ba0 * X1c * Y2s * Z2c + - dldp + * ( + torsion * (2 * X20 * Y1c + X2c * Y1c - 2 * X1c * Y20 - X1c * Y2c + X2s * Y1s) + + I2 * (2 * X20 * Y1c + X2c * Y1c - 2 * X1c * Y20 - X1c * Y2c + X2s * Y1s) + - 2 * curvature * X1c * Z20 + - curvature * X1c * Z2c + ) + + 2 * Ba0 * X20 * Y1c * Z2s + + 4 * Ba0 * X2c * Y1c * Z2s + - 2 * Ba0 * X1c * Y20 * Z2s + - 4 * Ba0 * X1c * Y2c * Z2s + + 2 * X1c * dX20 + + X1c * dX2c + + 2 * Y1c * dY20 + + Y1c * dY2c + + Y1s * dY2s + ) + ) + dZ31c = D @ Z31c + + Z31s = ( + 1 + / (3 * Ba0 * X1c * Y1s) + * ( + 2 * iota * (X1c * X2c + Y1c * Y2c + Y1s * Y2s) + - 2 * Ba0 * X2c * Y1c * Z20 + + 2 * Ba0 * X1c * Y2c * Z20 + - 2 * Ba0 * X2s * Y1s * Z20 + + 2 * Ba0 * X20 * Y1c * Z2c + - 2 * Ba0 * X1c * Y20 * Z2c + + 4 * Ba0 * X2s * Y1s * Z2c + + 2 * Ba0 * X20 * Y1s * Z2s + - 4 * Ba0 * X2c * Y1s * Z2s + + dldp + * ( + I2 * X2s * Y1c + + 2 * I2 * X20 * Y1s + - I2 * X2c * Y1s + - I2 * X1c * Y2s + + torsion * (X2s * Y1c + 2 * X20 * Y1s - X2c * Y1s - X1c * Y2s) + - curvature * X1c * Z2s + ) + - X1c * dX2s + - 2 * Y1s * dY20 + + Y1s * dY2c + - Y1c * dY2s + ) + ) + dZ31s = D @ Z31s + + X31c = ( + 1 + / (2 * dldp**2 * curvature) + * ( + -2 * Ba0 * Ba1 * B1c + - Ba0**2 * B31c + + 2 * dldp**2 * torsion**2 * X1c * X20 + + 2 * iota**2 * X1c * X2c + + dldp**2 * torsion**2 * X1c * X2c + + dldp**2 * curvature**2 * X1c * (2 * X20 + X2c) + + 3 * dldp * iota * torsion * X2s * Y1c + + 2 * dldp**2 * torsion**2 * Y1c * Y20 + + 2 * iota**2 * Y1c * Y2c + + dldp**2 * torsion**2 * Y1c * Y2c + - 2 * dldp * iota * torsion * X20 * Y1s + - 3 * dldp * iota * torsion * X2c * Y1s + - 3 * dldp * iota * torsion * X1c * Y2s + + 2 * iota**2 * Y1s * Y2s + + dldp**2 * torsion**2 * Y1s * Y2s + + 2 * dldp * iota * Z31s + + 2 * iota * X2s * dX1c + + 2 * dldp * torsion * Y20 * dX1c + + dldp * torsion * Y2c * dX1c + + 2 * dldp * torsion * Y1c * dX20 + + 2 * dX1c * dX20 + + dldp * torsion * Y1c * dX2c + + dX1c * dX2c + - iota * X1c * dX2s + + dldp * torsion * Y1s * dX2s + - 2 * dldp * torsion * X20 * dY1c + - dldp * torsion * X2c * dY1c + + 2 * iota * Y2s * dY1c + - 2 * dldp * torsion * X1c * dY20 + + 2 * iota * Y1s * dY20 + + 2 * dY1c * dY20 + - dldp * torsion * X1c * dY2c + + iota * Y1s * dY2c + + dY1c * dY2c + - dldp * torsion * X2s * dY1s + - 2 * iota * Y2c * dY1s + - iota * Y1c * dY2s + + dY1s * dY2s + + dldp + * curvature + * ( + -3 * iota * X1c * Z2s + + dldp * torsion * (Y1c * (2 * Z20 + Z2c) + Y1s * Z2s) + + 2 * Z20 * dX1c + + Z2c * dX1c + - 2 * X1c * dZ20 + - X1c * dZ2c + ) + + 2 * dldp * dZ31c + ) + ) + + X31s = ( + 1 + / (2 * dldp**2 * curvature) + * ( + -(Ba0**2) * B31s + + dldp**2 * curvature**2 * X1c * X2s + + dldp**2 * torsion**2 * X1c * X2s + + 2 * dldp**2 * torsion**2 * Y20 * Y1s + - dldp**2 * torsion**2 * Y2c * Y1s + + dldp**2 * torsion**2 * Y1c * Y2s + + 2 * iota**2 * (X1c * X2s - Y2c * Y1s + Y1c * Y2s) + + 2 * dldp**2 * curvature * torsion * Y1s * Z20 + - dldp**2 * curvature * torsion * Y1s * Z2c + + dldp**2 * curvature * torsion * Y1c * Z2s + + dldp * torsion * Y2s * dX1c + + dldp * curvature * Z2s * dX1c + + 2 * dldp * torsion * Y1s * dX20 + - dldp * torsion * Y1s * dX2c + + dldp * torsion * Y1c * dX2s + + dX1c * dX2s + - dldp * torsion * X2s * dY1c + - 2 * dldp * torsion * X20 * dY1s + + dldp * torsion * X2c * dY1s + + 2 * dY20 * dY1s + - dY2c * dY1s + - dldp * torsion * X1c * dY2s + + dY1c * dY2s + + iota + * ( + dldp + * torsion + * (2 * X20 * Y1c - 3 * X2c * Y1c - 2 * X1c * Y20 + 3 * X1c * Y2c - 3 * X2s * Y1s) + + dldp * curvature * X1c * (-2 * Z20 + 3 * Z2c) + - 2 * dldp * Z31c + - 2 * X2c * dX1c + - 2 * X1c * dX20 + + X1c * dX2c + - 2 * Y2c * dY1c + - 2 * Y1c * dY20 + + Y1c * dY2c + - 2 * Y2s * dY1s + + Y1s * dY2s + ) + - dldp * curvature * X1c * dZ2s + + 2 * dldp * dZ31s + ) + ) + dX31s = D @ X31s + + Y31s = ( + 1 + / (4 * Ba0 * X1c) + * ( + -2 * Ba1 * X1c * Y1s + + 2 * iota * I2 * X1c * Y1s + - dldp * (4 * curvature * X20 + torsion * I2 * (X1c**2 + Y1c**2 + Y1s**2)) + + 4 * Ba0 * (X31s * Y1c + 2 * X2s * Y2c - X31c * Y1s - 2 * X2c * Y2s) + - I2 * Y1c * dX1c + + I2 * X1c * dY1c + + 4 * dZ20 + ) + ) + dY31s = D @ Y31s + + Lambda_tilde = 2 / Y1s**2 * ( + Ba0 * inverse_B0_squared * I4 + (Ba1 * inverse_B0_squared + Ba0 * B20) * I2 + ) + 1 / Y1s**2 * ( + -2 + * iota + * ( + 2 * X2c**2 + + X1c * X31c + + 2 * X2s**2 + + 2 * Y2c**2 + + 2 * Y2s**2 + + Y1s * Y31s + + 2 * Z2c**2 + + 2 * Z2s**2 + ) + + 2 + * dldp + * ( + torsion * (-X31s * Y1c - 2 * X2s * Y2c + X31c * Y1s + 2 * X2c * Y2s + X1c * Y31s) + + curvature * (-2 * X2s * Z2c + 2 * X2c * Z2s + X1c * Z31s) + ) + - X31s * dX1c + - 2 * X2s * dX2c + + 2 * X2c * dX2s + + X1c * dX31s + - Y31s * dY1c + - 2 * Y2s * dY2c + + 2 * Y2c * dY2s + + Y1c * dY31s + - 2 * Z2s * dZ2c + + 2 * Z2c * dZ2s + ) + + reduced_derivative = D[1:, 1:] + symmetric_integral = jnp.concatenate( + ( + jnp.zeros(1, dtype=solution.sigma.dtype), + jnp.linalg.solve(reduced_derivative, solution.sigma[1:]), + ) + ) + symmetric_factor = jnp.exp(2 * iota * symmetric_integral) + denominator_factor = (X1c**2 + Y1c**2 + Y1s**2) / Y1s**2 + symmetric_numerator = jnp.sum( + symmetric_factor * Lambda_tilde * solution.geometry.d_varphi_d_phi + ) + symmetric_denominator = jnp.sum( + symmetric_factor * denominator_factor * solution.geometry.d_varphi_d_phi + ) + + sigma_average = jnp.sum(solution.sigma * solution.geometry.d_varphi_d_phi) / inputs.nphi + periodic_integral = jnp.linalg.solve( + reduced_derivative, + solution.sigma[1:] - sigma_average, + ) + general_integral = jnp.concatenate( + ( + jnp.zeros(1, dtype=solution.sigma.dtype), + periodic_integral + sigma_average * solution.varphi[1:], + ) + ) + general_factor = jnp.exp(2 * iota * general_integral) + period = 2 * jnp.pi / inputs.axis.nfp + factor_extended = jnp.concatenate( + (general_factor, jnp.exp(jnp.asarray([2 * iota * sigma_average * period]))) + ) + Lambda_extended = jnp.concatenate((Lambda_tilde, Lambda_tilde[:1])) + denominator_extended = jnp.concatenate((denominator_factor, denominator_factor[:1])) + varphi_extended = jnp.concatenate((solution.varphi, jnp.asarray([period]))) + general_numerator = jnp.trapezoid(factor_extended * Lambda_extended, varphi_extended) + general_denominator = jnp.trapezoid( + factor_extended * denominator_extended, + varphi_extended, + ) + + stellarator_symmetric = ( + (inputs.sigma0 == 0) + & (jnp.max(jnp.abs(inputs.axis.rs)) == 0) + & (jnp.max(jnp.abs(inputs.axis.zc)) == 0) + ) + numerator = jnp.where(stellarator_symmetric, symmetric_numerator, general_numerator) + denominator = jnp.where(stellarator_symmetric, symmetric_denominator, general_denominator) + integrating_factor = jnp.where(stellarator_symmetric, symmetric_factor, general_factor) + iota2 = inputs.B0 * numerator / (2 * denominator) + + shear = ShearData( + B31c=B31c, + iota2=iota2, + numerator=numerator, + denominator=denominator, + Lambda_tilde=Lambda_tilde, + integrating_factor=integrating_factor, + sigma_average=sigma_average, + Z31c=Z31c, + Z31s=Z31s, + X31c=X31c, + X31s=X31s, + Y31s=Y31s, + stellarator_symmetric=stellarator_symmetric, + ) + return replace(solution, shear=shear) diff --git a/src/pyqsc_jax/singularity.py b/src/pyqsc_jax/singularity.py new file mode 100644 index 0000000..c25e7b0 --- /dev/null +++ b/src/pyqsc_jax/singularity.py @@ -0,0 +1,215 @@ +"""Regular-coordinate singular-radius diagnostics.""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp + +from pyqsc_jax.field import ( + _differentiate_cylindrical_vector, + _regular_map_vectors, + _to_cartesian, +) +from pyqsc_jax.models import NearAxisSolution, SingularityDiagnostics + + +def _determinant_coefficients( + solution: NearAxisSolution, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + geometry = solution.geometry + length_scale = geometry.abs_G0_over_B0 + d1, d2, h11, h12, h22 = _regular_map_vectors(solution) + derivative = lambda vector: _differentiate_cylindrical_vector(vector, solution) # noqa: E731 + cartesian = lambda vector: _to_cartesian(vector, solution) # noqa: E731 + arguments = ( + cartesian(length_scale * geometry.tangent_cylindrical), + cartesian(d1), + cartesian(d2), + cartesian(derivative(d1)), + cartesian(derivative(d2)), + cartesian(h11), + cartesian(h12), + cartesian(h22), + cartesian(derivative(h11)), + cartesian(derivative(h12)), + cartesian(derivative(h22)), + ) + ( + axis_tangent, + d1_value, + d2_value, + d1_prime, + d2_prime, + h11_value, + h12_value, + h22_value, + h11_prime, + h12_prime, + h22_prime, + ) = arguments + + determinant = lambda first, second, third: jnp.einsum( # noqa: E731 + "ni,ni->n", first, jnp.cross(second, third) + ) + g0 = determinant(axis_tangent, d1_value, d2_value) + g1c = ( + determinant(d1_prime, d1_value, d2_value) + + determinant(axis_tangent, h11_value, d2_value) + + determinant(axis_tangent, d1_value, h12_value) + ) + g1s = ( + determinant(d2_prime, d1_value, d2_value) + + determinant(axis_tangent, h12_value, d2_value) + + determinant(axis_tangent, d1_value, h22_value) + ) + q1_squared = ( + 0.5 * determinant(h11_prime, d1_value, d2_value) + + determinant(d1_prime, h11_value, d2_value) + + determinant(d1_prime, d1_value, h12_value) + + determinant(axis_tangent, h11_value, h12_value) + ) + q2_squared = ( + 0.5 * determinant(h22_prime, d1_value, d2_value) + + determinant(d2_prime, h12_value, d2_value) + + determinant(d2_prime, d1_value, h22_value) + + determinant(axis_tangent, h12_value, h22_value) + ) + q1_q2 = ( + determinant(h12_prime, d1_value, d2_value) + + determinant(d1_prime, h12_value, d2_value) + + determinant(d2_prime, h11_value, d2_value) + + determinant(d1_prime, d1_value, h22_value) + + determinant(d2_prime, d1_value, h12_value) + + determinant(axis_tangent, h11_value, h22_value) + + determinant(axis_tangent, h12_value, h12_value) + ) + hessian_11 = 2 * q1_squared + hessian_22 = 2 * q2_squared + g20 = (hessian_11 + hessian_22) / 4 + g2c = (hessian_11 - hessian_22) / 4 + g2s = q1_q2 / 2 + return g0, g1c, g1s, g20, g2s, g2c + + +def _positive_quadratic_root( + constant: jax.Array, + linear: jax.Array, + quadratic: jax.Array, +) -> jax.Array: + tolerance = 100 * jnp.finfo(quadratic.dtype).eps + discriminant = linear**2 - 4 * quadratic * constant + square_root = jnp.sqrt(jnp.maximum(discriminant, 0)) + denominator = 2 * quadratic + root_minus = (-linear - square_root) / denominator + root_plus = (-linear + square_root) / denominator + infinity = jnp.asarray(jnp.inf, dtype=quadratic.dtype) + root_minus = jnp.where((discriminant >= 0) & (root_minus > 0), root_minus, infinity) + root_plus = jnp.where((discriminant >= 0) & (root_plus > 0), root_plus, infinity) + quadratic_root = jnp.minimum(root_minus, root_plus) + linear_root = -constant / linear + linear_root = jnp.where(linear_root > 0, linear_root, infinity) + return jnp.where(jnp.abs(quadratic) <= tolerance, linear_root, quadratic_root) + + +def singularity_diagnostics( + solution: NearAxisSolution, + *, + angular_resolution: int = 256, + newton_iterations: int = 8, +) -> SingularityDiagnostics: + """Locate the first singularity of the quadratic near-axis map. + + A uniform angular scan enumerates all positive roots of the quadratic + Jacobian approximation. A vectorized Newton solve then refines the + simultaneous conditions ``det(X) = 0`` and ``d det(X) / d theta = 0``. + """ + + if solution.second_order is None: + raise ValueError("A second-order solution is required for singular-radius diagnostics.") + if not isinstance(angular_resolution, int) or angular_resolution < 8: + raise ValueError("angular_resolution must be an integer >= 8.") + if not isinstance(newton_iterations, int) or newton_iterations < 0: + raise ValueError("newton_iterations must be a nonnegative integer.") + + g0, g1c, g1s, g20, g2s, g2c = _determinant_coefficients(solution) + theta_grid = jnp.arange(angular_resolution, dtype=g0.dtype) * (2 * jnp.pi / angular_resolution) + sine = jnp.sin(theta_grid) + cosine = jnp.cos(theta_grid) + sine_2 = jnp.sin(2 * theta_grid) + cosine_2 = jnp.cos(2 * theta_grid) + linear = g1c[:, None] * cosine + g1s[:, None] * sine + quadratic = g20[:, None] + g2s[:, None] * sine_2 + g2c[:, None] * cosine_2 + root_grid = _positive_quadratic_root(g0[:, None], linear, quadratic) + minimum_indices = jnp.argmin(root_grid, axis=1) + initial_radius = jnp.take_along_axis(root_grid, minimum_indices[:, None], axis=1)[:, 0] + initial_theta = theta_grid[minimum_indices] + finite_seed = jnp.isfinite(initial_radius) + radius = jnp.where(finite_seed, initial_radius, 1) + theta = initial_theta + + def newton_step(_, state): + radius, theta = state + sine = jnp.sin(theta) + cosine = jnp.cos(theta) + sine_2 = jnp.sin(2 * theta) + cosine_2 = jnp.cos(2 * theta) + linear = g1c * cosine + g1s * sine + linear_prime = -g1c * sine + g1s * cosine + linear_second = -linear + quadratic = g20 + g2s * sine_2 + g2c * cosine_2 + quadratic_prime = 2 * g2s * cosine_2 - 2 * g2c * sine_2 + quadratic_second = -4 * (g2s * sine_2 + g2c * cosine_2) + residual_0 = g0 + radius * linear + radius**2 * quadratic + residual_1 = radius * linear_prime + radius**2 * quadratic_prime + jacobian_00 = linear + 2 * radius * quadratic + jacobian_01 = residual_1 + jacobian_10 = linear_prime + 2 * radius * quadratic_prime + jacobian_11 = radius * linear_second + radius**2 * quadratic_second + determinant = jacobian_00 * jacobian_11 - jacobian_01 * jacobian_10 + safe_determinant = jnp.where( + jnp.abs(determinant) > jnp.finfo(radius.dtype).eps, + determinant, + jnp.inf, + ) + delta_radius = (-residual_0 * jacobian_11 + jacobian_01 * residual_1) / safe_determinant + delta_theta = (-jacobian_00 * residual_1 + residual_0 * jacobian_10) / safe_determinant + candidate_radius = radius + delta_radius + radius = jnp.where(candidate_radius > 0, candidate_radius, radius / 2) + return radius, theta + delta_theta + + radius, theta = jax.lax.fori_loop( + 0, + newton_iterations, + newton_step, + (radius, theta), + ) + sine = jnp.sin(theta) + cosine = jnp.cos(theta) + sine_2 = jnp.sin(2 * theta) + cosine_2 = jnp.cos(2 * theta) + linear = g1c * cosine + g1s * sine + linear_prime = -g1c * sine + g1s * cosine + quadratic = g20 + g2s * sine_2 + g2c * cosine_2 + quadratic_prime = 2 * g2s * cosine_2 - 2 * g2c * sine_2 + residual_0 = g0 + radius * linear + radius**2 * quadratic + residual_1 = radius * linear_prime + radius**2 * quadratic_prime + residual_norm = jnp.sqrt(residual_0**2 + residual_1**2) + valid = finite_seed & jnp.isfinite(radius) & jnp.isfinite(residual_norm) + radius = jnp.where(valid, radius, jnp.inf) + residual_norm = jnp.where(valid, residual_norm, jnp.inf) + return SingularityDiagnostics( + r_singularity=jnp.min(radius), + r_singularity_vs_varphi=radius, + inv_r_singularity_vs_varphi=1 / radius, + theta_singularity_vs_varphi=jnp.mod(theta, 2 * jnp.pi), + residual_norm_vs_varphi=residual_norm, + maximum_residual_norm=jnp.max(jnp.where(valid, residual_norm, 0)), + g0=g0, + g1c=g1c, + g1s=g1s, + g20=g20, + g2s=g2s, + g2c=g2c, + angular_resolution=angular_resolution, + newton_iterations=newton_iterations, + ) diff --git a/src/pyqsc_jax/solvers.py b/src/pyqsc_jax/solvers.py new file mode 100644 index 0000000..5fbfd09 --- /dev/null +++ b/src/pyqsc_jax/solvers.py @@ -0,0 +1,248 @@ +"""pyQSC_JAX-specific nonlinear solver policy and reports.""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import jax +import jax.numpy as jnp +from solvax import linear_solve, root_solve + +from pyqsc_jax.models import LinearSolveReport, RootSolveReport + +ArrayLike = Any + + +@dataclass(frozen=True) +class RootSolveOptions: + """Tolerance and globalization policy for dense Newton root solves.""" + + atol: float = 1e-13 + rtol: float = 1e-13 + step_tolerance: float = 1e-13 + max_steps: int = 20 + max_backtracking_steps: int = 12 + + def __post_init__(self) -> None: + if self.atol < 0 or self.rtol < 0 or self.step_tolerance < 0: + raise ValueError("Root tolerances must be nonnegative.") + if self.max_steps < 0 or self.max_backtracking_steps < 0: + raise ValueError("Root iteration limits must be nonnegative.") + + +DEFAULT_ROOT_OPTIONS = RootSolveOptions() + + +def _infinity_norm(value: jax.Array) -> jax.Array: + return jnp.max(jnp.abs(value)) + + +def dense_newton_root( + residual_function: Callable[[jax.Array], jax.Array], + initial_guess: ArrayLike, + *, + options: RootSolveOptions = DEFAULT_ROOT_OPTIONS, +) -> tuple[jax.Array, RootSolveReport]: + """Solve a small dense nonlinear system with damped Newton updates.""" + + initial_guess = jnp.asarray(initial_guess) + initial_residual = residual_function(initial_guess) + initial_residual_norm = _infinity_norm(initial_residual) + tolerance = jnp.maximum(options.atol, options.rtol * initial_residual_norm) + dtype = initial_guess.dtype + + initial_state = ( + initial_guess, + initial_residual, + initial_residual_norm, + jnp.asarray(jnp.inf, dtype=dtype), + jnp.int32(0), + jnp.int32(0), + ) + + def continue_iteration(state): + x, residual, residual_norm, step_norm, iterations, _ = state + finite = ( + jnp.all(jnp.isfinite(x)) & jnp.all(jnp.isfinite(residual)) & jnp.isfinite(residual_norm) + ) + step_threshold = options.step_tolerance * (1 + _infinity_norm(x)) + return ( + (residual_norm > tolerance) + & (iterations < options.max_steps) + & (step_norm > step_threshold) + & finite + ) + + def newton_step(state): + x, residual, residual_norm, _, iterations, total_backtracking = state + jacobian = jax.jacfwd(residual_function)(x) + if initial_guess.ndim == 0: + full_step = -residual / jacobian + else: + full_step = jnp.linalg.solve(jacobian, -residual) + + def candidate(damping): + candidate_x = x + damping * full_step + candidate_residual = residual_function(candidate_x) + candidate_norm = _infinity_norm(candidate_residual) + return candidate_x, candidate_residual, candidate_norm + + damping0 = jnp.asarray(1.0, dtype=dtype) + candidate_x0, candidate_residual0, candidate_norm0 = candidate(damping0) + line_state0 = ( + damping0, + candidate_x0, + candidate_residual0, + candidate_norm0, + jnp.int32(0), + ) + + def continue_backtracking(line_state): + _, _, candidate_residual, candidate_norm, backtracking = line_state + candidate_finite = jnp.all(jnp.isfinite(candidate_residual)) & jnp.isfinite( + candidate_norm + ) + return ((candidate_norm >= residual_norm) | ~candidate_finite) & ( + backtracking < options.max_backtracking_steps + ) + + def backtrack(line_state): + damping, _, _, _, backtracking = line_state + damping = 0.5 * damping + candidate_x, candidate_residual, candidate_norm = candidate(damping) + return ( + damping, + candidate_x, + candidate_residual, + candidate_norm, + backtracking + 1, + ) + + damping, candidate_x, candidate_residual, candidate_norm, backtracking = jax.lax.while_loop( + continue_backtracking, backtrack, line_state0 + ) + step_norm = _infinity_norm(damping * full_step) + return ( + candidate_x, + candidate_residual, + candidate_norm, + step_norm, + iterations + 1, + total_backtracking + backtracking, + ) + + x, residual, residual_norm, step_norm, iterations, backtracking_steps = jax.lax.while_loop( + continue_iteration, newton_step, initial_state + ) + finite = ( + jnp.all(jnp.isfinite(x)) & jnp.all(jnp.isfinite(residual)) & jnp.isfinite(residual_norm) + ) + converged = finite & (residual_norm <= tolerance) + step_threshold = options.step_tolerance * (1 + _infinity_norm(x)) + stagnated = finite & ~converged & (step_norm <= step_threshold) + final_jacobian = jax.jacfwd(residual_function)(x) + if initial_guess.ndim == 0: + jacobian_condition_number = jnp.where(final_jacobian == 0, jnp.inf, 1.0) + else: + jacobian_condition_number = jnp.linalg.cond(final_jacobian) + report = RootSolveReport( + initial_residual_norm=initial_residual_norm, + residual_norm=residual_norm, + tolerance=tolerance, + step_norm=step_norm, + iterations=iterations, + backtracking_steps=backtracking_steps, + jacobian_condition_number=jacobian_condition_number, + converged=converged, + finite=finite, + stagnated=stagnated, + ) + return x, report + + +def implicit_dense_root( + residual_function: Callable[[jax.Array], jax.Array], + initial_guess: ArrayLike, + *, + options: RootSolveOptions = DEFAULT_ROOT_OPTIONS, +) -> tuple[jax.Array, RootSolveReport]: + """Solve once with dense Newton and differentiate the converged equation. + + The Newton candidate is stopped before it is supplied to SOLVAX's custom + root. Consequently JVPs and VJPs use the implicit function theorem rather + than differentiating the iteration history. + """ + + candidate, report = dense_newton_root( + residual_function, + initial_guess, + options=options, + ) + candidate = jax.lax.stop_gradient(candidate) + root = root_solve( + residual_function, + candidate, + lambda _function, supplied_candidate: supplied_candidate, + ) + return root, report + + +def implicit_dense_linear_solve( + matrix: ArrayLike, + right_hand_side: ArrayLike, + *, + residual_tolerance: float = 1e-11, + condition_limit: float = 1e12, +) -> tuple[jax.Array, LinearSolveReport]: + """Solve a dense system with implicit JVP/VJP rules and diagnostics.""" + + if residual_tolerance < 0: + raise ValueError("residual_tolerance must be nonnegative.") + if condition_limit <= 0: + raise ValueError("condition_limit must be positive.") + + matrix = jnp.asarray(matrix) + right_hand_side = jnp.asarray(right_hand_side) + if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]: + raise ValueError("matrix must be square.") + if right_hand_side.ndim != 1 or right_hand_side.shape[0] != matrix.shape[0]: + raise ValueError("right_hand_side must match the matrix dimension.") + + matvec = lambda value: matrix @ value # noqa: E731 + transpose_matvec = lambda value: matrix.T @ value # noqa: E731 + primal_solver = lambda _operator, value: jnp.linalg.solve(matrix, value) # noqa: E731 + transpose_solver = lambda _operator, value: jnp.linalg.solve( # noqa: E731 + matrix.T, + value, + ) + solution = linear_solve( + matvec, + right_hand_side, + primal_solver, + transpose_matvec=transpose_matvec, + transpose_solver=transpose_solver, + ) + residual = matrix @ solution - right_hand_side + residual_norm = _infinity_norm(residual) + right_hand_side_norm = _infinity_norm(right_hand_side) + relative_residual_norm = residual_norm / jnp.maximum( + right_hand_side_norm, + jnp.finfo(right_hand_side.dtype).tiny, + ) + condition_number = jnp.linalg.cond(matrix) + finite = ( + jnp.all(jnp.isfinite(solution)) + & jnp.isfinite(residual_norm) + & jnp.isfinite(condition_number) + ) + converged = finite & (relative_residual_norm <= residual_tolerance) + well_conditioned = finite & (condition_number <= condition_limit) + report = LinearSolveReport( + residual_norm=residual_norm, + relative_residual_norm=relative_residual_norm, + matrix_condition_number=condition_number, + finite=finite, + converged=converged, + well_conditioned=well_conditioned, + ) + return solution, report diff --git a/src/pyqsc_jax/spectral.py b/src/pyqsc_jax/spectral.py new file mode 100644 index 0000000..dd63fdd --- /dev/null +++ b/src/pyqsc_jax/spectral.py @@ -0,0 +1,86 @@ +"""Spectral operations on uniform periodic grids.""" + +from typing import Any + +import jax +import jax.numpy as jnp + +ArrayLike = Any + + +def periodic_grid(n: int, *, period: ArrayLike = 2 * jnp.pi) -> jax.Array: + """Return ``n`` uniform points on ``[0, period)``.""" + + if not isinstance(n, int) or isinstance(n, bool) or n < 2: + raise ValueError("A periodic grid requires an integer n >= 2.") + return jnp.arange(n) * (jnp.asarray(period) / n) + + +def differentiation_matrix(n: int, *, period: ArrayLike = 2 * jnp.pi) -> jax.Array: + """Return the first-derivative Fourier collocation matrix. + + The even and odd formulas follow the trigonometric differentiation + matrices of Weideman and Reddy. Rows act on values sampled by + :func:`periodic_grid`. + """ + + if not isinstance(n, int) or isinstance(n, bool) or n < 2: + raise ValueError("A differentiation matrix requires an integer n >= 2.") + + index = jnp.arange(n) + difference = index[:, None] - index[None, :] + angle = jnp.pi * difference / n + sign = jnp.where(jnp.mod(difference, 2) == 0, 1.0, -1.0) + off_diagonal = difference != 0 + if n % 2 == 0: + entries = 0.5 * sign / jnp.tan(angle) + else: + entries = 0.5 * sign / jnp.sin(angle) + dimensionless = jnp.where(off_diagonal, entries, 0.0) + return dimensionless * (2 * jnp.pi / jnp.asarray(period)) + + +def differentiate(values: ArrayLike, *, period: ArrayLike = 2 * jnp.pi) -> jax.Array: + """Differentiate values along their final periodic axis.""" + + values = jnp.asarray(values) + matrix = differentiation_matrix(values.shape[-1], period=period) + return jnp.einsum("jk,...k->...j", matrix, values) + + +def periodic_integral(values: ArrayLike, *, period: ArrayLike = 2 * jnp.pi) -> jax.Array: + """Integrate uniform periodic samples along their final axis.""" + + return jnp.asarray(period) * jnp.mean(jnp.asarray(values), axis=-1) + + +def fourier_coefficients(values: ArrayLike) -> tuple[jax.Array, jax.Array]: + """Return integer frequencies and complex Fourier coefficients.""" + + values = jnp.asarray(values) + n = values.shape[-1] + frequency = jnp.fft.fftfreq(n, d=1 / n) + coefficients = jnp.fft.fft(values, axis=-1) / n + return frequency, coefficients + + +def fourier_interpolate( + values: ArrayLike, x: ArrayLike, *, period: ArrayLike = 2 * jnp.pi +) -> jax.Array: + """Evaluate the trigonometric interpolant of one-dimensional real samples.""" + + values = jnp.asarray(values) + if values.ndim != 1: + raise ValueError("fourier_interpolate currently accepts one-dimensional samples.") + + frequency, coefficients = fourier_coefficients(values) + x = jnp.asarray(x) + angle = 2 * jnp.pi * x / jnp.asarray(period) + phase = jnp.exp(1j * angle[..., None] * frequency) + if values.size % 2 == 0: + nyquist = values.size // 2 + phase = phase.at[..., nyquist].set(jnp.cos(nyquist * angle)) + result = jnp.sum(coefficients * phase, axis=-1) + if jnp.issubdtype(values.dtype, jnp.floating): + return jnp.real(result) + return result diff --git a/src/pyqsc_jax/third_order.py b/src/pyqsc_jax/third_order.py new file mode 100644 index 0000000..0353480 --- /dev/null +++ b/src/pyqsc_jax/third_order.py @@ -0,0 +1,129 @@ +"""Third-order flux-constraint corrections. + +The concise relation is from Landreman & Sengupta (2019) and is cross-checked +against ``landreman/pyQSC:qsc/calculate_r3.py`` at the audited BSD-2-Clause +upstream commit recorded in the refactor baseline. +""" + +from __future__ import annotations + +from dataclasses import replace + +import jax.numpy as jnp + +from pyqsc_jax.models import NearAxisSolution, ThirdOrderData +from pyqsc_jax.second_order import MU0 + + +def solve_third_order(solution: NearAxisSolution) -> NearAxisSolution: + """Add the r3 corrections required for consistency through second order.""" + + r2 = solution.second_order + if r2 is None: + raise ValueError("A second-order solution is required before the third-order correction.") + inputs = solution.inputs + length_scale = solution.geometry.abs_G0_over_B0 + D = solution.geometry.d_d_varphi + d_X1c = D @ solution.X1c + d_Y1c = D @ solution.Y1c + first_order_norm = solution.X1c**2 + solution.Y1c**2 + solution.Y1s**2 + Q = ( + -inputs.spsi + * inputs.B0 + * length_scale + / (2 * solution.G0**2) + * (solution.iotaN * inputs.I2 + MU0 * inputs.p2 * solution.G0 / inputs.B0**2) + + 2 * (r2.X2c * r2.Y2s - r2.X2s * r2.Y2c) + + inputs.spsi + * inputs.B0 + / (2 * solution.G0) + * (length_scale * r2.X20 * solution.curvature - r2.d_Z20_d_varphi) + + inputs.I2 + / (4 * solution.G0) + * ( + -length_scale * solution.torsion * first_order_norm + + solution.Y1c * d_X1c + - solution.X1c * d_Y1c + ) + ) + sign_product = inputs.sG * inputs.spsi + coefficient = -Q / (2 * sign_product) + N_helicity = solution.iota - solution.iotaN + B0_correction = ( + -inputs.sG + * inputs.B0**2 + * (r2.G2 + inputs.I2 * N_helicity) + * length_scale + / (2 * solution.G0**2) + - inputs.sG * inputs.spsi * inputs.B0 * 2 * (r2.X2c * r2.Y2s - r2.X2s * r2.Y2c) + - inputs.sG + * inputs.B0**2 + / (2 * solution.G0) + * (length_scale * r2.X20 * solution.curvature - r2.d_Z20_d_varphi) + - inputs.sG + * inputs.spsi + * inputs.B0 + * inputs.I2 + / (4 * solution.G0) + * ( + -length_scale * solution.torsion * first_order_norm + + solution.Y1c * d_X1c + - solution.X1c * d_Y1c + ) + ) + + zeros = jnp.zeros_like(solution.X1c) + X3s1 = zeros + X3c1 = solution.X1c * coefficient + Y3s1 = solution.Y1s * coefficient + Y3c1 = solution.Y1c * coefficient + angle = -solution.helicity * inputs.axis.nfp * solution.varphi + + def untwist(sine_coefficient, cosine_coefficient, harmonic): + sine = jnp.sin(harmonic * angle) + cosine = jnp.cos(harmonic * angle) + return ( + sine_coefficient * cosine + cosine_coefficient * sine, + -sine_coefficient * sine + cosine_coefficient * cosine, + ) + + X3s1_untwisted, X3c1_untwisted = untwist(X3s1, X3c1, 1) + Y3s1_untwisted, Y3c1_untwisted = untwist(Y3s1, Y3c1, 1) + Z3s1_untwisted, Z3c1_untwisted = untwist(zeros, zeros, 1) + X3s3_untwisted, X3c3_untwisted = untwist(zeros, zeros, 3) + Y3s3_untwisted, Y3c3_untwisted = untwist(zeros, zeros, 3) + Z3s3_untwisted, Z3c3_untwisted = untwist(zeros, zeros, 3) + third_order = ThirdOrderData( + flux_constraint_coefficient=coefficient, + B0_order_a_squared_to_cancel=B0_correction, + flux_constraint_residual=jnp.max(jnp.abs(Q + 2 * sign_product * coefficient)), + consistency_error=jnp.max(jnp.abs(coefficient - B0_correction / (2 * inputs.B0))), + X3s1=X3s1, + X3c1=X3c1, + Y3s1=Y3s1, + Y3c1=Y3c1, + Z3s1=zeros, + Z3c1=zeros, + X3s3=zeros, + X3c3=zeros, + Y3s3=zeros, + Y3c3=zeros, + Z3s3=zeros, + Z3c3=zeros, + d_X3c1_d_varphi=D @ X3c1, + d_Y3s1_d_varphi=D @ Y3s1, + d_Y3c1_d_varphi=D @ Y3c1, + X3s1_untwisted=X3s1_untwisted, + X3c1_untwisted=X3c1_untwisted, + Y3s1_untwisted=Y3s1_untwisted, + Y3c1_untwisted=Y3c1_untwisted, + Z3s1_untwisted=Z3s1_untwisted, + Z3c1_untwisted=Z3c1_untwisted, + X3s3_untwisted=X3s3_untwisted, + X3c3_untwisted=X3c3_untwisted, + Y3s3_untwisted=Y3s3_untwisted, + Y3c3_untwisted=Y3c3_untwisted, + Z3s3_untwisted=Z3s3_untwisted, + Z3c3_untwisted=Z3c3_untwisted, + ) + return replace(solution, third_order=third_order) diff --git a/src/pyqsc_jax/vmec.py b/src/pyqsc_jax/vmec.py new file mode 100644 index 0000000..43c0b98 --- /dev/null +++ b/src/pyqsc_jax/vmec.py @@ -0,0 +1,583 @@ +"""Fast, diagnosed conversion of near-axis surfaces to VMEC INDATA files.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from functools import partial +from math import isfinite +from pathlib import Path +from time import perf_counter +from typing import Any + +import jax +import jax.numpy as jnp + +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.second_order import MU0 + +ArrayLike = Any + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class VmecBoundary: + """Uniform-phi surface samples and VMEC Fourier coefficients. + + Coefficient arrays have shape ``(2 * ntor + 1, mpol + 1)``. The first + index is ordered from toroidal mode ``-ntor`` through ``+ntor``. + """ + + R: jax.Array + Z: jax.Array + phi0: jax.Array + RBC: jax.Array + RBS: jax.Array + ZBC: jax.Array + ZBS: jax.Array + maximum_toroidal_angle_residual: jax.Array + toroidal_angle_tolerance: jax.Array + toroidal_angle_converged: jax.Array + maximum_R_reconstruction_error: jax.Array + maximum_Z_reconstruction_error: jax.Array + ntheta: int = field(metadata={"static": True}) + mpol: int = field(metadata={"static": True}) + ntor: int = field(metadata={"static": True}) + newton_iterations: int = field(metadata={"static": True}) + + +@dataclass(frozen=True) +class VmecInputParameters: + """Numerical controls written to a VMEC ``&INDATA`` namelist.""" + + delt: float = 0.9 + nstep: int = 200 + tcon0: float = 2.0 + ns_array: tuple[int, ...] = (15, 31, 61) + ftol_array: tuple[float, ...] = (1.0e-10, 1.0e-11, 1.0e-12) + niter_array: tuple[int, ...] = (2000, 3000, 5000) + + def __post_init__(self) -> None: + lengths = (len(self.ns_array), len(self.ftol_array), len(self.niter_array)) + if not self.ns_array or len(set(lengths)) != 1: + raise ValueError("ns_array, ftol_array, and niter_array must be nonempty and aligned.") + if self.delt <= 0 or self.nstep < 1 or self.tcon0 <= 0: + raise ValueError("VMEC damping, step interval, and constraint factor must be positive.") + if any(value < 3 for value in self.ns_array): + raise ValueError("Every VMEC radial resolution must be at least 3.") + if any(value <= 0 for value in self.ftol_array): + raise ValueError("Every VMEC force tolerance must be positive.") + if any(value < 1 for value in self.niter_array): + raise ValueError("Every VMEC iteration limit must be positive.") + + +@dataclass(frozen=True) +class VmecExport: + """Result of writing a diagnosed VMEC input file.""" + + path: Path + boundary: VmecBoundary + phiedge: float + curtor: float + pressure_axis: float + lasym: bool + conversion_seconds: float + + +def _periodic_interpolate( + query: jax.Array, + grid: jax.Array, + values: jax.Array, + period: jax.Array, +) -> jax.Array: + period = jnp.asarray(period) + wrapped = jnp.mod(query, period) + extended_grid = jnp.concatenate((grid, period[None])) + extended_values = jnp.concatenate((values, values[:1])) + return jnp.interp(wrapped, extended_grid, extended_values) + + +def frenet_displacements( + solution: NearAxisSolution, + radius: jax.Array, + theta: jax.Array, + phi0: jax.Array, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Evaluate all available near-axis Frenet displacements.""" + + period = 2 * jnp.pi / solution.inputs.axis.nfp + grid = solution.phi + interpolate = lambda values: _periodic_interpolate(phi0, grid, values, period) # noqa: E731 + cosine = jnp.cos(theta) + sine = jnp.sin(theta) + X = radius * ( + interpolate(solution.X1c_untwisted) * cosine + interpolate(solution.X1s_untwisted) * sine + ) + Y = radius * ( + interpolate(solution.Y1c_untwisted) * cosine + interpolate(solution.Y1s_untwisted) * sine + ) + Z = jnp.zeros_like(X) + if solution.second_order is not None: + cosine2 = jnp.cos(2 * theta) + sine2 = jnp.sin(2 * theta) + X = X + radius**2 * ( + interpolate(solution.X20_untwisted) + + interpolate(solution.X2c_untwisted) * cosine2 + + interpolate(solution.X2s_untwisted) * sine2 + ) + Y = Y + radius**2 * ( + interpolate(solution.Y20_untwisted) + + interpolate(solution.Y2c_untwisted) * cosine2 + + interpolate(solution.Y2s_untwisted) * sine2 + ) + Z = Z + radius**2 * ( + interpolate(solution.Z20_untwisted) + + interpolate(solution.Z2c_untwisted) * cosine2 + + interpolate(solution.Z2s_untwisted) * sine2 + ) + if solution.third_order is not None: + cosine3 = jnp.cos(3 * theta) + sine3 = jnp.sin(3 * theta) + X = X + radius**3 * ( + interpolate(solution.X3c1_untwisted) * cosine + + interpolate(solution.X3s1_untwisted) * sine + + interpolate(solution.X3c3_untwisted) * cosine3 + + interpolate(solution.X3s3_untwisted) * sine3 + ) + Y = Y + radius**3 * ( + interpolate(solution.Y3c1_untwisted) * cosine + + interpolate(solution.Y3s1_untwisted) * sine + + interpolate(solution.Y3c3_untwisted) * cosine3 + + interpolate(solution.Y3s3_untwisted) * sine3 + ) + Z = Z + radius**3 * ( + interpolate(solution.Z3c1_untwisted) * cosine + + interpolate(solution.Z3s1_untwisted) * sine + + interpolate(solution.Z3c3_untwisted) * cosine3 + + interpolate(solution.Z3s3_untwisted) * sine3 + ) + return X, Y, Z + + +def _surface_at_axis_angle( + solution: NearAxisSolution, + radius: jax.Array, + theta: jax.Array, + phi0: jax.Array, +) -> tuple[jax.Array, jax.Array, jax.Array]: + period = 2 * jnp.pi / solution.inputs.axis.nfp + grid = solution.phi + interpolate = lambda values: _periodic_interpolate(phi0, grid, values, period) # noqa: E731 + X, Y, Z = frenet_displacements(solution, radius, theta, phi0) + normal = solution.geometry.normal_cylindrical + binormal = solution.geometry.binormal_cylindrical + tangent = solution.geometry.tangent_cylindrical + delta_R = ( + X * interpolate(normal[:, 0]) + + Y * interpolate(binormal[:, 0]) + + Z * interpolate(tangent[:, 0]) + ) + delta_phi = ( + X * interpolate(normal[:, 1]) + + Y * interpolate(binormal[:, 1]) + + Z * interpolate(tangent[:, 1]) + ) + delta_Z = ( + X * interpolate(normal[:, 2]) + + Y * interpolate(binormal[:, 2]) + + Z * interpolate(tangent[:, 2]) + ) + axis_R = interpolate(solution.R0) + R = jnp.hypot(axis_R + delta_R, delta_phi) + cylindrical_phi = phi0 + jnp.arctan2(delta_phi, axis_R + delta_R) + cylindrical_Z = interpolate(solution.Z0) + delta_Z + return R, cylindrical_Z, cylindrical_phi + + +@partial( + jax.jit, + static_argnames=("ntheta", "newton_iterations"), +) +def uniform_cylindrical_surface( + solution: NearAxisSolution, + radius: ArrayLike, + *, + ntheta: int = 32, + newton_iterations: int = 6, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + """Map a near-axis boundary to a uniform cylindrical-toroidal grid.""" + + radius = jnp.asarray(radius) + theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False)[:, None] + target_phi = jnp.broadcast_to(solution.phi[None, :], (ntheta, solution.inputs.nphi)) + phi0 = target_phi + + def newton_step(_iteration, current_phi0): + angle_function = lambda value: _surface_at_axis_angle( # noqa: E731 + solution, + radius, + theta, + value, + )[2] + cylindrical_phi, derivative = jax.jvp( + angle_function, + (current_phi0,), + (jnp.ones_like(current_phi0),), + ) + derivative_floor = jnp.sqrt(jnp.finfo(current_phi0.dtype).eps) + safe_derivative = jnp.where( + jnp.abs(derivative) > derivative_floor, + derivative, + jnp.where(derivative >= 0, derivative_floor, -derivative_floor), + ) + return current_phi0 - (cylindrical_phi - target_phi) / safe_derivative + + phi0 = jax.lax.fori_loop(0, newton_iterations, newton_step, phi0) + R, Z, cylindrical_phi = _surface_at_axis_angle( + solution, + radius, + theta, + phi0, + ) + residual = jnp.max(jnp.abs(cylindrical_phi - target_phi)) + return R, Z, phi0, residual + + +def _fft_coefficients( + values: jax.Array, + mpol: int, + ntor: int, +) -> tuple[jax.Array, jax.Array]: + ntheta, nphi = values.shape + spectrum = jnp.fft.fft2(values) / (ntheta * nphi) + poloidal_modes = jnp.arange(mpol + 1) + toroidal_modes = jnp.arange(-ntor, ntor + 1) + selected = spectrum[ + poloidal_modes[None, :], + jnp.mod(-toroidal_modes[:, None], nphi), + ] + cosine = 2 * jnp.real(selected) + sine = -2 * jnp.imag(selected) + constant_index = ntor + cosine = cosine.at[constant_index, 0].set(jnp.real(spectrum[0, 0])) + sine = sine.at[constant_index, 0].set(0) + cosine = cosine.at[:ntor, 0].set(0) + sine = sine.at[:ntor, 0].set(0) + return cosine, sine + + +def _reconstruct_surface( + cosine_coefficients: jax.Array, + sine_coefficients: jax.Array, + *, + ntheta: int, + nphi: int, + nfp: int, + mpol: int, + ntor: int, +) -> jax.Array: + theta = jnp.linspace(0, 2 * jnp.pi, ntheta, endpoint=False) + phi = jnp.linspace(0, 2 * jnp.pi / nfp, nphi, endpoint=False) + phi2d, theta2d = jnp.meshgrid(phi, theta, indexing="xy") + poloidal_modes = jnp.arange(mpol + 1) + toroidal_modes = jnp.arange(-ntor, ntor + 1) + angle = ( + poloidal_modes[None, :, None, None] * theta2d[None, None, :, :] + - toroidal_modes[:, None, None, None] * nfp * phi2d[None, None, :, :] + ) + return jnp.sum( + cosine_coefficients[:, :, None, None] * jnp.cos(angle) + + sine_coefficients[:, :, None, None] * jnp.sin(angle), + axis=(0, 1), + ) + + +@partial( + jax.jit, + static_argnames=("ntheta", "mpol", "ntor", "newton_iterations"), +) +def vmec_boundary( + solution: NearAxisSolution, + radius: ArrayLike, + *, + ntheta: int = 40, + mpol: int = 12, + ntor: int = 14, + newton_iterations: int = 6, + toroidal_angle_tolerance: float = 0.0, +) -> VmecBoundary: + """Return a uniformly sampled, FFT-projected VMEC boundary. + + A zero ``toroidal_angle_tolerance`` selects 100 machine epsilons + for the active JAX dtype. The returned convergence flag is data, so this + function remains JIT-compatible; :func:`to_vmec` turns a false flag into + a hard failure before writing an input file. + """ + + R, Z, phi0, angle_residual = uniform_cylindrical_surface( + solution, + radius, + ntheta=ntheta, + newton_iterations=newton_iterations, + ) + RBC, RBS = _fft_coefficients(R, mpol, ntor) + ZBC, ZBS = _fft_coefficients(Z, mpol, ntor) + reconstructed_R = _reconstruct_surface( + RBC, + RBS, + ntheta=ntheta, + nphi=solution.inputs.nphi, + nfp=solution.inputs.axis.nfp, + mpol=mpol, + ntor=ntor, + ) + reconstructed_Z = _reconstruct_surface( + ZBC, + ZBS, + ntheta=ntheta, + nphi=solution.inputs.nphi, + nfp=solution.inputs.axis.nfp, + mpol=mpol, + ntor=ntor, + ) + requested_tolerance = jnp.asarray(toroidal_angle_tolerance, dtype=R.dtype) + automatic_tolerance = 100 * jnp.finfo(R.dtype).eps + angle_tolerance = jnp.where( + requested_tolerance == 0, + automatic_tolerance, + requested_tolerance, + ) + return VmecBoundary( + R=R, + Z=Z, + phi0=phi0, + RBC=RBC, + RBS=RBS, + ZBC=ZBC, + ZBS=ZBS, + maximum_toroidal_angle_residual=angle_residual, + toroidal_angle_tolerance=angle_tolerance, + toroidal_angle_converged=angle_residual <= angle_tolerance, + maximum_R_reconstruction_error=jnp.max(jnp.abs(reconstructed_R - R)), + maximum_Z_reconstruction_error=jnp.max(jnp.abs(reconstructed_Z - Z)), + ntheta=ntheta, + mpol=mpol, + ntor=ntor, + newton_iterations=newton_iterations, + ) + + +def _validated_resolution( + solution: NearAxisSolution, + *, + ntheta: int, + mpol: int, + ntor: int, + newton_iterations: int, +) -> None: + integers = { + "ntheta": ntheta, + "mpol": mpol, + "ntor": ntor, + "newton_iterations": newton_iterations, + } + if any(not isinstance(value, int) or isinstance(value, bool) for value in integers.values()): + raise ValueError("VMEC conversion resolutions must be integers.") + if ntheta < 2 * (mpol + 1): + raise ValueError("ntheta must be at least 2 * (mpol + 1).") + if mpol < 1 or ntor < 0 or newton_iterations < 1: + raise ValueError("mpol and Newton iterations must be positive; ntor must be nonnegative.") + if 2 * ntor + 1 > solution.inputs.nphi: + raise ValueError("The solution nphi must be at least 2 * ntor + 1.") + + +def _input_parameters( + parameters: VmecInputParameters | Mapping[str, Any] | None, +) -> VmecInputParameters: + if parameters is None: + return VmecInputParameters() + if isinstance(parameters, VmecInputParameters): + return parameters + allowed = set(VmecInputParameters.__dataclass_fields__) + unknown = set(parameters) - allowed + if unknown: + names = ", ".join(sorted(unknown)) + raise ValueError(f"Unknown VMEC input parameter(s): {names}.") + converted = dict(parameters) + for name in ("ns_array", "ftol_array", "niter_array"): + if name in converted: + converted[name] = tuple(converted[name]) + return VmecInputParameters(**converted) + + +def _format_sequence(values: tuple[Any, ...] | jax.Array) -> str: + return ", ".join(f"{float(value):.16e}" for value in values) + + +def _format_integer_sequence(values: tuple[int, ...]) -> str: + return ", ".join(str(int(value)) for value in values) + + +def _coefficient_lines( + boundary: VmecBoundary, + *, + lasym: bool, + coefficient_tolerance: float, +) -> list[str]: + arrays = tuple( + jax.device_get(value) for value in (boundary.RBC, boundary.RBS, boundary.ZBC, boundary.ZBS) + ) + RBC, RBS, ZBC, ZBS = arrays + lines: list[str] = [] + for m in range(boundary.mpol + 1): + for n in range(-boundary.ntor, boundary.ntor + 1): + index = n + boundary.ntor + symmetric_nonzero = ( + abs(RBC[index, m]) > coefficient_tolerance + or abs(ZBS[index, m]) > coefficient_tolerance + ) + asymmetric_nonzero = lasym and ( + abs(RBS[index, m]) > coefficient_tolerance + or abs(ZBC[index, m]) > coefficient_tolerance + ) + if symmetric_nonzero or asymmetric_nonzero: + lines.append( + f" RBC({n:03d},{m:03d}) = {RBC[index, m]:+.16e}," + f" ZBS({n:03d},{m:03d}) = {ZBS[index, m]:+.16e}" + ) + if lasym: + lines.append( + f" RBS({n:03d},{m:03d}) = {RBS[index, m]:+.16e}," + f" ZBC({n:03d},{m:03d}) = {ZBC[index, m]:+.16e}" + ) + return lines + + +def to_vmec( + solution: NearAxisSolution, + filename: str | Path, + *, + r: float = 0.1, + parameters: VmecInputParameters | Mapping[str, Any] | None = None, + ntheta: int = 40, + mpol: int = 12, + ntor: int = 14, + ntor_max: int = 14, + newton_iterations: int = 6, + toroidal_angle_tolerance: float = 0.0, + coefficient_tolerance: float = 1.0e-14, +) -> VmecExport: + """Write a VMEC fixed-boundary input and return conversion diagnostics.""" + + radius = float(r) + if not isfinite(radius) or radius <= 0: + raise ValueError("r must be positive and finite.") + if not isinstance(ntor_max, int) or isinstance(ntor_max, bool) or ntor_max < 0: + raise ValueError("ntor_max must be a nonnegative integer.") + if not isfinite(toroidal_angle_tolerance) or toroidal_angle_tolerance < 0: + raise ValueError("toroidal_angle_tolerance must be nonnegative and finite.") + if not isfinite(coefficient_tolerance) or coefficient_tolerance < 0: + raise ValueError("coefficient_tolerance must be nonnegative and finite.") + effective_ntor = min(ntor, ntor_max) + _validated_resolution( + solution, + ntheta=ntheta, + mpol=mpol, + ntor=effective_ntor, + newton_iterations=newton_iterations, + ) + controls = _input_parameters(parameters) + + start = perf_counter() + boundary = vmec_boundary( + solution, + radius, + ntheta=ntheta, + mpol=mpol, + ntor=effective_ntor, + newton_iterations=newton_iterations, + toroidal_angle_tolerance=toroidal_angle_tolerance, + ) + jax.tree.map(lambda value: value.block_until_ready(), boundary) + conversion_seconds = perf_counter() - start + if not bool(boundary.toroidal_angle_converged): + residual = float(boundary.maximum_toroidal_angle_residual) + tolerance = float(boundary.toroidal_angle_tolerance) + raise RuntimeError( + "The near-axis surface could not be represented on a uniform " + "cylindrical-toroidal grid: the maximum angle residual " + f"{residual:.6e} exceeds {tolerance:.6e}. Reduce r or increase " + "newton_iterations; no VMEC input was written." + ) + + asymmetric_amplitude = max( + float(jnp.max(jnp.abs(boundary.RBS))), + float(jnp.max(jnp.abs(boundary.ZBC))), + ) + lasym = asymmetric_amplitude > coefficient_tolerance + inputs = solution.inputs + axis = inputs.axis + phiedge = float(jnp.pi * radius**2 * inputs.B0) + curtor = float(2 * jnp.pi * inputs.I2 * radius**2 / MU0) + pressure_axis = float(-inputs.p2 * radius**2) + lines = [ + "! Generated deterministically by pyQSC_JAX.", + f"! Near-axis radius r = {radius:.16e}; etabar = {float(inputs.etabar):.16e}.", + ( + f"! nphi = {inputs.nphi}; order = r{inputs.order}; ntheta = {ntheta};" + f" mpol = {mpol}; ntor = {effective_ntor}." + ), + ( + "! Conversion diagnostics:" + f" max_phi_residual = {float(boundary.maximum_toroidal_angle_residual):.6e};" + f" phi_tolerance = {float(boundary.toroidal_angle_tolerance):.6e};" + " phi_converged = true;" + f" max_R_error = {float(boundary.maximum_R_reconstruction_error):.6e};" + f" max_Z_error = {float(boundary.maximum_Z_reconstruction_error):.6e}." + ), + "&INDATA", + f" DELT = {controls.delt:.16e}", + f" NSTEP = {controls.nstep}", + f" TCON0 = {controls.tcon0:.16e}", + f" NS_ARRAY = {_format_integer_sequence(controls.ns_array)}", + f" FTOL_ARRAY = {_format_sequence(controls.ftol_array)}", + f" NITER_ARRAY = {_format_integer_sequence(controls.niter_array)}", + f" LASYM = {'T' if lasym else 'F'}", + " LFREEB = F", + f" NFP = {axis.nfp}", + f" MPOL = {mpol}", + f" NTOR = {effective_ntor}", + f" PHIEDGE = {phiedge:.16e}", + " PRES_SCALE = 1.0000000000000000e+00", + " PMASS_TYPE = 'power_series'", + f" AM = {pressure_axis:.16e}, {-pressure_axis:.16e}", + f" CURTOR = {curtor:.16e}", + " NCURR = 1", + " PCURR_TYPE = 'power_series'", + " AC = 1.0000000000000000e+00", + f" RAXIS_CC = {_format_sequence(axis.rc)}", + f" RAXIS_CS = {_format_sequence(-axis.rs)}", + f" ZAXIS_CC = {_format_sequence(axis.zc)}", + f" ZAXIS_CS = {_format_sequence(-axis.zs)}", + "! Boundary coefficients", + ] + lines.extend( + _coefficient_lines( + boundary, + lasym=lasym, + coefficient_tolerance=coefficient_tolerance, + ) + ) + lines.append("/") + output_path = Path(filename) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return VmecExport( + path=output_path, + boundary=boundary, + phiedge=phiedge, + curtor=curtor, + pressure_axis=pressure_axis, + lasym=lasym, + conversion_seconds=conversion_seconds, + ) diff --git a/src/pyqsc_jax/vmex.py b/src/pyqsc_jax/vmex.py new file mode 100644 index 0000000..99ca6ef --- /dev/null +++ b/src/pyqsc_jax/vmex.py @@ -0,0 +1,445 @@ +"""Optional differentiable fixed-boundary equilibrium interface to VMEX.""" + +from __future__ import annotations + +import dataclasses +import importlib +from dataclasses import dataclass +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np + +from pyqsc_jax.models import NearAxisSolution +from pyqsc_jax.second_order import MU0 +from pyqsc_jax.vmec import VmecBoundary, _validated_resolution, vmec_boundary + +VMEX_VALIDATED_COMMIT = "2a40d7566be083070ea3ea534fa5d1fc44ad733a" + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class VmexRadialQuantities: + """Differentiable radial and scalar quantities from a converged VMEX state. + + ``iota_vmec`` retains VMEC's native sign convention. ``iota`` uses the + pyQSC_JAX toroidal-angle convention, which is the negative of VMEC's for + boundaries produced by :func:`pyqsc_jax.vmec_boundary`. + """ + + s: jax.Array + iota: jax.Array + iota_vmec: jax.Array + qs_surfaces: jax.Array + quasisymmetry: jax.Array + magnetic_well: jax.Array + aspect: jax.Array + volume: jax.Array + magnetic_energy: jax.Array + thermal_energy: jax.Array + + +@dataclass(frozen=True) +class VmexEquilibrium: + """A VMEX implicit solution together with user-facing radial quantities.""" + + solution: Any + quantities: VmexRadialQuantities + + +@dataclass(frozen=True) +class VmexProblem: + """Reusable static VMEX problem and its differentiable parameter pytree.""" + + input: Any + parameters: Any + boundary: VmecBoundary + radius: float + qs_surfaces: tuple[float, ...] + helicity_m: int + helicity_n: int + ntheta: int + mpol: int + ntor: int + newton_iterations: int + toroidal_angle_tolerance: float + ftol: float + max_iterations: int + adjoint_tol: float + multigrid: bool + device: Any + vmex_version: str + validated_commit: str = VMEX_VALIDATED_COMMIT + + @property + def finite_beta(self) -> bool: + """Whether the static reference input has a nonzero pressure profile.""" + + return bool(np.any(np.asarray(self.input.am))) + + def parameters_for( + self, + solution: NearAxisSolution, + *, + radius: Any | None = None, + ) -> Any: + """Map another near-axis solution into this problem's parameter pytree.""" + + return vmex_parameters_from_solution(self, solution, radius=radius) + + def quantities(self, parameters: Any | None = None) -> VmexRadialQuantities: + """Solve and return differentiable radial quantities.""" + + return vmex_radial_quantities(self, parameters) + + def solve(self, parameters: Any | None = None) -> VmexEquilibrium: + """Solve the fixed-boundary equilibrium and retain the raw VMEX result.""" + + return solve_vmex(self, parameters) + + +def _import_vmex(): + try: + module = importlib.import_module("vmex") + except ImportError as error: + raise ImportError( + "The differentiable equilibrium interface requires VMEX. Install " + "pyqsc-jax with the 'vmex' extra, or install the current source with " + "'python -m pip install git+https://github.com/uwplasma/vmex.git'." + ) from error + missing = tuple( + name for name in ("VmecInput", "implicit", "optimize") if not hasattr(module, name) + ) + if missing: + names = ", ".join(missing) + raise ImportError(f"The installed VMEX does not provide the required public API: {names}.") + return module + + +def _validated_surfaces(surfaces: Any) -> tuple[float, ...]: + values = tuple(float(value) for value in surfaces) + if any(not np.isfinite(value) or value <= 0 or value > 1 for value in values): + raise ValueError("Every VMEX quasisymmetry surface must satisfy 0 < s <= 1.") + if any(right <= left for left, right in zip(values, values[1:], strict=False)): + raise ValueError("VMEX quasisymmetry surfaces must be strictly increasing.") + return values + + +def _validated_radial_controls( + ns_array: Any, + ftol_array: Any | None, + *, + ftol: float, + max_iterations: int, +) -> tuple[tuple[int, ...], tuple[float, ...], tuple[int, ...]]: + ns = tuple(int(value) for value in ns_array) + if not ns or any(value < 3 for value in ns): + raise ValueError("ns_array must be nonempty and every radial resolution must be >= 3.") + if any(right <= left for left, right in zip(ns, ns[1:], strict=False)): + raise ValueError("ns_array must be strictly increasing.") + if not np.isfinite(ftol) or ftol <= 0: + raise ValueError("ftol must be positive and finite.") + if ( + not isinstance(max_iterations, int) + or isinstance(max_iterations, bool) + or max_iterations < 1 + ): + raise ValueError("max_iterations must be a positive integer.") + if ftol_array is None: + if len(ns) == 1: + tolerances = (float(ftol),) + else: + start = max(float(ftol), 1.0e-8) + tolerances = tuple(float(value) for value in np.geomspace(start, ftol, len(ns))) + else: + tolerances = tuple(float(value) for value in ftol_array) + if len(tolerances) != len(ns) or any( + not np.isfinite(value) or value <= 0 for value in tolerances + ): + raise ValueError("ftol_array must contain one positive finite value per ns_array stage.") + return ns, tolerances, (int(max_iterations),) * len(ns) + + +def _is_asymmetric(boundary: VmecBoundary, tolerance: float = 1.0e-13) -> bool: + return ( + max( + float(jnp.max(jnp.abs(boundary.RBS))), + float(jnp.max(jnp.abs(boundary.ZBC))), + ) + > tolerance + ) + + +def _require_converged_boundary(boundary: VmecBoundary) -> None: + """Reject concrete failures while remaining usable inside JAX tracing.""" + + try: + converged = bool(boundary.toroidal_angle_converged) + except jax.errors.TracerBoolConversionError: + return + if not converged: + raise RuntimeError( + "The VMEX boundary angle inversion did not converge. Reduce r or " + "increase newton_iterations before solving the radial equilibrium." + ) + + +def _profile_arrays(parameters: Any, solution: NearAxisSolution, radius: Any): + radius = jnp.asarray(radius) + pressure_axis = -solution.inputs.p2 * radius**2 + am = jnp.zeros_like(parameters.am) + am = am.at[0].set(pressure_axis) + if am.shape[0] > 1: + am = am.at[1].set(-pressure_axis) + ac = jnp.zeros_like(parameters.ac).at[0].set(1.0) + phiedge = jnp.pi * radius**2 * solution.inputs.B0 + curtor = 2 * jnp.pi * solution.inputs.I2 * radius**2 / MU0 + return am, ac, phiedge, curtor + + +def vmex_parameters_from_solution( + problem: VmexProblem, + solution: NearAxisSolution, + *, + radius: Any | None = None, +) -> Any: + """Traceably map a near-axis boundary and profiles to VMEX parameters. + + The problem fixes discrete resolution, topology, and solver controls. + Boundary coefficients, toroidal flux, pressure, and enclosed current + remain JAX values, so gradients can propagate from a newly constructed + :class:`NearAxisSolution` through VMEX's converged fixed point. + """ + + if solution.inputs.axis.nfp != problem.input.nfp: + raise ValueError("The new solution must have the problem's number of field periods.") + selected_radius = problem.radius if radius is None else radius + boundary = vmec_boundary( + solution, + selected_radius, + ntheta=problem.ntheta, + mpol=problem.mpol, + ntor=problem.ntor, + newton_iterations=problem.newton_iterations, + toroidal_angle_tolerance=problem.toroidal_angle_tolerance, + ) + _require_converged_boundary(boundary) + boundary_mask = boundary.toroidal_angle_converged + rbc = jnp.where(boundary_mask, boundary.RBC, jnp.nan) + rbs = jnp.where(boundary_mask, boundary.RBS, jnp.nan) + zbc = jnp.where(boundary_mask, boundary.ZBC, jnp.nan) + zbs = jnp.where(boundary_mask, boundary.ZBS, jnp.nan) + am, ac, phiedge, curtor = _profile_arrays(problem.parameters, solution, selected_radius) + return dataclasses.replace( + problem.parameters, + rbc=rbc, + rbs=rbs, + zbc=zbc, + zbs=zbs, + phiedge=phiedge, + curtor=curtor, + pres_scale=jnp.asarray(1.0, dtype=phiedge.dtype), + am=am, + ac=ac, + ) + + +def to_vmex_problem( + solution: NearAxisSolution, + *, + r: float = 0.03, + qs_surfaces: Any = (0.25, 0.5, 0.75, 1.0), + helicity_m: int = 1, + helicity_n: int | None = None, + ntheta: int = 24, + mpol: int = 6, + ntor: int = 6, + newton_iterations: int = 6, + toroidal_angle_tolerance: float = 0.0, + ns_array: Any = (15, 31), + ftol_array: Any | None = None, + ftol: float = 1.0e-10, + max_iterations: int = 5000, + adjoint_tol: float = 1.0e-11, + multigrid: bool = True, + device: Any = None, +) -> VmexProblem: + """Create a differentiable fixed-boundary VMEX problem without disk I/O. + + ``mpol`` is the maximum retained poloidal Fourier index in the + pyQSC_JAX conversion. VMEX therefore receives ``MPOL = mpol + 1``. + Pressure and current use the same near-axis-consistent profiles as + :func:`pyqsc_jax.to_vmec`. + """ + + vmex = _import_vmex() + radius = float(r) + if not np.isfinite(radius) or radius <= 0: + raise ValueError("r must be positive and finite.") + surfaces = _validated_surfaces(qs_surfaces) + ns, tolerances, iteration_limits = _validated_radial_controls( + ns_array, + ftol_array, + ftol=ftol, + max_iterations=max_iterations, + ) + if not isinstance(helicity_m, int) or isinstance(helicity_m, bool): + raise ValueError("helicity_m must be an integer.") + if helicity_n is None: + helicity_n = -int(np.asarray(solution.helicity)) + if not isinstance(helicity_n, int) or isinstance(helicity_n, bool): + raise ValueError("helicity_n must be an integer.") + if not np.isfinite(adjoint_tol) or adjoint_tol <= 0: + raise ValueError("adjoint_tol must be positive and finite.") + if not np.isfinite(toroidal_angle_tolerance) or toroidal_angle_tolerance < 0: + raise ValueError("toroidal_angle_tolerance must be nonnegative and finite.") + _validated_resolution( + solution, + ntheta=ntheta, + mpol=mpol, + ntor=ntor, + newton_iterations=newton_iterations, + ) + + boundary = vmec_boundary( + solution, + radius, + ntheta=ntheta, + mpol=mpol, + ntor=ntor, + newton_iterations=newton_iterations, + toroidal_angle_tolerance=toroidal_angle_tolerance, + ) + _require_converged_boundary(boundary) + lasym = _is_asymmetric(boundary) + if lasym and surfaces: + raise NotImplementedError( + "VMEX's traceable quasisymmetry profile currently supports " + "stellarator-symmetric equilibria only; pass qs_surfaces=() to " + "solve an asymmetric equilibrium without that diagnostic." + ) + + inputs = solution.inputs + axis = inputs.axis + arrays = tuple( + np.asarray(jax.device_get(value)) + for value in (boundary.RBC, boundary.RBS, boundary.ZBC, boundary.ZBS) + ) + rbc, rbs, zbc, zbs = arrays + pressure_axis = -float(inputs.p2) * radius**2 + am = np.zeros(21) + am[:2] = (pressure_axis, -pressure_axis) + ac = np.zeros(21) + ac[0] = 1.0 + vmex_input = vmex.VmecInput( + lasym=lasym, + lfreeb=False, + nfp=axis.nfp, + mpol=boundary.RBC.shape[1], + ntor=boundary.ntor, + ns_array=np.asarray(ns), + ftol_array=np.asarray(tolerances), + niter_array=np.asarray(iteration_limits), + delt=0.9, + tcon0=2.0, + phiedge=np.pi * radius**2 * float(inputs.B0), + pres_scale=1.0, + am=am, + ncurr=1, + ac=ac, + curtor=2 * np.pi * float(inputs.I2) * radius**2 / MU0, + raxis_c=np.asarray(axis.rc), + raxis_s=-np.asarray(axis.rs), + zaxis_c=np.asarray(axis.zc), + zaxis_s=-np.asarray(axis.zs), + rbc=rbc, + rbs=rbs, + zbc=zbc, + zbs=zbs, + ) + parameters = vmex.implicit.params_from_input(vmex_input, device=device) + problem = VmexProblem( + input=vmex_input, + parameters=parameters, + boundary=boundary, + radius=radius, + qs_surfaces=surfaces, + helicity_m=helicity_m, + helicity_n=helicity_n, + ntheta=ntheta, + mpol=mpol, + ntor=ntor, + newton_iterations=newton_iterations, + toroidal_angle_tolerance=float(toroidal_angle_tolerance), + ftol=float(ftol), + max_iterations=max_iterations, + adjoint_tol=float(adjoint_tol), + multigrid=bool(multigrid), + device=device, + vmex_version=str(vmex.__version__), + ) + return dataclasses.replace( + problem, + parameters=vmex_parameters_from_solution(problem, solution), + ) + + +def _radial_quantities(problem: VmexProblem, vmex_solution: Any) -> VmexRadialQuantities: + vmex = _import_vmex() + runtime = vmex_solution.runtime + if runtime is None: + raise RuntimeError("VMEX did not retain the runtime required for radial diagnostics.") + iota_vmec = vmex.implicit.iota_profile(vmex_solution.state, runtime) + s = jnp.linspace(0.0, 1.0, iota_vmec.shape[0], dtype=iota_vmec.dtype) + surfaces = jnp.asarray(problem.qs_surfaces, dtype=iota_vmec.dtype) + if problem.qs_surfaces: + qs = vmex.optimize.QuasisymmetryRatioResidual( + problem.qs_surfaces, + problem.helicity_m, + problem.helicity_n, + ) + quasisymmetry = qs.profile_state(vmex_solution.state, runtime) + else: + quasisymmetry = jnp.zeros((0,), dtype=iota_vmec.dtype) + return VmexRadialQuantities( + s=s, + iota=-iota_vmec, + iota_vmec=iota_vmec, + qs_surfaces=surfaces, + quasisymmetry=quasisymmetry, + magnetic_well=vmex.optimize.magnetic_well(vmex_solution.state, runtime), + aspect=vmex.optimize.aspect_ratio(vmex_solution.state, runtime), + volume=vmex.optimize.volume(vmex_solution.state, runtime), + magnetic_energy=vmex_solution.wb, + thermal_energy=vmex_solution.wp, + ) + + +def solve_vmex(problem: VmexProblem, parameters: Any | None = None) -> VmexEquilibrium: + """Solve one VMEX problem with implicit-AD-compatible parameters.""" + + vmex = _import_vmex() + selected_parameters = problem.parameters if parameters is None else parameters + solution = vmex.implicit.run( + problem.input, + selected_parameters, + ftol=problem.ftol, + max_iterations=problem.max_iterations, + adjoint_tol=problem.adjoint_tol, + multigrid=problem.multigrid, + device=problem.device, + ) + return VmexEquilibrium( + solution=solution, + quantities=_radial_quantities(problem, solution), + ) + + +def vmex_radial_quantities( + problem: VmexProblem, + parameters: Any | None = None, +) -> VmexRadialQuantities: + """Solve and return only the differentiable VMEX diagnostic pytree.""" + + return solve_vmex(problem, parameters).quantities diff --git a/tests/compatibility/test_essos_contract.py b/tests/compatibility/test_essos_contract.py new file mode 100644 index 0000000..5f74e40 --- /dev/null +++ b/tests/compatibility/test_essos_contract.py @@ -0,0 +1,59 @@ +from pyqsc_jax.near_axis import near_axis + + +def test_legacy_import_and_essos_on_axis_contract(): + field = near_axis( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + etabar=-0.9, + nfp=3, + nphi=31, + ) + + attributes = { + "B0", + "nfp", + "nphi", + "phi", + "varphi", + "R0", + "Z0", + "B_axis", + "grad_B_axis", + "axis_length", + "iota", + "iotaN", + "curvature", + "torsion", + "elongation", + "L_grad_B", + "x", + "dofs", + "rc", + "zs", + "etabar", + "sigma0", + "I2", + "spsi", + "sG", + } + methods = { + "AbsB", + "B_covariant", + "B_contravariant", + "B_mag", + "jacobian", + "get_boundary", + "Frenet_to_cylindrical", + "phi_of_theta_varphi", + "to_vmec", + } + + assert all(hasattr(field, name) for name in attributes) + assert all(callable(getattr(field, name)) for name in methods) + + +def test_legacy_plot_method_is_available(): + field = near_axis() + + assert callable(field.plot) diff --git a/tests/compatibility/test_legacy_adapter.py b/tests/compatibility/test_legacy_adapter.py new file mode 100644 index 0000000..c398607 --- /dev/null +++ b/tests/compatibility/test_legacy_adapter.py @@ -0,0 +1,153 @@ +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from pyqsc_jax.near_axis import near_axis + + +def standard_field(**kwargs): + parameters = { + "rc": [1.0, 0.045], + "zs": [0.0, -0.045], + "etabar": -0.9, + "nfp": 3, + "nphi": 15, + } + parameters.update(kwargs) + return near_axis(**parameters) + + +def test_legacy_dofs_validation_and_x_alias(): + field = standard_field() + updated = field.x.at[-1].set(-0.85) + field.x = updated + + np.testing.assert_array_equal(field.x, updated) + np.testing.assert_allclose(field.etabar, -0.85) + with pytest.raises(ValueError, match="dofs must have shape"): + field.dofs = jnp.zeros(2) + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"nphi": 16}, "odd integer"), + ({"nphi": True}, "odd integer"), + ({"rc": [[1.0]], "zs": [[0.0]]}, "one-dimensional"), + ({"rc": [1.0, 0.1], "zs": [0.0]}, "equal length"), + ], +) +def test_legacy_constructor_validation(kwargs, message): + with pytest.raises(ValueError, match=message): + standard_field(**kwargs) + + +def test_calculate_and_pytree_round_trip(): + field = standard_field() + calculated = field.calculate(field.rc, field.zs, field.etabar) + np.testing.assert_allclose(calculated[7], field.iota) + + leaves, structure = jax.tree_util.tree_flatten(field) + restored = jax.tree_util.tree_unflatten(structure, leaves) + np.testing.assert_allclose(restored.iota, field.iota) + assert restored.order == field.order + + +def test_coordinate_helpers_and_boundary_are_finite(): + field = standard_field() + X = 0.01 * field.X1c_untwisted + Y = 0.01 * field.Y1s_untwisted + R, Z, phi = field.Frenet_to_cylindrical_1_point(0.0, X, Y) + assert jnp.all(jnp.isfinite(jnp.asarray((R, Z, phi)))) + + residual = field.Frenet_to_cylindrical_residual_func(0.0, phi, X, Y) + np.testing.assert_allclose(residual, 0.0, atol=2e-14) + interpolated = field.interpolated_array_at_point(field.R0, 2 * jnp.pi / field.nfp) + np.testing.assert_allclose(interpolated, field.R0[0]) + + R_period, Z_period, phi0 = field.Frenet_to_cylindrical(0.01, ntheta=5) + assert R_period.shape == Z_period.shape == phi0.shape == (5, field.nphi) + assert jnp.all(jnp.isfinite(R_period)) + RBC, ZBS = field.to_Fourier(R_period, Z_period, field.nfp, mpol=2, ntor=2) + assert RBC.shape == ZBS.shape == (5, 3) + + x, y, z, cylindrical_R = field.get_boundary( + r=0.01, + ntheta=6, + nphi=8, + ntheta_fourier=5, + mpol=2, + ntor=2, + ) + assert x.shape == y.shape == z.shape == cylindrical_R.shape == (6, 8) + assert jnp.all(jnp.isfinite(jnp.stack((x, y, z)))) + + +def test_varphi_coordinate_inversion(): + field = standard_field() + r = 0.005 + theta = 0.3 + varphi = 0.1 + phi = field.phi_of_theta_varphi(r, theta, varphi) + assert jnp.isfinite(phi) + + R, Z, phi0 = field.Frenet_to_cylindrical(r, ntheta=3, phi_is_varphi=True) + assert R.shape == Z.shape == phi0.shape == (3, field.nphi) + assert jnp.all(jnp.isfinite(R)) + + x, y, z, _ = field.get_boundary( + r=r, + ntheta=2, + nphi=3, + ntheta_fourier=3, + mpol=1, + ntor=1, + phi_is_varphi=True, + ) + assert jnp.all(jnp.isfinite(jnp.stack((x, y, z)))) + + +def test_legacy_field_magnitude_and_jitted_methods(): + field = standard_field(I2=0.2) + point = jnp.asarray((0.01, 0.2, 0.1)) + expected = field.B0 * (1 + point[0] * field.etabar * jnp.cos(point[1])) + np.testing.assert_allclose(field.AbsB(point), expected) + assert jnp.all(jnp.isfinite(jax.jit(field.B_covariant)(point))) + assert jnp.all(jnp.isfinite(jax.jit(field.B_contravariant)(point))) + assert jnp.isfinite(jax.jit(field.jacobian)(point)) + assert jnp.isfinite(field.B_mag(*point)) + + +def test_plot_supports_created_and_supplied_axes(monkeypatch): + matplotlib = pytest.importorskip("matplotlib") + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + monkeypatch.setattr(plt, "show", lambda: None) + field = standard_field() + figure, axes = field.plot( + r=0.005, + ntheta=3, + nphi=4, + ntheta_fourier=3, + show=False, + close=True, + ) + assert axes.figure is figure + + supplied_figure = plt.figure() + supplied_axes = supplied_figure.add_subplot(projection="3d") + returned_figure, returned_axes = field.plot( + r=0.005, + ntheta=3, + nphi=4, + ntheta_fourier=3, + ax=supplied_axes, + show=True, + close=True, + axis_equal=False, + ) + assert returned_figure is supplied_figure + assert returned_axes is supplied_axes diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..ac63b78 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,3 @@ +import jax + +jax.config.update("jax_enable_x64", True) diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py new file mode 100644 index 0000000..d8f6328 --- /dev/null +++ b/tests/examples/test_examples.py @@ -0,0 +1,42 @@ +"""Every checked-in example must execute as a direct, headless script.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +EXAMPLE_SCRIPTS = tuple(sorted((REPOSITORY_ROOT / "examples").glob("[0-9][0-9]_*.py"))) +PUBLICATION_SCRIPTS = tuple( + sorted((REPOSITORY_ROOT / "examples" / "publication").glob("figure_*.py")) +) + + +@pytest.mark.slow +@pytest.mark.parametrize( + "script", + EXAMPLE_SCRIPTS + PUBLICATION_SCRIPTS, + ids=lambda path: path.stem, +) +def test_example_executes_as_direct_script(script: Path, tmp_path: Path) -> None: + """Execute the public script contract without relying on repository cwd.""" + + environment = os.environ.copy() + environment.update( + { + "JAX_ENABLE_X64": "true", + "MPLBACKEND": "Agg", + "PYTHONPATH": str(REPOSITORY_ROOT / "src"), + } + ) + subprocess.run( + (sys.executable, str(script)), + cwd=tmp_path, + env=environment, + check=True, + timeout=180, + ) diff --git a/tests/integration/test_vmec_equilibrium.py b/tests/integration/test_vmec_equilibrium.py new file mode 100644 index 0000000..78db0aa --- /dev/null +++ b/tests/integration/test_vmec_equilibrium.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import hashlib +import json +import os +import subprocess +from pathlib import Path + +import numpy as np +import pytest +from scipy.io import netcdf_file + +import pyqsc_jax as qsc + +REFERENCE_DIRECTORY = Path(__file__).resolve().parents[1] / "reference" / "vmec" + + +def _read_vmec_result(path: Path) -> dict[str, float | int]: + with netcdf_file(path, "r", mmap=False) as dataset: + return { + "ier_flag": int(dataset.variables["ier_flag"].data), + "iota_axis": float(-dataset.variables["iotaf"].data[0]), + "fsqr": float(dataset.variables["fsqr"].data), + "fsqz": float(dataset.variables["fsqz"].data), + "fsql": float(dataset.variables["fsql"].data), + "pressure_axis": float(dataset.variables["presf"].data[0]), + } + + +@pytest.mark.integration +def test_frozen_vmec_wout_recovers_near_axis_iota(): + manifest = json.loads((REFERENCE_DIRECTORY / "manifest.json").read_text(encoding="utf-8")) + vmec_input = REFERENCE_DIRECTORY / "input.qa_r0025" + wout = REFERENCE_DIRECTORY / "wout_qa_r0025.nc" + input_digest = hashlib.sha256(vmec_input.read_bytes()).hexdigest() + digest = hashlib.sha256(wout.read_bytes()).hexdigest() + result = _read_vmec_result(wout) + solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=121, + order="r2", + ) + relative_iota_error = abs(result["iota_axis"] - float(solution.iota)) / abs( + float(solution.iota) + ) + + assert input_digest == manifest["files"]["input.qa_r0025"] + assert digest == manifest["files"]["wout_qa_r0025.nc"] + assert result["ier_flag"] == 0 + assert max(result["fsqr"], result["fsqz"], result["fsql"]) < 1.0e-9 + assert relative_iota_error < 1.0e-3 + np.testing.assert_allclose( + result["iota_axis"], + manifest["vmec_result"]["iota_axis"], + rtol=0, + atol=2.0e-15, + ) + + +@pytest.mark.integration +@pytest.mark.slow +def test_local_vmec_rerun_when_executable_is_requested(tmp_path): + executable = os.environ.get("PYQSC_VMEC_EXECUTABLE") + if not executable: + pytest.skip("Set PYQSC_VMEC_EXECUTABLE to rerun the fixed-boundary VMEC case.") + solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=121, + order="r2", + ) + qsc.to_vmec( + solution, + tmp_path / "input.qa_r0025", + r=0.0025, + ntheta=32, + mpol=6, + ntor=6, + parameters={ + "ns_array": (31,), + "ftol_array": (1.0e-10,), + "niter_array": (3000,), + }, + ) + subprocess.run( + [executable, "input.qa_r0025"], + cwd=tmp_path, + check=True, + timeout=120, + ) + result = _read_vmec_result(tmp_path / "wout_qa_r0025.nc") + + assert result["ier_flag"] == 0 + assert max(result["fsqr"], result["fsqz"], result["fsql"]) < 1.0e-9 + np.testing.assert_allclose(result["iota_axis"], solution.iota, rtol=1.0e-3) + + +@pytest.mark.integration +@pytest.mark.slow +def test_local_vmec_finite_pressure_zero_current_database_case(tmp_path): + executable = os.environ.get("PYQSC_VMEC_EXECUTABLE") + if not executable: + pytest.skip("Set PYQSC_VMEC_EXECUTABLE to rerun the finite-beta VMEC case.") + solution = qsc.solve_configuration("database_qa_139524", nphi=241, order="r3") + export = qsc.to_vmec( + solution, + tmp_path / "input.qa139524_beta_r0015", + r=0.0015, + ntheta=40, + mpol=8, + ntor=8, + parameters={ + "ns_array": (31, 61), + "ftol_array": (1.0e-9, 1.0e-11), + "niter_array": (3000, 5000), + }, + ) + subprocess.run( + [executable, "input.qa139524_beta_r0015"], + cwd=tmp_path, + check=True, + timeout=120, + ) + result = _read_vmec_result(tmp_path / "wout_qa139524_beta_r0015.nc") + torsion_rms = np.sqrt( + np.sum(np.asarray(solution.torsion**2 * solution.geometry.d_l_d_phi)) + / np.sum(np.asarray(solution.geometry.d_l_d_phi)) + ) + + assert solution.inputs.I2 == 0 + assert solution.inputs.p2 != 0 + assert export.curtor == 0 + assert export.pressure_axis > 0 + assert result["pressure_axis"] == pytest.approx(export.pressure_axis) + assert torsion_rms > 1.0 + assert result["ier_flag"] == 0 + assert max(result["fsqr"], result["fsqz"], result["fsql"]) < 1.0e-9 + np.testing.assert_allclose(result["iota_axis"], solution.iota, rtol=5.0e-4) diff --git a/tests/integration/test_vmex_live.py b/tests/integration/test_vmex_live.py new file mode 100644 index 0000000..d86b8bb --- /dev/null +++ b/tests/integration/test_vmex_live.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import os + +import jax +import numpy as np +import pytest + +import pyqsc_jax as qsc + +RUN_VMEX = os.environ.get("PYQSC_RUN_VMEX") == "1" +pytestmark = [ + pytest.mark.integration, + pytest.mark.slow, + pytest.mark.skipif(not RUN_VMEX, reason="Set PYQSC_RUN_VMEX=1 for live VMEX tests."), +] + + +def _problem(solution): + return qsc.to_vmex_problem( + solution, + r=0.02, + ntheta=8, + mpol=3, + ntor=2, + ns_array=(7,), + ftol=1.0e-7, + max_iterations=1200, + adjoint_tol=1.0e-8, + multigrid=False, + qs_surfaces=(0.5, 1.0), + ) + + +def test_vmex_vacuum_profiles_and_implicit_gradient(): + solution = qsc.solve_configuration("qa", nphi=31) + problem = _problem(solution) + result = problem.solve() + quantities = result.quantities + + assert problem.vmex_version + assert quantities.iota.shape == quantities.s.shape == (7,) + assert quantities.quasisymmetry.shape == (2,) + assert np.all(np.isfinite(np.asarray(quantities.iota))) + assert np.all(np.isfinite(np.asarray(quantities.quasisymmetry))) + np.testing.assert_allclose(quantities.iota[0], solution.iota, rtol=8.0e-3) + assert float(quantities.thermal_energy) == pytest.approx(0.0, abs=1.0e-14) + + value, gradient = jax.value_and_grad( + lambda parameters: ( + qsc.vmex_radial_quantities( + problem, + parameters, + ).magnetic_well + ) + )(problem.parameters) + assert np.isfinite(float(value)) + assert np.all(np.isfinite(np.asarray(gradient.rbc))) + assert float(np.linalg.norm(np.asarray(gradient.rbc))) > 0 + + +def test_vmex_finite_beta_profiles(): + solution = qsc.solve_configuration("plasma_stellarator", nphi=31) + problem = _problem(solution) + quantities = problem.quantities() + + assert problem.finite_beta + assert float(quantities.thermal_energy) > 0 + assert np.all(np.isfinite(np.asarray(quantities.iota))) + assert np.all(np.isfinite(np.asarray(quantities.quasisymmetry))) + + value, gradient = jax.value_and_grad( + lambda parameters: ( + qsc.vmex_radial_quantities( + problem, + parameters, + ).magnetic_well + ) + )(problem.parameters) + assert np.isfinite(float(value)) + assert np.isfinite(float(gradient.pres_scale)) + assert float(abs(gradient.pres_scale)) > 0 + assert np.all(np.isfinite(np.asarray(gradient.rbc))) + assert float(np.linalg.norm(np.asarray(gradient.rbc))) > 0 diff --git a/tests/literature/test_plasma_volume_reference.py b/tests/literature/test_plasma_volume_reference.py new file mode 100644 index 0000000..2311246 --- /dev/null +++ b/tests/literature/test_plasma_volume_reference.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +def resolved_volume_biot_savart( + solution, + formal_radius, + *, + radial_resolution=16, + angular_resolution=64, +): + """Independent midpoint/Gauss volume integral at the first axis point.""" + + source = qsc.plasma_current_source( + solution, + formal_radius=formal_radius, + ) + geometry = solution.geometry + radial_nodes, radial_weights = np.polynomial.legendre.leggauss(radial_resolution) + radial = 0.5 * formal_radius * (radial_nodes + 1) + radial_weights = 0.5 * formal_radius * radial_weights + theta = 2 * np.pi * np.arange(angular_resolution) / angular_resolution + theta_weight = 2 * np.pi / angular_resolution + cosine = np.cos(theta) + sine = np.sin(theta) + cosine2 = np.cos(2 * theta) + sine2 = np.sin(2 * theta) + + tangent = np.asarray(geometry.tangent_cartesian) + normal = np.asarray(geometry.normal_cartesian) + binormal = np.asarray(geometry.binormal_cartesian) + first_order = ( + np.asarray(solution.X1c)[:, None, None] * cosine[None, :, None] * normal[:, None, :] + + ( + np.asarray(solution.Y1s)[:, None, None] * sine[None, :, None] + + np.asarray(solution.Y1c)[:, None, None] * cosine[None, :, None] + ) + * binormal[:, None, :] + ) + X2 = ( + np.asarray(solution.X20)[:, None] + + np.asarray(solution.X2c)[:, None] * cosine2 + + np.asarray(solution.X2s)[:, None] * sine2 + ) + Y2 = ( + np.asarray(solution.Y20)[:, None] + + np.asarray(solution.Y2c)[:, None] * cosine2 + + np.asarray(solution.Y2s)[:, None] * sine2 + ) + Z2 = ( + np.asarray(solution.Z20)[:, None] + + np.asarray(solution.Z2c)[:, None] * cosine2 + + np.asarray(solution.Z2s)[:, None] * sine2 + ) + second_order = ( + X2[:, :, None] * normal[:, None, :] + + Y2[:, :, None] * binormal[:, None, :] + + Z2[:, :, None] * tangent[:, None, :] + ) + source_position = ( + np.asarray(geometry.position_cartesian)[:, None, None, :] + + radial[None, :, None, None] * first_order[:, None, :, :] + + radial[None, :, None, None] ** 2 * second_order[:, None, :, :] + ) + weighted_current = float(source.axis_length_per_radian) * ( + radial[None, :, None, None] * np.asarray(source.w1)[:, None, None, :] + + radial[None, :, None, None] ** 2 + * ( + cosine[None, None, :, None] * np.asarray(source.w2_cosine)[:, None, None, :] + + sine[None, None, :, None] * np.asarray(source.w2_sine)[:, None, None, :] + ) + ) + displacement = np.asarray(geometry.position_cartesian[0]) - source_position + kernel = ( + np.cross(weighted_current, displacement) + / np.linalg.norm( + displacement, + axis=-1, + )[..., None] + ** 3 + ) + cylindrical_period = 2 * np.pi / solution.inputs.axis.nfp + d_phi = cylindrical_period / solution.inputs.nphi + varphi_weights = np.asarray(geometry.d_varphi_d_phi) * d_phi + return np.sum( + kernel + * varphi_weights[:, None, None, None] + * radial_weights[None, :, None, None] + * theta_weight, + axis=(0, 1, 2), + ) / (4 * np.pi) + + +@pytest.mark.literature +@pytest.mark.slow +def test_matched_field_converges_to_resolved_volume_current_biot_savart(): + solution = qsc.Qsc( + rc=[1.0], + zs=[0.0], + nfp=1, + etabar=1.0, + I2=0.1, + nphi=301, + order="r2", + ) + radii = (0.1, 0.07) + errors = [] + scaled_errors = [] + for radius in radii: + direct = resolved_volume_biot_savart(solution, radius) + asymptotic = np.asarray( + qsc.plasma_field_on_axis( + solution, + formal_radius=radius, + ).field[0] + ) + error = np.linalg.norm(direct - asymptotic) + errors.append(error) + scaled_errors.append(error / (radius**4 * abs(np.log(radius)))) + + assert errors[1] < 0.35 * errors[0] + assert max(scaled_errors) / min(scaled_errors) < 1.2 + assert errors[0] < 1.0e-5 diff --git a/tests/numerics/test_first_order_autodiff.py b/tests/numerics/test_first_order_autodiff.py new file mode 100644 index 0000000..3d0d6f8 --- /dev/null +++ b/tests/numerics/test_first_order_autodiff.py @@ -0,0 +1,51 @@ +import jax +import jax.numpy as jnp +import numpy as np + +import pyqsc_jax as qsc + +AXIS = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) + + +def iota_from_etabar(etabar): + return qsc.solve(axis=AXIS, etabar=etabar, nphi=31).iota + + +def test_jit_vmap_jvp_and_vjp(): + etabar = -0.9 + eager = iota_from_etabar(etabar) + compiled = jax.jit(iota_from_etabar)(etabar) + np.testing.assert_allclose(compiled, eager, rtol=2e-13) + + batched = jax.vmap(iota_from_etabar)(jnp.array([-0.95, -0.9, -0.85])) + assert batched.shape == (3,) + assert jnp.all(jnp.isfinite(batched)) + + _, jvp = jax.jvp(iota_from_etabar, (etabar,), (1.0,)) + _, pullback = jax.vjp(iota_from_etabar, etabar) + (vjp,) = pullback(jnp.asarray(1.0)) + np.testing.assert_allclose(jvp, vjp, rtol=2e-12, atol=2e-12) + + +def test_iota_derivative_matches_finite_difference(): + etabar = -0.9 + derivative = jax.grad(iota_from_etabar)(etabar) + step = 2e-5 + finite_difference = (iota_from_etabar(etabar + step) - iota_from_etabar(etabar - step)) / ( + 2 * step + ) + + np.testing.assert_allclose(derivative, finite_difference, rtol=2e-7, atol=2e-9) + + +def test_axis_coefficient_derivative_matches_finite_difference(): + def iota_from_rc1(rc1): + axis = qsc.Axis(rc=jnp.array([1.0, rc1]), zs=AXIS.zs, nfp=AXIS.nfp) + return qsc.solve(axis=axis, etabar=-0.9, nphi=31).iota + + rc1 = 0.045 + derivative = jax.grad(iota_from_rc1)(rc1) + step = 2e-6 + finite_difference = (iota_from_rc1(rc1 + step) - iota_from_rc1(rc1 - step)) / (2 * step) + + np.testing.assert_allclose(derivative, finite_difference, rtol=3e-6, atol=3e-8) diff --git a/tests/numerics/test_second_order_autodiff.py b/tests/numerics/test_second_order_autodiff.py new file mode 100644 index 0000000..139f102 --- /dev/null +++ b/tests/numerics/test_second_order_autodiff.py @@ -0,0 +1,42 @@ +import jax +import jax.numpy as jnp +import numpy as np + +import pyqsc_jax as qsc + + +def b20_residual(etabar): + return qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=etabar, + B2c=-0.00322, + nphi=15, + order="r2", + ).B20_residual + + +def test_second_order_jit_jvp_and_vjp(): + etabar = 0.64 + eager = b20_residual(etabar) + np.testing.assert_allclose(jax.jit(b20_residual)(etabar), eager, rtol=2e-13) + + _, jvp = jax.jvp(b20_residual, (etabar,), (1.0,)) + _, pullback = jax.vjp(b20_residual, etabar) + (vjp,) = pullback(jnp.asarray(1.0)) + np.testing.assert_allclose(jvp, vjp, rtol=3e-11, atol=3e-11) + + +def test_second_order_gradient_matches_finite_difference(): + etabar = 0.64 + derivative = jax.grad(b20_residual)(etabar) + step = 2e-5 + finite_difference = (b20_residual(etabar + step) - b20_residual(etabar - step)) / (2 * step) + np.testing.assert_allclose(derivative, finite_difference, rtol=2e-6, atol=2e-8) + + +def test_second_order_vmap(): + values = jax.vmap(b20_residual)(jnp.asarray([0.61, 0.64, 0.67])) + assert values.shape == (3,) + assert jnp.all(jnp.isfinite(values)) diff --git a/tests/physics/test_axis_optimization.py b/tests/physics/test_axis_optimization.py new file mode 100644 index 0000000..142718a --- /dev/null +++ b/tests/physics/test_axis_optimization.py @@ -0,0 +1,474 @@ +from __future__ import annotations + +from dataclasses import replace + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +import pyqsc_jax.axis_optimization as axis_optimization +from pyqsc_jax.axis_optimization import ( + _candidate_is_distinct, + _copy_retained_axis_modes, + _halton_box, + _internal_from_physical, + _levenberg_marquardt, + _physical_from_internal, + _selector_value, + _validate_problem, + _verification_passes, +) + + +def circular_problem(**overrides): + parameters = { + "axis": qsc.Axis(rc=[1.0], zs=[0.0], nfp=1), + "variable_indices": (0,), + "lower_bounds": jnp.asarray([0.8]), + "upper_bounds": jnp.asarray([1.2]), + "etabar": 1.0, + "I2": 0.1, + "nphi": 15, + } + parameters.update(overrides) + return qsc.AxisSearchProblem(**parameters) + + +def small_options(**overrides): + parameters = { + "coarse_samples": 2, + "local_starts": 1, + "maximum_iterations": 2, + "verification_multipliers": (1, 2), + "verified_zero_tolerance": 1.0e-10, + "verification_tail_tolerance": 1.0e-8, + } + parameters.update(overrides) + return qsc.AxisSearchOptions(**parameters) + + +def test_stellarator_symmetric_indices_and_low_discrepancy_box(): + axis = qsc.Axis(rc=[1.0, 0.1, 0.01], zs=[0.0, -0.1, 0.02], nfp=3) + + assert qsc.stellarator_symmetric_variable_indices(axis) == (1, 2, 10, 11) + assert qsc.stellarator_symmetric_variable_indices(axis, modes=(2,)) == (2, 11) + with pytest.raises(ValueError, match="modes"): + qsc.stellarator_symmetric_variable_indices(axis, modes=(0,)) + with pytest.raises(ValueError, match="modes"): + qsc.stellarator_symmetric_variable_indices(axis, modes=(3,)) + samples = _halton_box(3, 3, jnp.float64) + np.testing.assert_allclose( + samples, + [ + [0.5, 1 / 3, 0.2], + [0.25, 2 / 3, 0.4], + [0.75, 1 / 9, 0.6], + ], + ) + + +def test_bound_transform_round_trip_and_strict_interior(): + lower = jnp.asarray([-2.0, 1.0]) + upper = jnp.asarray([4.0, 5.0]) + physical = jnp.asarray([-1.0, 4.0]) + internal = _internal_from_physical(physical, lower, upper) + reconstructed = _physical_from_internal(internal, lower, upper) + + np.testing.assert_allclose(reconstructed, physical) + assert np.all(np.asarray(reconstructed) > np.asarray(lower)) + assert np.all(np.asarray(reconstructed) < np.asarray(upper)) + + +def test_local_solve_accepts_a_finite_rank_deficient_exact_zero(): + options = small_options(maximum_iterations=1) + initial = jnp.asarray([0.0, 0.0]) + residual = lambda value: jnp.zeros((3,), dtype=value.dtype) # noqa: E731 + + internal, report = _levenberg_marquardt(residual, initial, options) + + np.testing.assert_allclose(internal, initial) + assert report.converged + assert report.finite + assert np.isinf(float(report.jacobian_condition_number)) + + +def test_local_solve_accepts_a_finite_converged_residual_with_nan_derivative(): + @jax.custom_vjp + def residual(_value): + return jnp.zeros((3,), dtype=_value.dtype) + + def residual_forward(value): + return jnp.zeros((3,), dtype=value.dtype), value.shape + + def residual_backward(shape, cotangent): + return (jnp.full(shape, jnp.nan, dtype=cotangent.dtype),) + + residual.defvjp(residual_forward, residual_backward) + options = small_options(maximum_iterations=1) + initial = jnp.asarray([0.0, 0.0]) + + internal, report = _levenberg_marquardt(residual, initial, options) + + np.testing.assert_allclose(internal, initial) + assert report.converged + assert report.finite + assert np.isinf(float(report.jacobian_condition_number)) + + +def test_fourier_continuation_copies_retained_modes_and_runs_stages(): + low_axis = qsc.Axis(rc=[1.0], zs=[0.0], nfp=1) + high_axis = qsc.Axis(rc=[0.9, 0.0], zs=[0.0, 0.0], nfp=1) + copied = _copy_retained_axis_modes(low_axis, high_axis) + np.testing.assert_allclose(copied.rc, [1.0, 0.0]) + with pytest.raises(ValueError, match="same nfp"): + _copy_retained_axis_modes(low_axis, replace(high_axis, nfp=2)) + + stages = ( + circular_problem(axis=low_axis, nphi=15), + qsc.AxisSearchProblem( + axis=high_axis, + variable_indices=(1, 7), + lower_bounds=jnp.asarray([-0.01, -0.01]), + upper_bounds=jnp.asarray([0.01, 0.01]), + etabar=1.0, + I2=0.1, + nphi=15, + ), + ) + continuation = qsc.continue_axis_search( + stages, + options=small_options( + coarse_samples=1, + maximum_iterations=1, + verification_multipliers=(1,), + ), + ) + + assert continuation.complete + assert len(continuation.stages) == 2 + assert continuation.final is continuation.stages[-1] + assert continuation.final.best is not None + np.testing.assert_allclose(continuation.final.best.solution.axis.rc[0], 1.0) + + +def test_search_certifies_only_the_verified_nonnegative_zero(): + result = qsc.search_axis(circular_problem(), options=small_options()) + + assert result.status == "verified_zero" + assert result.global_certificate + assert "global lower bound" in result.message + assert result.best is not None + assert result.best.feasible + assert result.best.primary_residual < 1.0e-12 + assert result.best.local_report.converged + assert result.best.local_report.rejected_steps >= 1 + assert result.local_starts_attempted == 1 + assert result.distinct_basins == 1 + assert result.search_budget == 3 + int(result.best.local_report.function_evaluations) + np.testing.assert_allclose(result.coarse_variables[0], [1.0]) + assert np.all(np.asarray(result.coarse_feasible)) + assert not _candidate_is_distinct( + result.best.variables, + [result.best], + jnp.asarray([0.8]), + jnp.asarray([1.2]), + 1.0e-3, + ) + + +def test_search_failure_and_ill_conditioning_are_distinct_statuses(monkeypatch): + coarse_failure = qsc.search_axis( + circular_problem(etabar=0.0), + options=small_options( + coarse_samples=1, + verification_multipliers=(1,), + ), + ) + assert coarse_failure.status == "solver_failure" + assert coarse_failure.best is None + assert coarse_failure.local_starts_attempted == 0 + + original_local_solve = axis_optimization._levenberg_marquardt + + def forced_nonfinite_local_solve(*args, **kwargs): + internal, report = original_local_solve(*args, **kwargs) + return internal, replace(report, finite=jnp.asarray(False)) + + with monkeypatch.context() as context: + context.setattr( + axis_optimization, + "_levenberg_marquardt", + forced_nonfinite_local_solve, + ) + local_failure = qsc.search_axis( + circular_problem(), + options=small_options( + coarse_samples=1, + maximum_iterations=1, + verification_multipliers=(1,), + ), + ) + assert local_failure.status == "solver_failure" + assert local_failure.best is None + assert local_failure.local_starts_attempted == 1 + + ill_conditioned = qsc.search_axis( + circular_problem(), + options=small_options( + coarse_samples=1, + maximum_iterations=1, + verification_multipliers=(1,), + maximum_linear_condition_number=1.0e-30, + ), + ) + assert ill_conditioned.status == "ill_conditioned" + assert ill_conditioned.best is None + assert ill_conditioned.local_starts_attempted == 1 + + +def test_search_reports_no_feasible_candidate_for_failed_hard_profile(): + impossible = qsc.Criteria.from_curvo_2025(minimum_abs_iota=2.0) + result = qsc.search_axis( + circular_problem( + criteria=impossible, + selector="maximum_criteria_margin", + ), + options=small_options(verification_multipliers=(1,)), + seeds=(jnp.asarray([0.9]),), + ) + + assert result.status == "no_feasible_candidate" + assert not result.global_certificate + assert result.best is not None + assert not result.best.feasible + assert result.best.criteria_report is not None + assert not result.best.criteria_report.passed + assert result.coarse_variables.shape == (4, 1) + + +def test_nonzero_search_improves_locally_without_global_claim(): + axis = qsc.Axis( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + ) + indices = qsc.stellarator_symmetric_variable_indices(axis, modes=(1,)) + problem = qsc.AxisSearchProblem( + axis=axis, + variable_indices=indices, + lower_bounds=jnp.asarray([0.11, 0.11]), + upper_bounds=jnp.asarray([0.19, 0.19]), + etabar=0.64, + nphi=15, + selector="minimum_maximum_elongation", + ) + result = qsc.search_axis( + problem, + options=small_options( + coarse_samples=2, + maximum_iterations=3, + verified_zero_tolerance=1.0e-8, + verification_tail_tolerance=1.0, + verification_multipliers=(1,), + ), + ) + + assert result.status == "best_found" + assert not result.global_certificate + assert "no global-minimum certificate" in result.message + assert result.best.primary_residual < result.coarse_residuals[0] + assert result.best.local_report.accepted_steps >= 1 + assert result.best.selector_value < 0 + + +def test_verification_rejects_an_underresolved_nominal_zero(): + solution = qsc.optimize_B2c( + qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=0.64, + nphi=15, + order="r2", + ) + ).solution + verification = qsc.verify_B20_resolution(solution, multipliers=(1, 2)) + + assert not _verification_passes( + verification, + small_options( + verified_zero_tolerance=0.05, + verification_relative_tolerance=1.0e-6, + verification_tail_tolerance=1.0e-6, + ), + ) + + +def test_every_canonical_selector_has_documented_orientation(): + solution = qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=0.9, + p2=-600000.0, + B2c=-0.7, + nphi=15, + order="r2", + ) + criteria = qsc.Criteria.from_curvo_2025() + report = criteria.evaluate(solution) + base = circular_problem(axis=solution.axis, variable_indices=(1,)) + + np.testing.assert_allclose( + _selector_value( + replace(base, selector="maximum_singular_radius"), + solution, + report, + ), + solution.r_singularity, + ) + np.testing.assert_allclose( + _selector_value( + replace(base, selector="maximum_minimum_L_grad_B"), + solution, + report, + ), + np.min(solution.L_grad_B), + ) + np.testing.assert_allclose( + _selector_value( + replace(base, selector="maximum_minimum_L_grad_grad_B"), + solution, + report, + ), + np.min(solution.L_grad_grad_B), + ) + np.testing.assert_allclose( + _selector_value( + replace(base, selector="minimum_maximum_elongation"), + solution, + report, + ), + -np.max(solution.elongation), + ) + assert ( + _selector_value( + replace(base, selector="minimum_axis_sobolev_norm"), + solution, + report, + ) + < 0 + ) + assert ( + _selector_value( + replace( + base, + selector="target_axis_length", + target_axis_length=float(solution.axis_length), + ), + solution, + report, + ) + == 0 + ) + assert np.isfinite( + _selector_value( + replace( + base, + criteria=criteria, + selector="maximum_criteria_margin", + ), + solution, + report, + ) + ) + + +@pytest.mark.parametrize( + ("problem", "options", "message"), + [ + (circular_problem(variable_indices=()), small_options(), "variable_indices"), + ( + circular_problem(variable_indices=(0, 0)), + small_options(), + "variable_indices", + ), + (circular_problem(variable_indices=(4,)), small_options(), "variable_indices"), + (circular_problem(lower_bounds=[0.8, 0.9]), small_options(), "bounds"), + (circular_problem(lower_bounds=[1.2]), small_options(), "strictly ordered"), + (circular_problem(lower_bounds=[1.0]), small_options(), "strictly inside"), + (circular_problem(nphi=4), small_options(), "nphi"), + ( + circular_problem(target_iota=0.1), + small_options(), + "target_iota", + ), + ( + circular_problem(solve_for="etabar"), + small_options(), + "target_iota", + ), + ( + circular_problem(solve_for="bad"), + small_options(), + "solve_for", + ), + ( + circular_problem(selector="target_axis_length"), + small_options(), + "target_axis_length", + ), + ( + circular_problem(selector="maximum_criteria_margin"), + small_options(), + "criteria", + ), + ( + circular_problem(selector="bad"), + small_options(), + "selector", + ), + ( + circular_problem(), + small_options(coarse_samples=0), + "positive integers", + ), + ( + circular_problem(), + small_options(initial_damping=0.0), + "must be positive", + ), + ( + circular_problem(), + small_options(verification_multipliers=()), + "nonempty", + ), + ], +) +def test_search_policy_guards(problem, options, message): + with pytest.raises(ValueError, match=message): + _validate_problem(problem, options) + + +def test_valid_inverse_search_policy_reaches_post_branch_validation(): + _validate_problem( + circular_problem( + solve_for="I2", + target_iota=0.1, + ), + small_options(), + ) + + +def test_explicit_seed_guard(): + with pytest.raises(ValueError, match="explicit seed"): + qsc.search_axis( + circular_problem(), + options=small_options(), + seeds=(jnp.asarray([1.3]),), + ) + with pytest.raises(ValueError, match="At least one"): + qsc.continue_axis_search(()) diff --git a/tests/physics/test_b20_optimization.py b/tests/physics/test_b20_optimization.py new file mode 100644 index 0000000..661bfd4 --- /dev/null +++ b/tests/physics/test_b20_optimization.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +def qa_solution(nphi=31, order="r2", **overrides): + parameters = { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + "nphi": nphi, + "order": order, + } + parameters.update(overrides) + return qsc.Qsc(**parameters) + + +def test_dense_B20_diagnostics_match_direct_definitions(): + solution = qa_solution() + diagnostics = qsc.b20_diagnostics(solution, smooth_maximum_power=12) + weights = solution.geometry.d_l_d_phi + mean = np.sum(np.asarray(weights * solution.B20)) / np.sum(np.asarray(weights)) + anomaly = np.asarray(solution.B20) - mean + + np.testing.assert_allclose(diagnostics.weighted_mean, mean) + np.testing.assert_allclose(diagnostics.anomaly, anomaly) + np.testing.assert_allclose( + diagnostics.weighted_l2, + np.sqrt( + np.sum(np.asarray(weights) * (anomaly / solution.inputs.B0) ** 2) / np.sum(weights) + ), + ) + np.testing.assert_allclose(diagnostics.weighted_l2, solution.B20_residual) + np.testing.assert_allclose( + diagnostics.grid_maximum, + np.max(np.abs(anomaly / solution.inputs.B0)), + ) + np.testing.assert_allclose( + diagnostics.peak_to_peak, + solution.B20_variation / solution.inputs.B0, + ) + assert diagnostics.smooth_maximum_power == 12 + assert diagnostics.fourier_modes.shape == (15,) + assert diagnostics.fourier_coefficients.shape == (15,) + np.testing.assert_allclose( + diagnostics.nonzero_fourier_l1, + 2 * np.sum(np.abs(diagnostics.fourier_coefficients / solution.inputs.B0)), + ) + assert diagnostics.weighted_l2 <= diagnostics.smooth_maximum <= diagnostics.grid_maximum + assert 0 <= float(diagnostics.fourier_tail_ratio) <= 1 + + +def test_affine_B2c_elimination_is_exact_and_stationary(): + solution = qa_solution() + result = qsc.optimize_B2c(solution) + minus = qa_solution(B2c=result.B2c_optimal - 0.1) + plus = qa_solution(B2c=result.B2c_optimal + 0.1) + + np.testing.assert_allclose(result.B2c_optimal, -0.49172683641534204, rtol=3.0e-12) + assert result.affine_reconstruction_error < 5.0e-14 + assert not bool(result.degenerate) + assert result.diagnostics.weighted_l2 < solution.B20_residual + assert result.diagnostics.weighted_l2 < minus.B20_residual + assert result.diagnostics.weighted_l2 < plus.B20_residual + + def squared_residual(B2c): + return qa_solution(nphi=15, B2c=B2c).B20_residual ** 2 + + optimum = qsc.optimal_B2c_value(qa_solution(nphi=15)) + derivative = jax.grad(squared_residual)(optimum) + assert abs(float(derivative)) < 2.0e-11 + + +def test_B2c_optimum_is_jittable_and_differentiable(): + def optimum(etabar): + return qsc.optimal_B2c_value(qa_solution(nphi=15, etabar=etabar)) + + value = optimum(jnp.asarray(0.64)) + jitted = jax.jit(optimum)(jnp.asarray(0.64)) + tangent = jax.jvp(optimum, (jnp.asarray(0.64),), (jnp.asarray(1.0),))[1] + step = 1.0e-5 + finite_difference = (optimum(0.64 + step) - optimum(0.64 - step)) / (2 * step) + + np.testing.assert_allclose(jitted, value, rtol=2.0e-11) + np.testing.assert_allclose(tangent, finite_difference, rtol=2.0e-6) + + +def test_optimal_solution_recomputes_r3_and_shear_without_stale_data(): + original = qsc.solve_magnetic_shear(qa_solution(nphi=15, order="r3")) + result = qsc.optimize_B2c(original) + + assert result.solution.third_order is not None + assert result.solution.shear is not None + np.testing.assert_allclose(result.solution.B31c, original.B31c) + assert result.solution.flux_constraint_residual < 2.0e-13 + assert result.solution.consistency_error < 2.0e-10 + + +def test_circular_axis_has_degenerate_nonconstant_B2c_response(): + solution = qsc.Qsc( + rc=[1.0], + zs=[0.0], + nfp=1, + etabar=1.0, + I2=0.1, + B2c=0.2, + nphi=15, + order="r2", + ) + result = qsc.optimize_B2c(solution) + + assert bool(result.degenerate) + np.testing.assert_allclose(result.B2c_optimal, 0.2) + assert result.diagnostics.weighted_l2 < 2.0e-14 + + +def test_resolution_verification_recomputes_fixed_candidate(): + result = qsc.optimize_B2c(qa_solution()) + verification = qsc.verify_B20_resolution(result.solution, multipliers=(1, 2)) + + np.testing.assert_array_equal(verification.resolutions, [31, 61]) + np.testing.assert_allclose(verification.weighted_l2[0], result.diagnostics.weighted_l2) + assert verification.relative_weighted_l2_change[1] < 2.0e-6 + assert verification.fourier_tail_ratio[1] < verification.fourier_tail_ratio[0] + assert np.all(np.asarray(verification.nonzero_fourier_l1) >= 0) + assert verification.relative_grid_maximum_change[1] < 0.03 + + +@pytest.mark.physics +def test_documented_optimized_axis_has_nearly_constant_B20(): + stock = qsc.optimize_B2c(qsc.solve_configuration("database_low_b20_57409", nphi=121)) + optimized = qsc.solve_configuration("b20_optimized_good", nphi=121) + diagnostics = qsc.b20_diagnostics(optimized) + verification = qsc.verify_B20_resolution(optimized, multipliers=(1, 2)) + criteria = qsc.Criteria.from_curvo_2025(minimum_abs_iota=0.4) + + assert float(diagnostics.weighted_l2) < 1.4e-10 + assert float(diagnostics.grid_maximum) < 3.0e-10 + assert float(diagnostics.peak_to_peak) < 6.0e-10 + assert float(stock.diagnostics.weighted_l2 / diagnostics.weighted_l2) > 2.0e8 + assert float(verification.relative_weighted_l2_change[1]) < 2.0e-3 + assert criteria.evaluate(optimized).passed + assert abs(float(optimized.iota)) > 0.4 + + +def test_B20_optimization_input_guards(): + first_order = qa_solution(order="r1") + with pytest.raises(ValueError, match="second-order"): + qsc.b20_diagnostics(first_order) + with pytest.raises(ValueError, match="second-order"): + qsc.optimal_B2c_value(first_order) + with pytest.raises(ValueError, match="smooth_maximum_power"): + qsc.b20_diagnostics(qa_solution(), smooth_maximum_power=1) + with pytest.raises(ValueError, match="degeneracy_tolerance"): + qsc.optimal_B2c_value(qa_solution(), degeneracy_tolerance=-1.0) + with pytest.raises(ValueError, match="multipliers"): + qsc.verify_B20_resolution(qa_solution(), multipliers=()) + with pytest.raises(ValueError, match="multipliers"): + qsc.verify_B20_resolution(qa_solution(), multipliers=(True,)) diff --git a/tests/physics/test_continuation.py b/tests/physics/test_continuation.py new file mode 100644 index 0000000..0e6f138 --- /dev/null +++ b/tests/physics/test_continuation.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import numpy as np +import pytest + +import pyqsc_jax as qsc + +AXIS = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) + + +def test_pseudo_arclength_crosses_etabar_iota_fold(): + branch = qsc.continue_etabar_branch( + axis=AXIS, + etabar_start=-1.0, + etabar_next=-0.98, + num_points=15, + nphi=31, + ) + + assert branch.complete + assert branch.status == "complete" + assert len(branch.solutions) == 15 + assert branch.etabar.shape == (15,) + assert branch.iota.shape == (15,) + assert branch.sigma.shape == (15, 31) + assert branch.tangents.shape == (15, 2) + assert int(np.count_nonzero(np.asarray(branch.fold_detected))) == 1 + fold_index = int(np.flatnonzero(np.asarray(branch.fold_detected))[0]) + assert fold_index == 9 + assert float(branch.response_derivative[fold_index - 1]) > 0 + assert float(branch.response_derivative[fold_index]) < 0 + assert float(branch.iota[fold_index]) > float(branch.iota[fold_index - 2]) + assert float(branch.iota[fold_index]) > float(branch.iota[fold_index + 2]) + assert np.all(np.diff(np.asarray(branch.etabar)) > 0) + + for solution in branch.solutions: + assert bool(solution.root_report.converged) + assert solution.root_report.residual_norm < 1.0e-12 + + +def test_corrected_points_match_independent_forward_solves(): + branch = qsc.continue_etabar_branch( + axis=AXIS, + etabar_start=-0.9, + etabar_next=-0.88, + num_points=6, + step_size=0.018, + nphi=31, + ) + + for index in (2, 4, 5): + direct = qsc.solve(axis=AXIS, etabar=branch.etabar[index], nphi=31) + np.testing.assert_allclose(branch.iota[index], direct.iota, atol=2.0e-13) + np.testing.assert_allclose(branch.sigma[index], direct.sigma, atol=3.0e-12) + + norms = np.linalg.norm(np.diff(np.stack((branch.etabar, branch.iota), axis=1), axis=0), axis=1) + np.testing.assert_allclose(norms[1:], 0.018, rtol=8.0e-4) + + +def test_positive_branch_and_two_point_result(): + branch = qsc.continue_etabar_branch( + axis=AXIS, + etabar_start=0.9, + etabar_next=0.88, + num_points=2, + nphi=31, + ) + + assert branch.complete + assert np.all(np.asarray(branch.etabar) > 0) + np.testing.assert_allclose(branch.iota[0], 0.41830690943386584, rtol=2.0e-13) + + +def test_continuation_reports_fixed_sign_boundary(): + branch = qsc.continue_etabar_branch( + axis=AXIS, + etabar_start=-0.1, + etabar_next=-0.05, + num_points=5, + step_size=0.2, + nphi=31, + ) + + assert branch.status == "branch_zero_crossing" + assert not branch.complete + assert len(branch.solutions) == 2 + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"num_points": 1}, "num_points"), + ({"num_points": True}, "num_points"), + ({"step_size": 0.0}, "step_size"), + ({"fold_tolerance": -1.0}, "fold_tolerance"), + ({"etabar_start": [-1.0]}, "scalars"), + ({"etabar_start": 0.0}, "nonzero"), + ({"etabar_next": 0.9}, "same etabar sign"), + ({"etabar_next": -1.0}, "distinct"), + ], +) +def test_continuation_input_guards(overrides, message): + parameters = { + "axis": AXIS, + "etabar_start": -1.0, + "etabar_next": -0.98, + "num_points": 3, + "nphi": 15, + } + parameters.update(overrides) + with pytest.raises(ValueError, match=message): + qsc.continue_etabar_branch(**parameters) diff --git a/tests/physics/test_criteria.py b/tests/physics/test_criteria.py new file mode 100644 index 0000000..2f48344 --- /dev/null +++ b/tests/physics/test_criteria.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from dataclasses import replace + +import jax +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.second_order import MU0 + + +def finite_pressure_solution(order="r2", **overrides): + parameters = { + "rc": [1.0, 0.09], + "zs": [0.0, -0.09], + "nfp": 2, + "etabar": 0.95, + "I2": 0.9, + "p2": -600000.0, + "B2c": -0.7, + "nphi": 31, + "order": order, + } + parameters.update(overrides) + return qsc.Qsc(**parameters) + + +def test_curvo_profile_values_margins_and_scaling(): + solution = finite_pressure_solution() + criteria = qsc.Criteria.from_curvo_2025(major_radius=2.0, B0=3.0) + report = criteria.evaluate(solution) + + assert report.profile_name == "curvo_2025" + assert len(report.evaluations) == 10 + assert set(report.values) == set(report.margins) + np.testing.assert_allclose(criteria.minimum_L_grad_B, 0.2) + np.testing.assert_allclose(criteria.minimum_axis_radius, 0.6) + np.testing.assert_allclose(criteria.minimum_singular_radius, 0.1) + np.testing.assert_allclose(criteria.minimum_L_grad_grad_B, 0.2) + np.testing.assert_allclose(criteria.maximum_B20_variation, 3.75) + + beta = -MU0 * solution.inputs.p2 * solution.r_singularity**2 / solution.inputs.B0**2 + np.testing.assert_allclose(report["beta"].value, beta) + np.testing.assert_allclose( + report["B20_variation"].margin, + criteria.maximum_B20_variation - solution.B20_variation, + ) + np.testing.assert_allclose( + report["minimum_L_grad_B"].margin, + np.min(solution.L_grad_B) - criteria.minimum_L_grad_B, + ) + assert report["B20_variation"].sense == "max" + assert report["B20_variation"].units == "T/m^2" + + +def test_profile_is_fully_configurable_and_reports_aggregate_pass(): + solution = finite_pressure_solution() + report = qsc.Criteria( + minimum_axis_length=-1.0, + minimum_abs_iota=0.0, + maximum_elongation=1.0e6, + minimum_L_grad_B=0.0, + minimum_axis_radius=0.0, + minimum_singular_radius=0.0, + minimum_L_grad_grad_B=0.0, + maximum_B20_variation=1.0e6, + minimum_beta=0.0, + minimum_DMerc_times_r2=-1.0e6, + ).evaluate(solution) + + assert report.passed + assert all(bool(evaluation.passed) for evaluation in report.evaluations) + with pytest.raises(KeyError, match="missing"): + _ = report["missing"] + + +def test_strict_criteria_reject_equality_and_inclusive_criteria_accept_it(): + solution = finite_pressure_solution() + baseline = qsc.Criteria.from_curvo_2025().evaluate(solution) + criteria = qsc.Criteria.from_curvo_2025( + minimum_axis_length=float(solution.axis_length), + minimum_abs_iota=float(abs(solution.iota)), + minimum_DMerc_times_r2=float(solution.DMerc_times_r2), + ) + report = criteria.evaluate(solution) + + assert not bool(report["axis_length"].passed) + assert bool(report["abs_iota"].passed) + assert not bool(report["DMerc_times_r2"].passed) + assert baseline["axis_length"].sense == "strict_min" + + +def test_criterion_margin_is_jittable_and_differentiable(): + criteria = qsc.Criteria.from_curvo_2025() + + def beta_margin(p2): + return criteria.evaluate(finite_pressure_solution(nphi=15, p2=p2))["beta"].margin + + p2 = -600000.0 + value = beta_margin(p2) + jitted = jax.jit(beta_margin)(p2) + tangent = jax.jvp(beta_margin, (p2,), (1.0,))[1] + step = 1.0 + finite_difference = (beta_margin(p2 + step) - beta_margin(p2 - step)) / (2 * step) + + np.testing.assert_allclose(jitted, value) + np.testing.assert_allclose(tangent, finite_difference, rtol=2.0e-6) + + +def test_profile_guards_and_second_order_requirement(): + with pytest.raises(ValueError, match="major_radius"): + qsc.Criteria.from_curvo_2025(major_radius=0.0) + with pytest.raises(ValueError, match="B0"): + qsc.Criteria.from_curvo_2025(B0=0.0) + with pytest.raises(ValueError, match="Unknown criteria"): + qsc.Criteria.from_curvo_2025(unknown_threshold=1.0) + with pytest.raises(ValueError, match="second-order"): + qsc.Criteria.from_curvo_2025().evaluate(finite_pressure_solution(order="r1")) + + +def test_threshold_replacement_changes_only_selected_policy(): + criteria = qsc.Criteria.from_curvo_2025() + changed = replace(criteria, minimum_abs_iota=0.4) + + assert criteria.minimum_abs_iota == 0.2 + assert changed.minimum_abs_iota == 0.4 + assert changed.maximum_elongation == criteria.maximum_elongation diff --git a/tests/physics/test_field_jet.py b/tests/physics/test_field_jet.py new file mode 100644 index 0000000..afd1111 --- /dev/null +++ b/tests/physics/test_field_jet.py @@ -0,0 +1,151 @@ +import jax +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +def _vacuum_qa(nphi=61): + return qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=0.64, + B2c=-0.00322, + nphi=nphi, + order="r2", + ) + + +def _finite_pressure_current(nphi=61): + return qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=0.9, + B2c=-0.7, + p2=-600000.0, + nphi=nphi, + order="r2", + ) + + +@pytest.mark.parametrize("factory", [_vacuum_qa, _finite_pressure_current]) +def test_total_field_jet_satisfies_chain_rule_and_maxwell_identities(factory): + solution = factory() + jet = solution.field_jet + + assert jet is not None + assert jet.field.shape == (61, 3) + assert jet.gradient.shape == (61, 3, 3) + assert jet.hessian.shape == (61, 3, 3, 3) + assert jet.hessian_frenet.shape == (61, 3, 3, 3) + assert jet.coordinate_jacobian.shape == (61, 3, 3) + assert jet.coordinate_hessian.shape == (61, 3, 3, 3) + assert jet.minimum_absolute_coordinate_jacobian > 0.9 + assert jet.maximum_field_error < 5e-15 + assert jet.maximum_gradient_error < 2e-8 + assert jet.maximum_divergence < 3e-8 + assert jet.maximum_derivative_asymmetry < 2e-13 + assert jet.maximum_divergence_gradient < 5e-5 + np.testing.assert_allclose(solution.grad_grad_B_axis, jet.hessian) + np.testing.assert_allclose(solution.grad_grad_B, jet.hessian_frenet) + + +def test_vacuum_field_hessian_is_fully_symmetric_and_trace_free(): + solution = _vacuum_qa() + hessian = solution.field_jet.hessian + + np.testing.assert_allclose(hessian, np.swapaxes(hessian, 1, 2), rtol=0, atol=7e-7) + np.testing.assert_allclose(hessian, np.swapaxes(hessian, 1, 3), rtol=0, atol=7e-7) + np.testing.assert_allclose( + np.einsum("niik->nk", hessian), + 0, + rtol=0, + atol=7e-7, + ) + + +@pytest.mark.parametrize( + "factory, expected_components, expected_inverse_scale", + [ + ( + _vacuum_qa, + { + (0, 0, 0): [2.5058735043974697e-11, 1.3780320235299797, -3.0796390010619596], + (0, 2, 1): [2.5345018639059774e-14, -0.9042329447022055, 0.04480778724398558], + (1, 2, 2): [-1.0583277338937587, -1.0542622720207702, 1.4815487851166125], + }, + 2.8316255897272793, + ), + ( + _finite_pressure_current, + { + (0, 0, 0): [-1.0984223013409288e-11, -0.4764076511218376, 1.041130887765426], + (0, 2, 1): [-1.044435577486018e-14, -0.4447191502214616, 0.220359403292567], + (1, 2, 2): [1.2467274928318781, 1.4565818008885116, -0.5095116177456592], + }, + 2.031147679042703, + ), + ], +) +def test_field_hessian_matches_upstream_pyqsc( + factory, + expected_components, + expected_inverse_scale, +): + solution = factory() + indices = [0, 15, 30] + + for component, expected in expected_components.items(): + np.testing.assert_allclose( + np.asarray(solution.grad_grad_B)[indices, *component], + expected, + rtol=3e-6, + atol=2e-6, + ) + np.testing.assert_allclose( + solution.grad_grad_B_inverse_scale_length, + expected_inverse_scale, + rtol=3e-7, + ) + + +def test_field_hessian_resolution_convergence(): + medium = _vacuum_qa(nphi=61) + fine = _vacuum_qa(nphi=91) + + np.testing.assert_allclose( + np.asarray(medium.grad_grad_B)[0], + np.asarray(fine.grad_grad_B)[0], + rtol=2e-6, + atol=2e-6, + ) + np.testing.assert_allclose( + medium.grad_grad_B_inverse_scale_length_vs_varphi[0], + fine.grad_grad_B_inverse_scale_length_vs_varphi[0], + rtol=2e-6, + ) + + +def test_field_hessian_supports_jit_and_jvp(): + def component(etabar): + solution = qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=etabar, + B2c=-0.00322, + nphi=31, + order="r2", + ) + return solution.grad_grad_B_axis[7, 0, 1, 2] + + eager_value = component(0.64) + compiled_value = jax.jit(component)(0.64) + value, tangent = jax.jvp(component, (0.64,), (1.0,)) + + np.testing.assert_allclose(compiled_value, eager_value, rtol=2e-11, atol=2e-11) + np.testing.assert_allclose(value, eager_value, rtol=3e-12, atol=3e-12) + assert np.isfinite(tangent) diff --git a/tests/physics/test_first_order.py b/tests/physics/test_first_order.py new file mode 100644 index 0000000..c723e16 --- /dev/null +++ b/tests/physics/test_first_order.py @@ -0,0 +1,94 @@ +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.first_order import sigma_residual +from pyqsc_jax.models import NearAxisInputs + + +def standard_solution(nphi=31, **kwargs): + parameters = { + "axis": qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3), + "etabar": -0.9, + "nphi": nphi, + } + parameters.update(kwargs) + return qsc.solve(**parameters) + + +def test_sigma_residual_and_report(): + solution = standard_solution() + state = solution.sigma.at[0].set(solution.iota) + residual = sigma_residual( + state, + inputs=solution.inputs, + geometry=solution.geometry, + ) + + assert bool(solution.root_report.converged) + assert bool(solution.root_report.finite) + assert solution.root_report.residual_norm < 2e-13 + np.testing.assert_allclose(jnp.max(jnp.abs(residual)), solution.root_report.residual_norm) + np.testing.assert_allclose(solution.sigma[0], solution.inputs.sigma0) + + +def test_vacuum_field_gradient_satisfies_maxwell_identities_spectrally(): + solution = standard_solution(nphi=61) + gradient = solution.grad_B_axis + divergence = jnp.trace(gradient, axis1=-2, axis2=-1) + antisymmetric = gradient - jnp.swapaxes(gradient, -1, -2) + + np.testing.assert_allclose(jnp.linalg.norm(solution.B_axis, axis=-1), solution.inputs.B0) + assert jnp.max(jnp.abs(divergence)) < 2e-8 + assert jnp.max(jnp.abs(antisymmetric)) < 2e-8 + assert jnp.all(solution.L_grad_B > 0) + + +def test_first_order_shapes_and_coefficients(): + solution = standard_solution() + assert solution.axis is solution.inputs.axis + assert solution.B_axis.shape == (31, 3) + assert solution.grad_B_axis.shape == (31, 3, 3) + assert solution.sigma.shape == (31,) + np.testing.assert_allclose(solution.X1s, 0.0) + np.testing.assert_allclose( + solution.X1c * solution.curvature, + solution.inputs.etabar, + rtol=2e-13, + ) + area_jacobian = solution.X1c * solution.Y1s - solution.X1s * solution.Y1c + np.testing.assert_allclose( + area_jacobian, + solution.inputs.sG * solution.inputs.spsi, + rtol=2e-13, + ) + + +def test_invalid_order_and_inverse_mode_are_explicitly_rejected(): + with pytest.raises(ValueError, match="order must be"): + standard_solution(order="fourth") + with pytest.raises(ValueError, match="iota is required"): + standard_solution(solve_for="etabar") + + +@pytest.mark.parametrize( + "overrides, message", + [ + ({"nphi": 2}, "nphi"), + ({"nphi": True}, "nphi"), + ({"order": 4}, "order"), + ({"sG": 0}, "sG"), + ({"spsi": 0}, "spsi"), + ({"solve_for": "pressure"}, "solve_for"), + ({"etabar": jnp.ones(2)}, "etabar"), + ], +) +def test_input_model_validation(overrides, message): + parameters = { + "axis": qsc.Axis(rc=[1.0], zs=[0.0]), + "etabar": -0.9, + } + parameters.update(overrides) + with pytest.raises(ValueError, match=message): + NearAxisInputs(**parameters) diff --git a/tests/physics/test_geometry.py b/tests/physics/test_geometry.py new file mode 100644 index 0000000..44fc3ee --- /dev/null +++ b/tests/physics/test_geometry.py @@ -0,0 +1,78 @@ +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from pyqsc_jax import Axis +from pyqsc_jax.geometry import compute_axis_geometry + + +def test_circular_axis_geometry_is_analytic(): + major_radius = 2.0 + nphi = 31 + geometry = compute_axis_geometry(Axis(rc=[major_radius], zs=[0.0]), nphi=nphi) + + np.testing.assert_allclose(geometry.axis_length, 2 * jnp.pi * major_radius, rtol=2e-14) + np.testing.assert_allclose(geometry.curvature, 1 / major_radius, rtol=2e-14) + np.testing.assert_allclose(geometry.torsion, 0.0, atol=2e-14) + np.testing.assert_allclose(geometry.varphi, geometry.samples.phi, rtol=2e-14, atol=2e-14) + np.testing.assert_allclose( + geometry.tangent_cylindrical, + jnp.tile(jnp.array([0.0, 1.0, 0.0]), (nphi, 1)), + atol=2e-14, + ) + np.testing.assert_allclose( + geometry.normal_cylindrical, + jnp.tile(jnp.array([-1.0, 0.0, 0.0]), (nphi, 1)), + atol=2e-14, + ) + np.testing.assert_allclose( + geometry.binormal_cylindrical, + jnp.tile(jnp.array([0.0, 0.0, 1.0]), (nphi, 1)), + atol=2e-14, + ) + assert int(geometry.frame_helicity) == 0 + assert bool(geometry.diagnostics.frenet_valid) + assert bool(geometry.diagnostics.cylindrical_coordinates_valid) + + +def test_frame_is_right_handed_and_orthonormal(): + axis = Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) + geometry = compute_axis_geometry(axis, nphi=31) + frame = jnp.stack( + (geometry.tangent_cartesian, geometry.normal_cartesian, geometry.binormal_cartesian), + axis=-2, + ) + gram = frame @ jnp.swapaxes(frame, -1, -2) + + expected_gram = jnp.broadcast_to(jnp.eye(3), gram.shape) + np.testing.assert_allclose(gram, expected_gram, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(jnp.linalg.det(frame), 1.0, rtol=2e-13, atol=2e-13) + assert geometry.diagnostics.maximum_frame_orthogonality_error < 1e-12 + + +def test_invalid_axis_has_explicit_diagnostics(): + geometry = compute_axis_geometry(Axis(rc=[0.0], zs=[0.0]), nphi=15) + + assert not bool(geometry.diagnostics.frenet_valid) + assert not bool(geometry.diagnostics.cylindrical_coordinates_valid) + np.testing.assert_allclose(geometry.diagnostics.minimum_speed, 0.0) + + with pytest.raises(ValueError, match="nphi"): + compute_axis_geometry(Axis(rc=[1.0], zs=[0.0]), nphi=2) + + +def test_geometry_jit_and_axis_length_gradient(): + axis = Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) + compiled = jax.jit(compute_axis_geometry, static_argnames=("nphi",))(axis, nphi=31) + eager = compute_axis_geometry(axis, nphi=31) + np.testing.assert_allclose(compiled.curvature, eager.curvature, rtol=2e-13) + + def length(rc1): + varied = Axis(rc=jnp.array([1.0, rc1]), zs=axis.zs, nfp=axis.nfp) + return compute_axis_geometry(varied, nphi=31).axis_length + + derivative = jax.grad(length)(0.045) + step = 1e-5 + finite_difference = (length(0.045 + step) - length(0.045 - step)) / (2 * step) + np.testing.assert_allclose(derivative, finite_difference, rtol=2e-8, atol=2e-8) diff --git a/tests/physics/test_inverse.py b/tests/physics/test_inverse.py new file mode 100644 index 0000000..5ae5a14 --- /dev/null +++ b/tests/physics/test_inverse.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc + +AXIS = qsc.Axis(rc=[1.0, 0.045], zs=[0.0, -0.045], nfp=3) + + +def test_target_iota_etabar_round_trip_and_sign_branches(): + forward = qsc.solve(axis=AXIS, etabar=-0.9, nphi=31) + negative = qsc.solve( + axis=AXIS, + etabar=-1.0, + iota=forward.iota, + solve_for="etabar", + nphi=31, + ) + positive = qsc.solve( + axis=AXIS, + etabar=1.0, + iota=forward.iota, + solve_for="etabar", + nphi=31, + ) + + np.testing.assert_allclose(negative.inputs.etabar, -0.9, rtol=2.0e-12) + np.testing.assert_allclose(positive.inputs.etabar, 0.9, rtol=2.0e-12) + np.testing.assert_allclose(negative.iota, forward.iota) + np.testing.assert_allclose(negative.sigma, forward.sigma, atol=2.0e-12) + assert bool(negative.root_report.converged) + assert negative.root_report.residual_norm < 2.0e-13 + assert negative.inverse is not None + assert negative.parameter == "etabar" + assert int(negative.parameter_sign) == -1 + assert not bool(negative.branch_fold) + + +def test_target_iota_exposes_distinct_local_etabar_branches(): + target = qsc.solve(axis=AXIS, etabar=-0.9, nphi=31).iota + outer = qsc.solve( + axis=AXIS, + etabar=-1.0, + iota=target, + solve_for="etabar", + nphi=31, + ) + inner = qsc.solve( + axis=AXIS, + etabar=-0.8, + iota=target, + solve_for="etabar", + nphi=31, + ) + inner_forward = qsc.solve(axis=AXIS, etabar=inner.inputs.etabar, nphi=31) + + assert abs(float(outer.inputs.etabar - inner.inputs.etabar)) > 0.1 + np.testing.assert_allclose(inner.iota, target) + np.testing.assert_allclose(inner_forward.iota, target, atol=2.0e-13) + assert np.sign(float(outer.response_derivative)) != np.sign(float(inner.response_derivative)) + + +@pytest.mark.parametrize("actual_I2", [-0.5, 0.2, 1.0]) +def test_target_iota_I2_round_trip(actual_I2): + forward = qsc.solve(axis=AXIS, etabar=-0.9, I2=actual_I2, nphi=31) + inverse = qsc.solve( + axis=AXIS, + etabar=-0.9, + I2=0.0, + iota=forward.iota, + solve_for="I2", + nphi=31, + ) + + np.testing.assert_allclose(inverse.inputs.I2, actual_I2, rtol=2.0e-12, atol=2.0e-12) + np.testing.assert_allclose(inverse.iota, forward.iota) + assert inverse.parameter == "I2" + assert bool(inverse.root_report.converged) + + +def test_inverse_mode_propagates_to_second_order(): + forward = qsc.solve(axis=AXIS, etabar=-0.9, I2=0.2, nphi=31, order="r2") + inverse = qsc.solve( + axis=AXIS, + etabar=-1.0, + I2=0.2, + iota=forward.iota, + solve_for="etabar", + nphi=31, + order="r2", + ) + + np.testing.assert_allclose(inverse.inputs.etabar, -0.9, rtol=2.0e-12) + np.testing.assert_allclose(inverse.B20, forward.B20, rtol=3.0e-11, atol=3.0e-11) + assert bool(inverse.linear_report.converged) + + +def test_inverse_solve_is_jittable_and_implicitly_differentiable(): + def etabar_for_iota(target_iota): + return qsc.solve( + axis=AXIS, + etabar=-1.0, + iota=target_iota, + solve_for="etabar", + nphi=31, + ).inputs.etabar + + target = qsc.solve(axis=AXIS, etabar=-0.9, nphi=31).iota + solution = qsc.solve( + axis=AXIS, + etabar=-1.0, + iota=target, + solve_for="etabar", + nphi=31, + ) + value = etabar_for_iota(target) + jitted = jax.jit(etabar_for_iota)(target) + tangent = jax.jvp(etabar_for_iota, (target,), (jnp.asarray(1.0),))[1] + step = 1.0e-5 + finite_difference = (etabar_for_iota(target + step) - etabar_for_iota(target - step)) / ( + 2 * step + ) + + np.testing.assert_allclose(jitted, value, rtol=2.0e-12) + np.testing.assert_allclose(tangent, 1 / solution.response_derivative, rtol=2.0e-11) + np.testing.assert_allclose(tangent, finite_difference, rtol=2.0e-6) + + +def test_inverse_input_guards_and_forward_diagnostic_absence(): + forward = qsc.solve(axis=AXIS, etabar=-0.9, nphi=31) + with pytest.raises(AttributeError, match="Forward solution"): + _ = forward.target_iota + with pytest.raises(ValueError, match="etabar is required"): + qsc.solve(axis=AXIS) + with pytest.raises(ValueError, match="prescribed only"): + qsc.solve(axis=AXIS, etabar=-0.9, iota=0.4) + with pytest.raises(ValueError, match="iota is required"): + qsc.solve(axis=AXIS, etabar=-0.9, solve_for="I2") + with pytest.raises(ValueError, match="etabar is required"): + qsc.solve(axis=AXIS, iota=0.4, solve_for="I2") + with pytest.raises(ValueError, match="scalar"): + qsc.solve( + axis=AXIS, + etabar=-1.0, + iota=jnp.ones(2), + solve_for="etabar", + ) + with pytest.raises(ValueError, match="fold_tolerance"): + qsc.solve( + axis=AXIS, + etabar=-1.0, + iota=0.4, + solve_for="etabar", + fold_tolerance=-1.0, + ) diff --git a/tests/physics/test_mercier.py b/tests/physics/test_mercier.py new file mode 100644 index 0000000..03ac8d6 --- /dev/null +++ b/tests/physics/test_mercier.py @@ -0,0 +1,64 @@ +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +def test_mercier_diagnostics_match_upstream_finite_pressure_reference(): + solution = qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=0.9, + B2c=-0.7, + p2=-600000.0, + nphi=61, + order="r2", + ) + + expected = { + "d2_volume_d_psi2": -121.81059604927127, + "DGeod_times_r2": -0.1210731174740698, + "DWell_times_r2": 0.06028523900360924, + "DMerc_times_r2": -0.06078787847046056, + } + for name, value in expected.items(): + np.testing.assert_allclose(getattr(solution, name), value, rtol=3e-12, atol=3e-12) + np.testing.assert_allclose(getattr(solution.mercier, name), value, rtol=3e-12, atol=3e-12) + + +def test_vacuum_mercier_pressure_terms_vanish(): + solution = qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=0.64, + B2c=-0.00322, + nphi=61, + order="r2", + ) + + np.testing.assert_allclose(solution.d2_volume_d_psi2, 23.98845128286806, rtol=3e-12) + assert solution.DGeod_times_r2 == 0 + assert solution.DWell_times_r2 == 0 + assert solution.DMerc_times_r2 == 0 + + +def test_second_order_diagnostics_reject_first_order_solution(): + solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, + ) + + with pytest.raises(ValueError, match="second-order"): + qsc.mercier_diagnostics(solution) + with pytest.raises(ValueError, match="second-order"): + qsc.total_field_jet(solution) + with pytest.raises(AttributeError, match="Mercier"): + _ = solution.DMerc_times_r2 + with pytest.raises(AttributeError, match="second-derivative"): + _ = solution.grad_grad_B_axis diff --git a/tests/physics/test_plasma_current.py b/tests/physics/test_plasma_current.py new file mode 100644 index 0000000..0c5c1ff --- /dev/null +++ b/tests/physics/test_plasma_current.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.second_order import MU0 + + +def finite_current_solution(*, I2=0.9, p2=-600000.0, nphi=31): + return qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=I2, + p2=p2, + B2c=-0.7, + nphi=nphi, + order="r2", + ) + + +def test_covariant_and_enclosed_current_conversions_round_trip(): + I2 = 0.9 + radius = 0.08 + chi = -1 + current = qsc.enclosed_current_from_covariant( + I2, + formal_radius=radius, + chi=chi, + ) + + np.testing.assert_allclose(current, 2 * np.pi * chi * I2 * radius**2 / MU0) + np.testing.assert_allclose( + qsc.covariant_current_from_enclosed( + current, + formal_radius=radius, + chi=chi, + ), + I2, + ) + derivative = jax.grad( + lambda a: qsc.enclosed_current_from_covariant( + I2, + formal_radius=a, + chi=chi, + ) + )(radius) + np.testing.assert_allclose(derivative, 4 * np.pi * chi * I2 * radius / MU0) + + +def test_positive_volume_source_matches_manuscript_coefficients(): + solution = finite_current_solution() + radius = 0.07 + source = qsc.plasma_current_source(solution, formal_radius=radius) + chi = solution.inputs.sG * solution.inputs.spsi + expected_j = 2 * chi * solution.inputs.I2 + expected_C2 = solution.G2 + solution.N_helicity * solution.inputs.I2 + + np.testing.assert_allclose(source.parallel_current_mu0, expected_j) + np.testing.assert_allclose(source.C2, expected_C2) + np.testing.assert_allclose( + source.enclosed_toroidal_current, + np.pi * radius**2 * expected_j / MU0, + ) + np.testing.assert_allclose( + source.w1, + expected_j * solution.geometry.tangent_cartesian, + ) + np.testing.assert_allclose( + source.w2_cosine, + source.wstar2_cosine + - (expected_j * solution.geometry.curvature * solution.X1c)[:, None] + * solution.geometry.tangent_cartesian, + ) + np.testing.assert_allclose(source.w2_sine, source.wstar2_sine) + assert source.chi == chi + + +def test_weighted_source_evaluation_has_regular_radial_power_and_batch_shape(): + solution = finite_current_solution(nphi=15) + source = qsc.plasma_current_source(solution, formal_radius=0.05) + radial = jnp.asarray([0.0, 0.01]) + theta = jnp.asarray([0.2, 0.7]) + weighted = qsc.evaluate_weighted_current(source, radial, theta) + + assert weighted.shape == (2, 15, 3) + np.testing.assert_allclose(weighted[0], 0.0) + direct = source.axis_length_per_radian * ( + radial[1] * source.w1 + + radial[1] ** 2 * (np.cos(theta[1]) * source.w2_cosine + np.sin(theta[1]) * source.w2_sine) + ) + np.testing.assert_allclose(weighted[1], direct) + scalar = jax.jit(qsc.evaluate_weighted_current)(source, 0.01, 0.3) + assert scalar.shape == (15, 3) + + +def test_pressure_and_parallel_current_paths_remain_well_defined_at_I2_zero(): + pressure_only = finite_current_solution(I2=0.0) + source = qsc.plasma_current_source(pressure_only, formal_radius=0.06) + + np.testing.assert_allclose(source.parallel_current_mu0, 0.0) + np.testing.assert_allclose(source.w1, 0.0) + assert np.max(np.abs(source.wstar2_cosine)) > 0 + assert np.max(np.abs(source.wstar2_sine)) > 0 + assert np.all(np.isfinite(source.wstar2_cosine)) + + vacuum = finite_current_solution(I2=0.0, p2=0.0) + vacuum_source = qsc.plasma_current_source(vacuum, formal_radius=0.06) + np.testing.assert_allclose(vacuum_source.w1, 0.0) + np.testing.assert_allclose(vacuum_source.wstar2_cosine, 0.0, atol=1.0e-14) + np.testing.assert_allclose(vacuum_source.wstar2_sine, 0.0, atol=1.0e-14) + + +def test_current_source_guards(): + with pytest.raises(ValueError, match="second-order"): + qsc.plasma_current_source( + qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + order="r1", + ), + formal_radius=0.1, + ) + for value in (0.0, -0.1): + with pytest.raises(ValueError, match="positive"): + qsc.enclosed_current_from_covariant( + 1.0, + formal_radius=value, + chi=1, + ) + with pytest.raises(ValueError, match="scalar"): + qsc.covariant_current_from_enclosed( + 1.0, + formal_radius=[0.1], + chi=1, + ) + for function in ( + qsc.enclosed_current_from_covariant, + qsc.covariant_current_from_enclosed, + ): + with pytest.raises(ValueError, match="chi"): + if function is qsc.enclosed_current_from_covariant: + function(1.0, formal_radius=0.1, chi=0) + else: + function(1.0, formal_radius=0.1, chi=0) diff --git a/tests/physics/test_plasma_field.py b/tests/physics/test_plasma_field.py new file mode 100644 index 0000000..fdc8f43 --- /dev/null +++ b/tests/physics/test_plasma_field.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import jax +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +def circular_solution(*, I2=0.1, p2=0.0, nphi=31): + return qsc.Qsc( + rc=[1.0], + zs=[0.0], + nfp=1, + etabar=1.0, + I2=I2, + p2=p2, + nphi=nphi, + order="r2", + ) + + +def finite_current_solution(*, I2=0.9, p2=-600000.0, nphi=31): + return qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=I2, + p2=p2, + B2c=-0.7, + nphi=nphi, + order="r2", + ) + + +def test_periodic_regularized_integral_and_circular_local_induction_limit(): + solution = circular_solution() + radius = 0.05 + integral = qsc.regularized_axis_integral(solution) + plasma = qsc.plasma_field_on_axis(solution, formal_radius=radius) + binormal_kernel = np.sum( + np.asarray(plasma.matched_axis_and_core) * np.asarray(solution.geometry.binormal_cartesian), + axis=-1, + ) + + assert np.max(np.abs(integral)) < 1.0e-13 + np.testing.assert_allclose(plasma.core_binormal_constant, 0.0) + np.testing.assert_allclose(plasma.core_normal_constant, 0.0) + np.testing.assert_allclose(binormal_kernel, np.log(8 / radius), rtol=2.0e-13) + np.testing.assert_allclose( + solution.inputs.I2 * radius**2 / 2 * binormal_kernel, + solution.inputs.I2 * radius**2 / 2 * np.log(8 / radius), + ) + + +def test_matching_reference_length_cancels_exactly(): + solution = finite_current_solution() + source = qsc.plasma_current_source(solution, formal_radius=0.06) + integral = qsc.regularized_axis_integral(solution) + first = qsc.matched_plasma_field_kernel( + solution, + source, + integral, + reference_length=2.3, + ) + second = qsc.matched_plasma_field_kernel( + solution, + source, + integral, + reference_length=7.1, + ) + plasma = qsc.plasma_field_on_axis(solution, formal_radius=0.06) + + np.testing.assert_allclose(first, second, rtol=0, atol=2.0e-15) + assert plasma.maximum_matching_scale_error < 2.0e-15 + np.testing.assert_allclose(plasma.matched_axis_and_core, first, atol=2.0e-15) + + +def test_shape_correction_and_full_field_converge_in_angular_resolution(): + solution = finite_current_solution(nphi=15) + coarse = qsc.plasma_field_on_axis( + solution, + formal_radius=0.05, + angular_resolution=64, + ) + fine = qsc.plasma_field_on_axis( + solution, + formal_radius=0.05, + angular_resolution=128, + ) + + np.testing.assert_allclose( + coarse.second_order_shape_correction, + fine.second_order_shape_correction, + rtol=2.0e-9, + atol=5.0e-15, + ) + np.testing.assert_allclose(coarse.field, fine.field, rtol=2.0e-9, atol=5.0e-15) + assert fine.formal_radius_to_curvature_radius > 0 + assert fine.estimated_field_remainder > 0 + assert fine.angular_resolution == 128 + + +def test_vacuum_and_pressure_only_limits(): + vacuum = qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=0.64, + I2=0.0, + p2=0.0, + B2c=-0.00322, + nphi=31, + order="r2", + ) + vacuum_field = qsc.plasma_field_on_axis(vacuum, formal_radius=0.05) + np.testing.assert_allclose(vacuum_field.field, 0.0, atol=1.0e-14) + np.testing.assert_allclose( + vacuum_field.second_order_shape_correction, + 0.0, + atol=1.0e-14, + ) + + pressure_only = finite_current_solution(I2=0.0) + pressure_field = qsc.plasma_field_on_axis( + pressure_only, + formal_radius=0.05, + ) + assert np.max(np.abs(pressure_field.field)) > 0 + np.testing.assert_allclose( + pressure_field.field, + pressure_field.second_order_shape_correction, + ) + + +def test_field_is_jittable_and_radius_derivative_matches_finite_difference(): + solution = finite_current_solution(nphi=15) + + def field_component(radius): + return qsc.plasma_field_on_axis( + solution, + formal_radius=radius, + angular_resolution=16, + ).field[0, 2] + + radius = 0.05 + value = field_component(radius) + jitted = jax.jit(field_component)(radius) + tangent = jax.jvp(field_component, (radius,), (1.0,))[1] + step = 1.0e-5 + finite_difference = (field_component(radius + step) - field_component(radius - step)) / ( + 2 * step + ) + + np.testing.assert_allclose(jitted, value, rtol=2.0e-13) + np.testing.assert_allclose(tangent, finite_difference, rtol=2.0e-7) + + +def test_regularized_integral_resolution_convergence_for_nonplanar_axis(): + coarse = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, + order="r1", + ) + fine = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=61, + order="r1", + ) + coarse_integral = qsc.regularized_axis_integral(coarse) + fine_integral = qsc.regularized_axis_integral(fine) + + np.testing.assert_allclose( + coarse_integral[0], + fine_integral[0], + rtol=2.0e-3, + atol=2.0e-5, + ) + assert np.all(np.isfinite(coarse_integral)) + + +@pytest.mark.physics +def test_documented_plasma_dominant_case_exceeds_thirty_percent(): + solution = qsc.solve_configuration("plasma_dominant_channel", nphi=61) + formal_radius = 0.2 + plasma = qsc.plasma_field_on_axis( + solution, + formal_radius=formal_radius, + angular_resolution=64, + ) + plasma_norm = np.linalg.norm(np.asarray(plasma.field), axis=-1) + total_norm = np.linalg.norm(np.asarray(solution.B_axis), axis=-1) + fraction = plasma_norm / total_norm + + assert np.min(fraction) > 0.30 + np.testing.assert_allclose(np.mean(fraction), 0.3325737905856357, rtol=2.0e-12) + assert formal_radius < float(solution.r_singularity) + np.testing.assert_allclose( + plasma.current_source.enclosed_toroidal_current, + 840000.0, + rtol=2.0e-13, + ) + + +@pytest.mark.physics +def test_documented_plasma_stellarator_is_pressure_only_nonplanar_and_angle_dependent(): + solution = qsc.solve_configuration("plasma_stellarator", nphi=121) + formal_radius = 0.15 + plasma = qsc.plasma_field_on_axis( + solution, + formal_radius=formal_radius, + angular_resolution=128, + ) + plasma_norm = np.linalg.norm(np.asarray(plasma.field), axis=-1) + total_norm = np.linalg.norm(np.asarray(solution.B_axis), axis=-1) + fraction = plasma_norm / total_norm + peak_to_peak_over_mean = np.ptp(plasma_norm) / np.mean(plasma_norm) + torsion_rms = np.sqrt(np.mean(np.asarray(solution.torsion) ** 2)) + + assert float(solution.inputs.I2) == 0.0 + assert float(solution.inputs.p2) != 0.0 + assert torsion_rms > 0.5 + assert np.min(fraction) > 1.7e-3 + assert peak_to_peak_over_mean > 0.1 + assert formal_radius / float(solution.r_singularity) < 0.4 + assert qsc.Criteria.from_curvo_2025(minimum_abs_iota=0.4).evaluate(solution).passed + np.testing.assert_allclose(np.mean(fraction), 0.0018895978020422664, rtol=3.0e-12) + + +def test_plasma_field_guards(): + solution = circular_solution() + with pytest.raises(ValueError, match="angular_resolution"): + qsc.plasma_field_on_axis( + solution, + formal_radius=0.05, + angular_resolution=7, + ) + source = qsc.plasma_current_source(solution, formal_radius=0.05) + with pytest.raises(ValueError, match="reference_length"): + qsc.matched_plasma_field_kernel( + solution, + source, + qsc.regularized_axis_integral(solution), + reference_length=0.0, + ) diff --git a/tests/physics/test_plasma_gradient.py b/tests/physics/test_plasma_gradient.py new file mode 100644 index 0000000..c3e45d5 --- /dev/null +++ b/tests/physics/test_plasma_gradient.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +def finite_current_solution(*, I2=0.9, p2=-600000.0, nphi=61): + return qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=I2, + p2=p2, + B2c=-0.7, + nphi=nphi, + order="r2", + ) + + +def test_straight_circular_channel_gradient(): + current_density_mu0 = 1.7 + gradient = qsc.elliptical_channel_gradient( + 1.0, + 0.0, + parallel_current_mu0=current_density_mu0, + chi=1, + ) + expected = np.asarray( + [ + [0.0, 0.0, 0.0], + [0.0, 0.0, -current_density_mu0 / 2], + [0.0, current_density_mu0 / 2, 0.0], + ] + ) + + np.testing.assert_allclose(gradient, expected) + np.testing.assert_allclose(np.trace(gradient), 0.0) + np.testing.assert_allclose( + gradient[2, 1] - gradient[1, 2], + current_density_mu0, + ) + + +def test_straight_sheared_elliptical_channel_obeys_ampere_and_divergence(): + current_density_mu0 = -0.8 + gradient = qsc.elliptical_channel_gradient( + 1.6, + -0.35, + parallel_current_mu0=current_density_mu0, + chi=-1, + ) + + np.testing.assert_allclose(np.trace(gradient), 0.0, atol=1.0e-15) + np.testing.assert_allclose( + gradient[2, 1] - gradient[1, 2], + current_density_mu0, + ) + assert gradient[1, 1] == -gradient[2, 2] + assert not np.isclose(gradient[2, 1], -gradient[1, 2]) + + +def test_gradient_rotates_covariantly_from_frenet_to_cartesian(): + angle = 0.37 + rotation = jnp.asarray( + [ + [1.0, 0.0, 0.0], + [0.0, np.cos(angle), np.sin(angle)], + [0.0, -np.sin(angle), np.cos(angle)], + ] + ) + frenet = qsc.elliptical_channel_gradient( + 1.3, + 0.2, + parallel_current_mu0=0.7, + chi=1, + ) + cartesian = qsc.elliptical_channel_gradient( + 1.3, + 0.2, + parallel_current_mu0=0.7, + chi=1, + frame=rotation, + ) + + np.testing.assert_allclose( + cartesian, + rotation.T @ frenet @ rotation, + atol=2.0e-15, + ) + + +def test_external_gradient_is_symmetric_trace_free_and_ampere_cancels(): + solution = finite_current_solution() + result = qsc.plasma_gradient_on_axis( + solution, + formal_radius=0.05, + ) + + assert result.maximum_divergence < 2.0e-15 + assert result.maximum_ampere_error < 2.0e-15 + assert result.maximum_external_asymmetry < 5.0e-9 + assert result.maximum_external_trace < 2.0e-9 + np.testing.assert_allclose( + result.external_gradient, + result.external_gradient_stf, + atol=2.0e-9, + ) + np.testing.assert_allclose( + qsc.unpack_symmetric_trace_free_rank2(result.external_gradient_independent), + result.external_gradient_stf, + atol=2.0e-15, + ) + np.testing.assert_allclose( + result.external_field, + solution.B_axis - result.field.field, + ) + + +def test_vacuum_reduction_leaves_total_field_jet_unchanged(): + solution = qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=0.64, + I2=0.0, + p2=0.0, + B2c=-0.00322, + nphi=31, + order="r2", + ) + result = qsc.plasma_gradient_on_axis( + solution, + formal_radius=0.05, + ) + + np.testing.assert_allclose(result.field.field, 0.0, atol=1.0e-14) + np.testing.assert_allclose(result.gradient, 0.0) + np.testing.assert_allclose(result.external_field, solution.B_axis) + np.testing.assert_allclose( + result.external_gradient, + solution.grad_B_axis, + ) + + +def test_stf_rank_two_pack_unpack_and_projection(): + tensor = jnp.asarray( + [ + [2.0, 1.0, -3.0], + [3.0, -1.0, 4.0], + [5.0, 2.0, 7.0], + ] + ) + projected = qsc.project_symmetric_trace_free_rank2(tensor) + packed = qsc.pack_symmetric_trace_free_rank2(tensor) + unpacked = qsc.unpack_symmetric_trace_free_rank2(packed) + + np.testing.assert_allclose(projected, projected.T) + np.testing.assert_allclose(np.trace(projected), 0.0, atol=1.0e-15) + np.testing.assert_allclose(unpacked, projected) + assert packed.shape == (5,) + + +def test_elliptical_gradient_is_jittable_and_differentiable(): + def component(x): + return qsc.elliptical_channel_gradient( + x, + 0.2, + parallel_current_mu0=0.7, + chi=1, + )[1, 2] + + x = 1.3 + value = component(x) + jitted = jax.jit(component)(x) + tangent = jax.jvp(component, (x,), (1.0,))[1] + step = 1.0e-5 + finite_difference = (component(x + step) - component(x - step)) / (2 * step) + + np.testing.assert_allclose(jitted, value) + np.testing.assert_allclose(tangent, finite_difference, rtol=2.0e-10) + + +def test_gradient_and_stf_guards(): + with pytest.raises(ValueError, match="chi"): + qsc.elliptical_channel_gradient( + 1.0, + 0.0, + parallel_current_mu0=1.0, + chi=0, + ) + with pytest.raises(ValueError, match="frame"): + qsc.elliptical_channel_gradient( + 1.0, + 0.0, + parallel_current_mu0=1.0, + chi=1, + frame=[1.0, 2.0], + ) + with pytest.raises(ValueError, match="shape"): + qsc.project_symmetric_trace_free_rank2(jnp.ones((2, 2))) + with pytest.raises(ValueError, match="length 5"): + qsc.unpack_symmetric_trace_free_rank2(jnp.ones(4)) diff --git a/tests/physics/test_plasma_hessian.py b/tests/physics/test_plasma_hessian.py new file mode 100644 index 0000000..ed0353e --- /dev/null +++ b/tests/physics/test_plasma_hessian.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +from dataclasses import replace +from itertools import permutations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.plasma import _principal_transverse_plasma_hessian + + +def finite_pressure_current_solution(*, nphi=61, I2=0.9): + return qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=I2, + p2=-600000.0, + B2c=-0.7, + nphi=nphi, + order="r2", + ) + + +def vacuum_qa_solution(*, nphi=61): + return qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=0.64, + I2=0.0, + p2=0.0, + B2c=-0.00322, + nphi=nphi, + order="r2", + ) + + +def test_stf_rank_three_pack_unpack_and_projection(): + tensor = jnp.arange(54.0).reshape(2, 3, 3, 3) + projected = qsc.project_symmetric_trace_free_rank3(tensor) + packed = qsc.pack_symmetric_trace_free_rank3(tensor) + unpacked = qsc.unpack_symmetric_trace_free_rank3(packed) + + for permutation in permutations((1, 2, 3)): + np.testing.assert_allclose( + projected, + jnp.transpose(projected, (0, *permutation)), + atol=1.0e-14, + ) + np.testing.assert_allclose( + jnp.einsum("...iik->...k", projected), + 0.0, + atol=3.0e-14, + ) + np.testing.assert_allclose(unpacked, projected, atol=2.0e-14) + assert packed.shape == (2, 7) + + +def test_stf_rank_three_helpers_are_jittable(): + tensor = jnp.arange(27.0).reshape(3, 3, 3) + + eager = qsc.unpack_symmetric_trace_free_rank3(qsc.pack_symmetric_trace_free_rank3(tensor)) + compiled = jax.jit( + lambda value: qsc.unpack_symmetric_trace_free_rank3( + qsc.pack_symmetric_trace_free_rank3(value) + ) + )(tensor) + + np.testing.assert_allclose(compiled, eager) + + +def test_external_hessian_is_fully_symmetric_and_trace_free(): + solution = finite_pressure_current_solution(nphi=121) + formal_radius = 0.05 + result = qsc.plasma_hessian_on_axis( + solution, + formal_radius=0.05, + ) + + assert result.maximum_derivative_asymmetry < 2.0e-15 + assert result.maximum_external_symmetry_error < 8.0e-11 + assert result.maximum_external_trace < 2.0e-10 + np.testing.assert_allclose( + result.external_hessian, + result.external_hessian_stf, + atol=2.0e-10, + ) + np.testing.assert_allclose( + qsc.unpack_symmetric_trace_free_rank3(result.external_hessian_independent), + result.external_hessian_stf, + atol=2.0e-15, + ) + axis_scale = solution.geometry.abs_G0_over_B0 + expected_remainder = ( + jnp.max(jnp.abs(result.hessian)) + * (formal_radius / axis_scale) ** 2 + * (1 + jnp.abs(jnp.log(formal_radius / axis_scale))) + ) + np.testing.assert_allclose(result.estimated_hessian_remainder, expected_remainder) + + +def test_qh_external_hessian_preserves_oriented_ellipse_topology(): + solution = qsc.Qsc( + rc=[1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + zs=[0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + nfp=4, + etabar=1.569, + I2=0.2, + p2=-100000.0, + B2c=0.1348, + nphi=121, + order="r2", + ) + result = qsc.plasma_hessian_on_axis( + solution, + formal_radius=0.04, + ) + + assert int(solution.helicity) == 1 + assert result.maximum_derivative_asymmetry < 2.0e-15 + assert result.maximum_external_symmetry_error < 2.0e-9 + assert result.maximum_external_trace < 3.0e-9 + + +def test_circular_curved_channel_recovers_finite_conductor_limit(): + """Equation (215), with affine and second-order shape terms suppressed.""" + + solution = qsc.Qsc( + rc=[1.0], + zs=[0.0], + nfp=1, + etabar=1.0, + I2=0.1, + nphi=31, + order="r2", + ) + zeros = jnp.zeros_like(solution.X20) + circular_second_order = replace( + solution.second_order, + X20=zeros, + X2c=zeros, + X2s=zeros, + Y20=zeros, + Y2c=zeros, + Y2s=zeros, + Z20=zeros, + Z2c=zeros, + Z2s=zeros, + ) + curvature_only_solution = replace( + solution, + second_order=circular_second_order, + ) + source = qsc.plasma_current_source( + solution, + formal_radius=0.1, + ) + zero_vector = jnp.zeros_like(source.wstar2_cosine) + curvature_only_source = replace( + source, + wstar2_cosine=zero_vector, + wstar2_sine=zero_vector, + ) + + transverse, _ = _principal_transverse_plasma_hessian( + curvature_only_solution, + curvature_only_source, + ) + coefficient = source.parallel_current_mu0 * solution.geometry.curvature / 8 + expected = jnp.zeros_like(transverse) + expected = expected.at[:, 0, 0, 2].set(-coefficient) + expected = expected.at[:, 0, 1, 1].set(-coefficient) + expected = expected.at[:, 1, 0, 1].set(-coefficient) + expected = expected.at[:, 1, 1, 2].set(-3 * coefficient) + + np.testing.assert_allclose(transverse, expected, atol=2.0e-15) + + +def test_vacuum_reduction_leaves_total_hessian_unchanged(): + solution = vacuum_qa_solution() + result = qsc.plasma_hessian_on_axis( + solution, + formal_radius=0.05, + ) + + np.testing.assert_allclose(result.field.field.field, 0.0, atol=1.0e-14) + np.testing.assert_allclose(result.field.gradient, 0.0) + np.testing.assert_allclose(result.hessian, 0.0) + np.testing.assert_allclose( + result.external_hessian, + solution.grad_grad_B_axis, + ) + + +def test_plasma_hessian_converges_spectrally(): + medium = qsc.plasma_hessian_on_axis( + finite_pressure_current_solution(nphi=61), + formal_radius=0.05, + ) + fine = qsc.plasma_hessian_on_axis( + finite_pressure_current_solution(nphi=121), + formal_radius=0.05, + ) + + np.testing.assert_allclose( + medium.hessian_frenet[0], + fine.hessian_frenet[0], + atol=1.0e-8, + rtol=1.0e-8, + ) + assert fine.maximum_external_symmetry_error < 1.0e-4 * (medium.maximum_external_symmetry_error) + assert fine.maximum_external_trace < 1.0e-4 * medium.maximum_external_trace + + +def test_plasma_hessian_supports_jit_and_jvp(): + def component(I2): + result = qsc.plasma_hessian_on_axis( + finite_pressure_current_solution(nphi=15, I2=I2), + formal_radius=0.05, + angular_resolution=32, + ) + return result.hessian[4, 0, 1, 2] + + value = component(0.9) + compiled = jax.jit(component)(0.9) + tangent = jax.jvp(component, (0.9,), (1.0,))[1] + step = 1.0e-5 + finite_difference = (component(0.9 + step) - component(0.9 - step)) / (2 * step) + + np.testing.assert_allclose(compiled, value, rtol=2.0e-12, atol=2.0e-12) + np.testing.assert_allclose(tangent, finite_difference, rtol=2.0e-8, atol=2.0e-8) + + +def test_rank_three_guards(): + with pytest.raises(ValueError, match="shape"): + qsc.project_symmetric_trace_free_rank3(jnp.ones((3, 3))) + with pytest.raises(ValueError, match="length 7"): + qsc.unpack_symmetric_trace_free_rank3(jnp.ones(6)) diff --git a/tests/physics/test_second_order.py b/tests/physics/test_second_order.py new file mode 100644 index 0000000..939571d --- /dev/null +++ b/tests/physics/test_second_order.py @@ -0,0 +1,138 @@ +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.second_order import MU0 + + +def vacuum_qa(nphi=31, **kwargs): + parameters = { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + "nphi": nphi, + "order": "r2", + } + parameters.update(kwargs) + return qsc.Qsc(**parameters) + + +def test_complete_second_order_equations_and_linear_report(): + solution = vacuum_qa() + residuals = qsc.second_order_residuals(solution) + + assert solution.second_order is not None + assert bool(solution.linear_report.converged) + assert bool(solution.linear_report.finite) + assert bool(solution.linear_report.well_conditioned) + assert solution.linear_report.relative_residual_norm < 3e-14 + assert residuals.maximum_absolute < 1e-11 + for residual in ( + residuals.force_balance_1, + residuals.force_balance_2, + residuals.area_constraint_1, + residuals.area_constraint_2, + ): + assert residual.shape == (solution.inputs.nphi,) + + +def test_second_order_definitions_and_derivatives(): + solution = vacuum_qa() + weights = solution.geometry.d_l_d_phi + weighted_mean = jnp.sum(solution.B20 * weights) / jnp.sum(weights) + + np.testing.assert_allclose(solution.B20_mean, weighted_mean, rtol=2e-13) + np.testing.assert_allclose(solution.B20_anomaly, solution.B20 - weighted_mean) + np.testing.assert_allclose( + solution.B20_variation, + jnp.max(solution.B20) - jnp.min(solution.B20), + ) + np.testing.assert_allclose( + solution.d_X20_d_varphi, + solution.geometry.d_d_varphi @ solution.X20, + ) + np.testing.assert_allclose( + solution.d2_X1c_d_varphi2, + solution.geometry.d_d_varphi @ (solution.geometry.d_d_varphi @ solution.X1c), + ) + + +def test_finite_pressure_and_current_relations(): + solution = qsc.Qsc( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=0.9, + B2c=-0.7, + p2=-600000.0, + nphi=31, + order="r2", + ) + expected_beta_1s = ( + -4 + * solution.inputs.sG + * solution.inputs.spsi + * MU0 + * solution.inputs.p2 + * solution.inputs.etabar + * solution.geometry.abs_G0_over_B0 + / (solution.iotaN * solution.inputs.B0**2) + ) + expected_G2 = ( + -MU0 * solution.inputs.p2 * solution.G0 / solution.inputs.B0**2 + - solution.iota * solution.inputs.I2 + ) + + np.testing.assert_allclose(solution.beta_1s, expected_beta_1s) + np.testing.assert_allclose(solution.G2, expected_G2) + assert qsc.second_order_residuals(solution).maximum_absolute < 1e-11 + + +def test_qh_untwisting_preserves_harmonic_norm(): + solution = qsc.Qsc( + rc=[1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + zs=[0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + nfp=4, + etabar=1.569, + B2c=0.1348, + nphi=31, + order="r2", + ) + + assert int(solution.helicity) == 1 + np.testing.assert_allclose(solution.X20_untwisted, solution.X20) + np.testing.assert_allclose( + solution.X2s_untwisted**2 + solution.X2c_untwisted**2, + solution.X2s**2 + solution.X2c**2, + rtol=2e-13, + atol=2e-13, + ) + assert qsc.second_order_residuals(solution).maximum_absolute < 2e-11 + + +def test_second_order_residual_requires_r2_solution(): + first_order = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=15, + ) + + assert not hasattr(first_order, "X20") + with pytest.raises(AttributeError, match="no 'X20'"): + _ = first_order.X20 + with pytest.raises(ValueError, match="second-order"): + qsc.second_order_residuals(first_order) + + +def test_second_order_resolution_convergence(): + medium = vacuum_qa(nphi=31) + fine = vacuum_qa(nphi=61) + + np.testing.assert_allclose(medium.B20_mean, fine.B20_mean, rtol=6e-7) + np.testing.assert_allclose(medium.B20_residual, fine.B20_residual, rtol=2e-7) diff --git a/tests/physics/test_shear.py b/tests/physics/test_shear.py new file mode 100644 index 0000000..f839113 --- /dev/null +++ b/tests/physics/test_shear.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.near_axis import near_axis + +PAPER_CASES = ( + ( + { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + }, + 4.665537859250996, + ), + ( + { + "rc": [1.0, 0.09], + "zs": [0.0, -0.09], + "nfp": 2, + "etabar": 0.95, + "I2": 0.9, + "B2c": -0.7, + "p2": -600000.0, + }, + 0.2758333187134502, + ), + ( + { + "rc": [1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + "zs": [0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + "nfp": 4, + "etabar": 1.569, + "B2c": 0.1348, + }, + -1.1837987800965397, + ), +) + + +def qa_solution(nphi=61, **overrides): + parameters = { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + "order": "r3", + "nphi": nphi, + } + parameters.update(overrides) + return qsc.Qsc(**parameters) + + +@pytest.mark.parametrize(("parameters", "expected_iota2"), PAPER_CASES) +def test_magnetic_shear_matches_upstream_paper_cases(parameters, expected_iota2): + solution = qsc.solve_magnetic_shear(qsc.Qsc(**parameters, order="r3", nphi=61)) + + np.testing.assert_allclose(solution.iota2, expected_iota2, rtol=2.0e-12, atol=2.0e-12) + assert bool(solution.stellarator_symmetric) + np.testing.assert_allclose( + solution.iota2, + solution.inputs.B0 * solution.numerator / (2 * solution.denominator), + ) + assert np.all(np.isfinite(np.asarray(solution.Lambda_tilde))) + assert np.all(np.isfinite(np.asarray(solution.integrating_factor))) + + +def test_general_asymmetric_integration_matches_upstream(): + solution = qsc.Qsc( + rc=[1.0, 0.1], + rs=[0.0, 0.01], + zc=[0.0, -0.005], + zs=[0.0, 0.1], + nfp=2, + etabar=1.0, + sigma0=0.2, + I2=1.0, + order="r3", + nphi=61, + ) + solution = qsc.solve_magnetic_shear(solution) + + np.testing.assert_allclose(solution.iota2, -2131.8417260171145, rtol=3.0e-12) + assert not bool(solution.stellarator_symmetric) + assert abs(float(solution.sigma_average)) > 0.1 + + +def test_B31c_response_and_resolution_convergence(): + coarse = qsc.solve_magnetic_shear(qa_solution(nphi=31), B31c=0.23) + medium = qsc.solve_magnetic_shear(qa_solution(nphi=51), B31c=0.23) + fine = qsc.solve_magnetic_shear(qa_solution(nphi=91), B31c=0.23) + + np.testing.assert_allclose(medium.iota2, fine.iota2, rtol=5.0e-11) + np.testing.assert_allclose(fine.iota2, 4.676607885614116, rtol=3.0e-12) + assert abs(float(coarse.iota2 - fine.iota2)) < 1.2e-5 + np.testing.assert_allclose(fine.B31c, 0.23) + + +def test_magnetic_shear_is_jittable_and_differentiable(): + solution = qa_solution(nphi=31) + + def shear_for_B31c(B31c): + return qsc.solve_magnetic_shear(solution, B31c=B31c).iota2 + + value = shear_for_B31c(jnp.asarray(0.1)) + jitted = jax.jit(shear_for_B31c)(jnp.asarray(0.1)) + tangent = jax.jvp(shear_for_B31c, (jnp.asarray(0.1),), (jnp.asarray(1.0),))[1] + cotangent = jax.vjp(shear_for_B31c, jnp.asarray(0.1))[1](jnp.asarray(1.0))[0] + step = 1.0e-4 + finite_difference = (shear_for_B31c(0.1 + step) - shear_for_B31c(0.1 - step)) / (2 * step) + + np.testing.assert_allclose(jitted, value, rtol=2.0e-13) + np.testing.assert_allclose(tangent, cotangent, rtol=2.0e-12, atol=2.0e-12) + np.testing.assert_allclose(tangent, finite_difference, rtol=2.0e-8, atol=2.0e-10) + + +def test_magnetic_shear_differentiates_through_near_axis_solve(): + def shear_for_etabar(etabar): + solution = qa_solution(nphi=31, etabar=etabar) + return qsc.solve_magnetic_shear(solution).iota2 + + etabar = jnp.asarray(0.64) + tangent = jax.jvp(shear_for_etabar, (etabar,), (jnp.asarray(1.0),))[1] + step = 1.0e-5 + finite_difference = (shear_for_etabar(0.64 + step) - shear_for_etabar(0.64 - step)) / (2 * step) + + np.testing.assert_allclose(tangent, finite_difference, rtol=8.0e-8, atol=5.0e-8) + + +def test_legacy_calculate_shear_and_stale_value_removal(): + adapter = near_axis( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=0.64, + B2c=-0.00322, + order="r3", + nphi=61, + ) + assert adapter.calculate_shear() is None + np.testing.assert_allclose(adapter.iota2, 4.665537859250996, rtol=2.0e-12) + adapter.dofs = adapter.dofs.at[-1].set(0.65) + assert not hasattr(adapter, "iota2") + + +def test_magnetic_shear_input_guards_and_missing_result(): + with pytest.raises(ValueError, match="second-order"): + qsc.solve_magnetic_shear(qa_solution(order="r1")) + with pytest.raises(ValueError, match="scalar"): + qsc.solve_magnetic_shear(qa_solution(), B31c=jnp.ones(2)) + with pytest.raises(NotImplementedError, match="sG=spsi=1"): + qsc.solve_magnetic_shear(qa_solution(sG=-1)) + with pytest.raises(AttributeError, match="not been calculated"): + _ = qa_solution().iota2 diff --git a/tests/physics/test_singularity.py b/tests/physics/test_singularity.py new file mode 100644 index 0000000..eed476f --- /dev/null +++ b/tests/physics/test_singularity.py @@ -0,0 +1,177 @@ +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.near_axis import near_axis + +CASES = [ + ( + { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + }, + 0.2257896241404959, + [0.7304158315686611, 0.2505553226671735, 0.2935268282653716, 0.23909537193609554], + ), + ( + { + "rc": [1.0, 0.09], + "zs": [0.0, -0.09], + "nfp": 2, + "etabar": 0.95, + "I2": 0.9, + "B2c": -0.7, + "p2": -600000.0, + }, + 0.22161114270880408, + [0.7427937841922576, 0.3271954156702964, 0.2659539057552679, 0.3034006922051161], + ), + ( + { + "rc": [1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + "zs": [0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + "nfp": 4, + "etabar": 1.569, + "B2c": 0.1348, + }, + 0.3583276532138262, + [0.8681431131590702, 0.4034129514765815, 0.3728231585587465, 0.39141407123791455], + ), +] + + +@pytest.mark.parametrize("parameters, expected_minimum, expected_samples", CASES) +def test_singular_radius_matches_upstream_pyqsc( + parameters, + expected_minimum, + expected_samples, +): + solution = qsc.Qsc(**parameters, nphi=61, order="r2") + + np.testing.assert_allclose(solution.r_singularity, expected_minimum, rtol=0, atol=5e-8) + np.testing.assert_allclose( + np.asarray(solution.r_singularity_vs_varphi)[[0, 15, 30, 45]], + expected_samples, + rtol=0, + atol=5e-8, + ) + np.testing.assert_allclose( + solution.inv_r_singularity_vs_varphi, + 1 / solution.r_singularity_vs_varphi, + ) + np.testing.assert_allclose( + solution.r_singularity_basic_vs_varphi, + solution.r_singularity_vs_varphi, + ) + + +def test_determinant_coefficients_and_refined_residual(): + solution = qsc.Qsc(**CASES[0][0], nphi=61, order="r2") + diagnostics = solution.singularity + expected_g0 = solution.geometry.abs_G0_over_B0 * solution.X1c * solution.Y1s + + np.testing.assert_allclose(diagnostics.g0, expected_g0, rtol=2e-13, atol=2e-13) + assert jnp.max(jnp.abs(diagnostics.g1s)) < 2e-13 + assert diagnostics.maximum_residual_norm < 2e-13 + assert diagnostics.angular_resolution == 256 + assert diagnostics.newton_iterations == 8 + + radius = diagnostics.r_singularity_vs_varphi + theta = diagnostics.theta_singularity_vs_varphi + linear = diagnostics.g1c * jnp.cos(theta) + diagnostics.g1s * jnp.sin(theta) + linear_prime = -diagnostics.g1c * jnp.sin(theta) + diagnostics.g1s * jnp.cos(theta) + quadratic = ( + diagnostics.g20 + + diagnostics.g2s * jnp.sin(2 * theta) + + diagnostics.g2c * jnp.cos(2 * theta) + ) + quadratic_prime = 2 * diagnostics.g2s * jnp.cos(2 * theta) - 2 * diagnostics.g2c * jnp.sin( + 2 * theta + ) + np.testing.assert_allclose( + diagnostics.g0 + radius * linear + radius**2 * quadratic, + 0, + rtol=0, + atol=2e-13, + ) + np.testing.assert_allclose( + radius * linear_prime + radius**2 * quadratic_prime, + 0, + rtol=0, + atol=2e-13, + ) + + +def test_newton_refinement_removes_angular_grid_dependence(): + solution = qsc.Qsc(**CASES[0][0], nphi=31, order="r2") + coarse = qsc.singularity_diagnostics( + solution, + angular_resolution=32, + newton_iterations=8, + ) + fine = qsc.singularity_diagnostics( + solution, + angular_resolution=512, + newton_iterations=8, + ) + + np.testing.assert_allclose(coarse.r_singularity, fine.r_singularity, rtol=0, atol=2e-11) + assert coarse.maximum_residual_norm < 3e-13 + + +def test_singular_radius_supports_jit_and_jvp(): + def radius(etabar): + return qsc.Qsc( + rc=[1.0, 0.155, 0.0102], + zs=[0.0, 0.154, 0.0111], + nfp=2, + etabar=etabar, + B2c=-0.00322, + nphi=31, + order="r2", + ).r_singularity + + eager = radius(0.64) + compiled = jax.jit(radius)(0.64) + _, tangent = jax.jvp(radius, (0.64,), (1.0,)) + step = 1e-5 + finite_difference = (radius(0.64 + step) - radius(0.64 - step)) / (2 * step) + + np.testing.assert_allclose(compiled, eager, rtol=2e-11, atol=2e-11) + np.testing.assert_allclose(tangent, finite_difference, rtol=3e-7, atol=3e-8) + + +def test_singularity_validation_and_legacy_adapter(): + first_order = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, + ) + with pytest.raises(ValueError, match="second-order"): + qsc.singularity_diagnostics(first_order) + with pytest.raises(ValueError, match="angular_resolution"): + qsc.singularity_diagnostics( + qsc.Qsc(**CASES[0][0], nphi=15, order="r2"), + angular_resolution=4, + ) + with pytest.raises(ValueError, match="newton_iterations"): + qsc.singularity_diagnostics( + qsc.Qsc(**CASES[0][0], nphi=15, order="r2"), + newton_iterations=-1, + ) + with pytest.raises(AttributeError, match="singular-radius"): + _ = first_order.r_singularity + + legacy = near_axis(**CASES[0][0], nphi=31, order="r2") + np.testing.assert_allclose(legacy.r_singularity, legacy.solution.r_singularity) + np.testing.assert_allclose( + legacy.r_singularity_residual_sqnorm, + legacy.solution.r_singularity_residual_sqnorm, + ) diff --git a/tests/physics/test_third_order.py b/tests/physics/test_third_order.py new file mode 100644 index 0000000..5e0294a --- /dev/null +++ b/tests/physics/test_third_order.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.near_axis import near_axis + +# Frozen BSD-2-Clause pyQSC reference samples at the audited upstream commit. +CASES = ( + ( + "qa", + {}, + { + "X3c1": [0.06254885838364596, 1.1028068624859182, 0.9898672691339223], + "Y3c1": [0.0, 1.9611076246147592, 0.06254912908105097], + "Y3s1": [0.263992436104027, 3.1822277655813185, 0.6000377597469493], + "B0_order_a_squared_to_cancel": [ + 0.257001365727628, + 3.7466692503729218, + 1.5413730743953158, + ], + }, + ), + ( + "finite_pressure_current", + { + "rc": [1.0, 0.09], + "zs": [0.0, -0.09], + "nfp": 2, + "etabar": 0.95, + "I2": 0.9, + "B2c": -0.7, + "p2": -600000.0, + }, + { + "X3c1": [-0.7763858841222269, -0.2167234757709553, -0.4037215275576512], + "Y3c1": [0.0, 0.13599120989448577, 0.02104082225854222], + "Y3s1": [-1.2142017571662096, -0.27932779079566616, -0.18341957155186067], + "B0_order_a_squared_to_cancel": [ + -1.9418435614778733, + -0.4920849106382204, + -0.5442441717107536, + ], + }, + ), + ( + "qh", + { + "rc": [1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + "zs": [0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + "nfp": 4, + "etabar": 1.569, + "B2c": 0.1348, + }, + { + "X3c1": [-0.05934821467751999, -0.002180735261069853, -0.006774272974061107], + "Y3c1": [0.0, 0.002774228833885528, 0.00030670030088582816], + "Y3s1": [-0.15982340894062852, -0.002817420091098044, -0.002404111069876869], + "B0_order_a_squared_to_cancel": [ + -0.1947843318576025, + -0.0049574377809037395, + -0.008071209239579269, + ], + }, + ), +) + + +def standard_solution(nphi=31, **kwargs): + parameters = { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + "nphi": nphi, + } + parameters.update(kwargs) + return qsc.Qsc(**parameters) + + +@pytest.mark.parametrize(("_name", "kwargs", "reference"), CASES) +def test_third_order_matches_upstream_pyqsc(_name, kwargs, reference): + solution = standard_solution(order="r3", nphi=61, **kwargs) + indices = np.asarray([0, 15, 30]) + + for name, expected in reference.items(): + np.testing.assert_allclose( + np.asarray(getattr(solution, name))[indices], + expected, + rtol=3.0e-9, + atol=5.0e-10, + ) + + +def test_third_order_structure_and_independent_constraints(): + solution = standard_solution( + rc=[1.0, 0.09], + zs=[0.0, -0.09], + nfp=2, + etabar=0.95, + I2=0.9, + B2c=-0.7, + p2=-600000.0, + order="r3", + nphi=61, + ) + third = solution.third_order + assert third is not None + + np.testing.assert_allclose(third.X3c1, solution.X1c * third.flux_constraint_coefficient) + np.testing.assert_allclose(third.Y3c1, solution.Y1c * third.flux_constraint_coefficient) + np.testing.assert_allclose(third.Y3s1, solution.Y1s * third.flux_constraint_coefficient) + np.testing.assert_allclose( + third.B0_order_a_squared_to_cancel, + 2 * solution.inputs.B0 * third.flux_constraint_coefficient, + rtol=2.0e-10, + atol=2.0e-12, + ) + assert float(jnp.max(jnp.abs(third.flux_constraint_residual))) < 2.0e-13 + assert float(jnp.max(jnp.abs(third.consistency_error))) < 2.0e-10 + + for name in ( + "X3s1", + "Z3s1", + "Z3c1", + "X3s3", + "X3c3", + "Y3s3", + "Y3c3", + "Z3s3", + "Z3c3", + ): + np.testing.assert_array_equal(getattr(third, name), jnp.zeros(solution.inputs.nphi)) + + derivative = solution.geometry.d_d_varphi + np.testing.assert_allclose(third.d_X3c1_d_varphi, derivative @ third.X3c1, atol=2.0e-12) + np.testing.assert_allclose(third.d_Y3s1_d_varphi, derivative @ third.Y3s1, atol=2.0e-12) + np.testing.assert_allclose(third.d_Y3c1_d_varphi, derivative @ third.Y3c1, atol=2.0e-12) + + +def test_qh_untwisting_and_compatibility_surface_include_r3(): + r3 = qsc.Qsc( + rc=[1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + zs=[0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + nfp=4, + etabar=1.569, + B2c=0.1348, + order="r3", + nphi=61, + ) + + assert np.linalg.norm(np.asarray(r3.X3s1_untwisted)) > 0 + adapter = near_axis( + rc=[1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + zs=[0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + nfp=4, + etabar=1.569, + B2c=0.1348, + nphi=61, + order="r3", + ) + r = 0.02 + theta = 0.37 + x2, y2, z2 = near_axis( + rc=[1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + zs=[0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + nfp=4, + etabar=1.569, + B2c=0.1348, + nphi=61, + order="r2", + )._frenet_displacements(r, theta) + x3, y3, z3 = adapter._frenet_displacements(r, theta) + cosine = jnp.cos(theta) + sine = jnp.sin(theta) + cosine3 = jnp.cos(3 * theta) + sine3 = jnp.sin(3 * theta) + + np.testing.assert_allclose( + x3 - x2, + r**3 + * ( + r3.X3c1_untwisted * cosine + + r3.X3s1_untwisted * sine + + r3.X3c3_untwisted * cosine3 + + r3.X3s3_untwisted * sine3 + ), + atol=2.0e-15, + ) + np.testing.assert_allclose( + y3 - y2, + r**3 + * ( + r3.Y3c1_untwisted * cosine + + r3.Y3s1_untwisted * sine + + r3.Y3c3_untwisted * cosine3 + + r3.Y3s3_untwisted * sine3 + ), + atol=2.0e-15, + ) + np.testing.assert_allclose(z3, z2, atol=2.0e-15) + + +def test_third_order_is_jittable_and_differentiable(): + def coefficient(etabar): + return standard_solution(order="r3", nphi=31, etabar=etabar).X3c1[7] + + value = coefficient(jnp.asarray(-0.9)) + jitted = jax.jit(coefficient)(jnp.asarray(-0.9)) + tangent = jax.jvp(coefficient, (jnp.asarray(-0.9),), (jnp.asarray(1.0),))[1] + step = 2.0e-5 + finite_difference = (coefficient(-0.9 + step) - coefficient(-0.9 - step)) / (2 * step) + + np.testing.assert_allclose(jitted, value, rtol=2.0e-12, atol=2.0e-12) + np.testing.assert_allclose(tangent, finite_difference, rtol=3.0e-5, atol=3.0e-7) + + +def test_third_order_requires_second_order_and_valid_adapter_order(): + r1 = standard_solution(order="r1") + with pytest.raises(ValueError, match="second-order"): + qsc.solve_third_order(r1) + with pytest.raises(AttributeError, match="Lower-order"): + _ = r1.X3c1 + with pytest.raises(ValueError, match="order must be"): + near_axis(rc=[1.0], zs=[0.0], etabar=-0.9, order="r4") diff --git a/tests/reference/vmec/input.qa_r0025 b/tests/reference/vmec/input.qa_r0025 new file mode 100644 index 0000000..c5fe202 --- /dev/null +++ b/tests/reference/vmec/input.qa_r0025 @@ -0,0 +1,115 @@ +! Generated deterministically by pyQSC_JAX. +! Near-axis radius r = 2.5000000000000001e-03; etabar = -9.0000000000000002e-01. +! nphi = 121; order = r2; ntheta = 32; mpol = 6; ntor = 6. +! Conversion diagnostics: max_phi_residual = 1.387779e-17; max_R_error = 1.460998e-06; max_Z_error = 1.080278e-06. +&INDATA + DELT = 9.0000000000000002e-01 + NSTEP = 200 + TCON0 = 2.0000000000000000e+00 + NS_ARRAY = 31 + FTOL_ARRAY = 1.0000000000000000e-10 + NITER_ARRAY = 3000 + LASYM = F + LFREEB = F + NFP = 3 + MPOL = 6 + NTOR = 6 + PHIEDGE = 1.9634954084936207e-05 + PRES_SCALE = 1.0000000000000000e+00 + PMASS_TYPE = 'power_series' + AM = -0.0000000000000000e+00, 0.0000000000000000e+00 + CURTOR = 0.0000000000000000e+00 + NCURR = 1 + PCURR_TYPE = 'power_series' + AC = 1.0000000000000000e+00 + RAXIS_CC = 1.0000000000000000e+00, 4.4999999999999998e-02 + RAXIS_CS = -0.0000000000000000e+00, -0.0000000000000000e+00 + ZAXIS_CC = 0.0000000000000000e+00, 0.0000000000000000e+00 + ZAXIS_CS = -0.0000000000000000e+00, 4.4999999999999998e-02 +! Boundary coefficients + RBC(000,000) = +1.0000705829799887e+00, ZBS(000,000) = +0.0000000000000000e+00 + RBC(001,000) = +4.4972146504766455e-02, ZBS(001,000) = +4.4991499300713611e-02 + RBC(002,000) = +5.9873785623381428e-06, ZBS(002,000) = +4.3502961734281688e-07 + RBC(003,000) = -2.1517234573157090e-06, ZBS(003,000) = +4.7421860162078269e-07 + RBC(004,000) = +1.0991923800079042e-06, ZBS(004,000) = -4.8106974614858439e-07 + RBC(005,000) = -6.4130044677373424e-07, ZBS(005,000) = +3.7686082803356997e-07 + RBC(006,000) = +3.8870569264632445e-07, ZBS(006,000) = -2.6452556339482058e-07 + RBC(-06,001) = -2.4415059227033655e-07, ZBS(-06,001) = -1.6757477768105833e-07 + RBC(-05,001) = +6.4602884351962735e-07, ZBS(-05,001) = +3.9729116716134149e-07 + RBC(-04,001) = -1.8121635281841622e-06, ZBS(-04,001) = -9.1949094106951558e-07 + RBC(-03,001) = +5.5483444835936612e-06, ZBS(-03,001) = +1.9926099557692592e-06 + RBC(-02,001) = -1.6936463407602224e-05, ZBS(-02,001) = -5.9095719187370043e-06 + RBC(-01,001) = +3.3768733026137273e-05, ZBS(-01,001) = +4.9271499850138361e-05 + RBC(000,001) = +2.7275400006613565e-03, ZBS(000,001) = -2.7303052205015169e-03 + RBC(001,001) = -1.0843959315583243e-03, ZBS(001,001) = -1.0456522227157282e-03 + RBC(002,001) = +5.7683416839753723e-05, ZBS(002,001) = +7.7805475764700576e-05 + RBC(003,001) = +2.2469329177139640e-06, ZBS(003,001) = -5.2325258405908656e-06 + RBC(004,001) = -1.6896070240558097e-06, ZBS(004,001) = +1.0411745545182424e-06 + RBC(005,001) = +6.4216920053210053e-07, ZBS(005,001) = -4.0149036024405450e-07 + RBC(006,001) = -2.4405533518398577e-07, ZBS(006,001) = +1.6770724800463921e-07 + RBC(-06,002) = +3.1007833251658884e-07, ZBS(-06,002) = +2.4574641759511880e-07 + RBC(-05,002) = -5.2948451720773412e-07, ZBS(-05,002) = -3.9717792669642647e-07 + RBC(-04,002) = +9.1529307608356163e-07, ZBS(-04,002) = +6.1694907139047096e-07 + RBC(-03,002) = -1.7275408599925077e-06, ZBS(-03,002) = -8.7176040692811953e-07 + RBC(-02,002) = +4.3548900479616564e-06, ZBS(-02,002) = +6.2924084662730571e-07 + RBC(-01,002) = -1.5628821641526291e-05, ZBS(-01,002) = +1.8711493375401293e-06 + RBC(000,002) = -7.7093418228937785e-06, ZBS(000,002) = +8.0247944040313932e-05 + RBC(001,002) = +3.0483893018425361e-05, ZBS(001,002) = +3.5006771454329695e-05 + RBC(002,002) = -3.8692199524025512e-06, ZBS(002,002) = -6.5631891418695323e-06 + RBC(003,002) = -2.6021247701713909e-07, ZBS(003,002) = +8.2982903320988042e-07 + RBC(004,002) = +4.9602358763677354e-07, ZBS(004,002) = -3.5173158952779694e-08 + RBC(005,002) = -3.0005785856526528e-07, ZBS(005,002) = +1.1487820723356077e-08 + RBC(006,002) = +1.7354345575586541e-07, ZBS(006,002) = -2.4597539487659144e-08 + RBC(-06,003) = -3.3605899502702502e-10, ZBS(-06,003) = -4.9195863015844885e-10 + RBC(-05,003) = +5.1654029748188752e-10, ZBS(-05,003) = +7.4714165541172801e-10 + RBC(-04,003) = -7.5530762265820549e-10, ZBS(-04,003) = -1.1313424261651086e-09 + RBC(-03,003) = +9.8440354614450489e-10, ZBS(-03,003) = +1.7606542541457767e-09 + RBC(-02,003) = -3.4080094470854474e-10, ZBS(-02,003) = -3.5436337168678692e-09 + RBC(-01,003) = -2.1375749176700437e-09, ZBS(-01,003) = +8.9545429090448358e-09 + RBC(000,003) = -1.9612354192356898e-08, ZBS(000,003) = -1.3132081299764493e-09 + RBC(001,003) = -2.8655353251880292e-08, ZBS(001,003) = -1.0271426692348341e-08 + RBC(002,003) = -2.3747997161433985e-08, ZBS(002,003) = -4.4964079315489770e-08 + RBC(003,003) = -6.1455223117203370e-09, ZBS(003,003) = -2.0309373874284110e-09 + RBC(004,003) = +1.3717169981562259e-08, ZBS(004,003) = +9.6826085398181361e-09 + RBC(005,003) = -3.4488416158216233e-09, ZBS(005,003) = -9.0133423538389967e-10 + RBC(006,003) = +2.1503114674813830e-09, ZBS(006,003) = +4.9954874279097632e-10 + RBC(-06,004) = +4.6369779620584815e-11, ZBS(-06,004) = +5.8940424957078450e-11 + RBC(-05,004) = -2.4106668597179557e-11, ZBS(-05,004) = -6.5418755572823082e-11 + RBC(-04,004) = +4.0237650219534292e-11, ZBS(-04,004) = +8.6792187199893126e-11 + RBC(-03,004) = -1.8070051228621789e-11, ZBS(-03,004) = -1.0961627411167680e-10 + RBC(-02,004) = +7.8839647338186660e-11, ZBS(-02,004) = +5.5603632124865108e-11 + RBC(-01,004) = -3.9934554570480501e-10, ZBS(-01,004) = +3.6180781813635285e-10 + RBC(000,004) = -6.3301072225918542e-10, ZBS(000,004) = +1.2190921610862814e-09 + RBC(001,004) = +1.4350897724882337e-08, ZBS(001,004) = -1.4619286439620352e-08 + RBC(002,004) = +3.9707409493019856e-08, ZBS(002,004) = -4.0896719720885368e-08 + RBC(003,004) = +3.9202544088903947e-08, ZBS(003,004) = -9.8722393081329035e-09 + RBC(004,004) = +2.0277841430761762e-08, ZBS(004,004) = +5.9601236482370415e-08 + RBC(005,004) = +3.5861092221904178e-08, ZBS(005,004) = +1.3364988278544130e-08 + RBC(006,004) = -2.4269200734071315e-08, ZBS(006,004) = -1.4816477416940341e-08 + RBC(-06,005) = -2.6226002745625437e-12, ZBS(-06,005) = -7.2846539236555959e-12 + RBC(-05,005) = -1.3028100731259414e-11, ZBS(-05,005) = +7.1941878267675343e-12 + RBC(-04,005) = +2.8468464782067269e-11, ZBS(-04,005) = -2.7234203358021016e-11 + RBC(-03,005) = -5.3235556285850331e-11, ZBS(-03,005) = +3.6448043799500747e-11 + RBC(-02,005) = +1.2397062872034374e-10, ZBS(-02,005) = -7.0457689479103539e-11 + RBC(-01,005) = -7.1319150981371117e-11, ZBS(-01,005) = -1.9280850945643681e-11 + RBC(000,005) = -4.3987912406425169e-10, ZBS(000,005) = +6.3877670878631776e-10 + RBC(001,005) = -3.1939019251615725e-09, ZBS(001,005) = +3.1261830312733297e-09 + RBC(002,005) = -3.2643590893336510e-09, ZBS(002,005) = +2.2739485338332859e-09 + RBC(003,005) = -4.2118428941129741e-09, ZBS(003,005) = -2.4311639430023996e-09 + RBC(004,005) = +1.4610748856713604e-09, ZBS(004,005) = -8.2230180190818004e-09 + RBC(005,005) = -2.4678340978275323e-09, ZBS(005,005) = -2.7907292122839379e-11 + RBC(006,005) = +6.2121165377238212e-09, ZBS(006,005) = +5.2623472883796908e-09 + RBC(-06,006) = +3.8736295304419662e-11, ZBS(-06,006) = +5.7060497708605618e-12 + RBC(-05,006) = -4.7470189245640317e-12, ZBS(-05,006) = -7.9399004719218432e-12 + RBC(-04,006) = +1.8443338421304800e-11, ZBS(-04,006) = +1.8046986663601585e-11 + RBC(-03,006) = +1.1167396860698812e-11, ZBS(-03,006) = -1.2698250082275888e-11 + RBC(-02,006) = -6.9312333793738727e-12, ZBS(-02,006) = +1.3144183600258469e-11 + RBC(-01,006) = +6.7578184813644228e-11, ZBS(-01,006) = -5.1060449472660376e-11 + RBC(000,006) = -1.1230898093480489e-10, ZBS(000,006) = +1.0101530760459938e-10 + RBC(001,006) = -3.0274151016801464e-10, ZBS(001,006) = +4.0970273730041438e-10 + RBC(002,006) = +2.0430123815145075e-09, ZBS(002,006) = -2.2459041353843439e-09 + RBC(003,006) = +1.2097337288909390e-08, ZBS(003,006) = -1.2714596943081256e-08 + RBC(004,006) = +1.8328808588258400e-08, ZBS(004,006) = -1.3955119471723512e-08 + RBC(005,006) = +1.3074406926366760e-08, ZBS(005,006) = +1.1765439645369073e-08 + RBC(006,006) = +9.4881890380617359e-09, ZBS(006,006) = +2.2478558100516398e-08 +/ diff --git a/tests/reference/vmec/manifest.json b/tests/reference/vmec/manifest.json new file mode 100644 index 0000000..993e08c --- /dev/null +++ b/tests/reference/vmec/manifest.json @@ -0,0 +1,40 @@ +{ + "case": "qa_r0025", + "description": "Small-radius fixed-boundary VMEC validation of the pyQSC_JAX exporter.", + "generator": { + "executable": "/Users/rogerio/base_env/bin/xvmec", + "vmec_version": "9.0", + "platform": "macOS arm64" + }, + "near_axis": { + "B0": 1.0, + "etabar": -0.9, + "iota": 0.41830691021517735, + "nfp": 3, + "nphi": 121, + "order": "r2", + "radius": 0.0025, + "rc": [1.0, 0.045], + "zs": [0.0, -0.045] + }, + "export": { + "maximum_R_reconstruction_error": 1.460998e-06, + "maximum_Z_reconstruction_error": 1.080278e-06, + "maximum_toroidal_angle_residual": 1.387779e-17, + "mpol": 6, + "ntheta": 32, + "ntor": 6 + }, + "files": { + "input.qa_r0025": "5444b159de7562681363f65426e814fb46c11437cf3d1b31e5c20ba7c1fbedc1", + "wout_qa_r0025.nc": "0bbbb1e7681462aec1c6b726edc8d20dc1c18ce00116174159ae330d7ee73835" + }, + "vmec_result": { + "aspect": 398.2269295920157, + "fsql": 2.402075927993851e-11, + "fsqr": 7.589358862704765e-11, + "fsqz": 4.213307210656955e-11, + "iota_axis": 0.418543069691998, + "relative_iota_error": 0.0005645603050142573 + } +} diff --git a/tests/reference/vmec/wout_qa_r0025.nc b/tests/reference/vmec/wout_qa_r0025.nc new file mode 100644 index 0000000..9dc091e Binary files /dev/null and b/tests/reference/vmec/wout_qa_r0025.nc differ diff --git a/tests/regression/test_first_order_reference.py b/tests/regression/test_first_order_reference.py new file mode 100644 index 0000000..8243eb0 --- /dev/null +++ b/tests/regression/test_first_order_reference.py @@ -0,0 +1,77 @@ +import numpy as np + +import pyqsc_jax as qsc +from pyqsc_jax.near_axis import near_axis + + +def test_standard_solution_matches_legacy_and_upstream_reference(): + parameters = { + "rc": [1.0, 0.045], + "zs": [0.0, -0.045], + "nfp": 3, + "etabar": -0.9, + "nphi": 31, + } + solution = qsc.Qsc(**parameters) + legacy = near_axis(**parameters) + + np.testing.assert_allclose(solution.iota, 0.41830690943386617, rtol=2e-13) + np.testing.assert_allclose(solution.iota, legacy.iota, rtol=2e-13) + np.testing.assert_allclose(solution.sigma, legacy.sigma, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(solution.B_axis, legacy.B_axis.T, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose( + solution.grad_B_axis, + np.moveaxis(legacy.grad_B_axis, -1, 0), + rtol=2e-12, + atol=2e-12, + ) + np.testing.assert_allclose(solution.L_grad_B, legacy.L_grad_B, rtol=2e-13) + + +def test_finite_current_and_sigma0_match_upstream_reference(): + solution = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, + I2=0.3, + B0=1.2, + sigma0=0.1, + ) + + np.testing.assert_allclose(solution.iota, 0.5981040523602934, rtol=2e-13) + np.testing.assert_allclose(solution.iotaN, 0.5981040523602934, rtol=2e-13) + np.testing.assert_allclose(solution.G0, 1.2108964178133113, rtol=2e-13) + np.testing.assert_allclose(np.max(np.abs(solution.sigma)), 1.1527807285268685, rtol=2e-13) + np.testing.assert_allclose(solution.mean_elongation, 2.300022135471462, rtol=2e-13) + indices = [0, 7, 15] + expected_gradient_components = { + (0, 1): [-1.4189028123671708, -0.024458076547534513, 0.31866877632846324], + (1, 0): [-1.4958375512360236, -0.03652380306707692, 0.40219561346636035], + (0, 2): [1.248855714330292, 0.24413995843927835, -0.7230617247040809], + (2, 0): [0.6533238467899136, -0.2490021462405346, -1.0300929188685888], + } + for component, expected in expected_gradient_components.items(): + np.testing.assert_allclose( + np.asarray(solution.grad_B_axis)[indices, *component], + expected, + rtol=3e-12, + atol=3e-12, + ) + + +def test_qh_topology_matches_upstream_reference(): + solution = qsc.Qsc( + rc=[1.0, 0.265], + zs=[0.0, -0.21], + nfp=4, + etabar=-0.9, + nphi=31, + ) + + assert int(solution.helicity) == -1 + np.testing.assert_allclose(solution.iota, 3.0598175213150203, rtol=2e-13) + np.testing.assert_allclose(solution.iotaN, -0.9401824786849797, rtol=2e-13) + np.testing.assert_allclose(solution.G0, 1.3885479374943943, rtol=2e-13) + np.testing.assert_allclose(np.max(np.abs(solution.sigma)), 0.23608475507759125, rtol=2e-13) diff --git a/tests/regression/test_geometry_reference.py b/tests/regression/test_geometry_reference.py new file mode 100644 index 0000000..083b557 --- /dev/null +++ b/tests/regression/test_geometry_reference.py @@ -0,0 +1,104 @@ +import numpy as np + +from pyqsc_jax import Axis +from pyqsc_jax.geometry import compute_axis_geometry +from pyqsc_jax.near_axis import near_axis + + +def test_geometry_matches_legacy_first_order_baseline(): + parameters = { + "rc": [1.0, 0.045], + "zs": [0.0, -0.045], + "nfp": 3, + "nphi": 31, + } + geometry = compute_axis_geometry( + Axis(rc=parameters["rc"], zs=parameters["zs"], nfp=parameters["nfp"]), + nphi=parameters["nphi"], + ) + legacy = near_axis(etabar=-0.9, **parameters) + + np.testing.assert_allclose(geometry.samples.R, legacy.R0, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(geometry.samples.Z, legacy.Z0, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(geometry.axis_length, legacy.axis_length, rtol=2e-13) + np.testing.assert_allclose(geometry.curvature, legacy.curvature, rtol=2e-13) + np.testing.assert_allclose(geometry.torsion, legacy.torsion, rtol=2e-13) + np.testing.assert_allclose(geometry.varphi, legacy.varphi, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(geometry.normal_cylindrical[:, 0], legacy.normal_R, rtol=2e-13) + np.testing.assert_allclose(geometry.normal_cylindrical[:, 1], legacy.normal_phi, rtol=2e-13) + np.testing.assert_allclose(geometry.normal_cylindrical[:, 2], legacy.normal_z, rtol=2e-13) + + +def test_asymmetric_geometry_matches_independent_fortran_reference(): + axis = Axis( + rc=[1.3, 0.3, 0.01, -0.001], + zs=[0.0, 0.4, -0.02, -0.003], + rs=[0.0, -0.1, -0.03, 0.002], + zc=[0.3, 0.2, 0.04, 0.004], + nfp=5, + ) + geometry = compute_axis_geometry(axis, nphi=15) + curvature = [ + 2.10743037699653, + 2.33190181686696, + 1.83273654023051, + 1.81062232906827, + 2.28640008392347, + 1.76919841474321, + 0.919988560478029, + 0.741327470169023, + 1.37147330126897, + 2.64680884158075, + 3.39786486424852, + 2.47005615416209, + 1.50865425515356, + 1.18136509189105, + 1.42042418970102, + ] + torsion = [ + -0.167822738386845, + -0.0785778346620885, + -1.02205137493593, + -2.05213528002946, + -0.964613202459108, + -0.593496282035916, + -2.15852857178204, + -3.72911055219339, + -1.9330792779459, + -1.53882290974916, + -1.42156496444929, + -1.11381642382793, + -0.92608309386204, + -0.868339812017432, + -0.57696266498748, + ] + varphi = [ + 0.0, + 0.084185130335249, + 0.160931495903817, + 0.232881563535092, + 0.300551168190665, + 0.368933497012765, + 0.444686439112853, + 0.528001290336008, + 0.612254611059372, + 0.691096975269652, + 0.765820243301147, + 0.846373713025902, + 0.941973362938683, + 1.05053459351092, + 1.15941650366667, + ] + + np.testing.assert_allclose(geometry.curvature, curvature, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(geometry.torsion, torsion, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(geometry.varphi, varphi, rtol=2e-13, atol=2e-13) + + +def test_qh_frame_helicity_matches_upstream_convention(): + geometry = compute_axis_geometry( + Axis(rc=[1.0, 0.265], zs=[0.0, -0.21], nfp=4), + nphi=31, + ) + + assert int(geometry.frame_helicity) == -1 diff --git a/tests/regression/test_legacy_baseline.py b/tests/regression/test_legacy_baseline.py new file mode 100644 index 0000000..7198e2f --- /dev/null +++ b/tests/regression/test_legacy_baseline.py @@ -0,0 +1,71 @@ +import jax.numpy as jnp +import numpy as np + +from pyqsc_jax.near_axis import near_axis + + +def standard_field(**kwargs): + parameters = { + "rc": [1.0, 0.045], + "zs": [0.0, -0.045], + "etabar": -0.9, + "nfp": 3, + "nphi": 31, + } + parameters.update(kwargs) + return near_axis(**parameters) + + +def test_standard_first_order_baseline(): + field = standard_field() + + np.testing.assert_allclose(field.iota, 0.41830690943386617, rtol=2e-13) + np.testing.assert_allclose(field.axis_length, 6.340238817434161, rtol=2e-13) + np.testing.assert_allclose( + [jnp.min(field.curvature), jnp.max(field.curvature)], + [0.5956163516982903, 1.3060121594235536], + rtol=2e-13, + ) + np.testing.assert_allclose( + [jnp.min(field.torsion), jnp.max(field.torsion)], + [-2.272197039914583, 0.6057564303422814], + rtol=2e-13, + ) + np.testing.assert_allclose(jnp.max(jnp.abs(field.sigma)), 0.9920706836908618, rtol=2e-13) + assert field.B_axis.shape == (3, 31) + assert field.grad_B_axis.shape == (3, 3, 31) + + +def test_first_order_field_methods_are_finite(): + field = standard_field() + point = jnp.array([0.02, 0.4, 0.1]) + + assert jnp.isfinite(field.AbsB(point)) + assert jnp.isfinite(field.jacobian(point)) + assert jnp.all(jnp.isfinite(field.B_covariant(point))) + assert jnp.all(jnp.isfinite(field.B_contravariant(point))) + + +def test_dofs_noop_preserves_derived_frame(): + field = standard_field() + normal_before = jnp.stack([field.normal_R, field.normal_phi, field.normal_z], axis=1) + binormal_before = jnp.stack([field.binormal_R, field.binormal_phi, field.binormal_z], axis=1) + + field.dofs = field.dofs + + normal_after = jnp.stack([field.normal_R, field.normal_phi, field.normal_z], axis=1) + binormal_after = jnp.stack([field.binormal_R, field.binormal_phi, field.binormal_z], axis=1) + np.testing.assert_array_equal(normal_after, normal_before) + np.testing.assert_array_equal(binormal_after, binormal_before) + + +def test_r2_request_produces_second_order_solution(): + field = standard_field(order="r2", B2c=0.01, p2=-1.0e3) + + assert field.B20.shape == (field.nphi,) + assert jnp.all(jnp.isfinite(field.B20)) + assert bool(field.linear_report.converged) + R, Z, phi0 = field.Frenet_to_cylindrical(0.005, ntheta=3) + assert R.shape == Z.shape == phi0.shape == (3, field.nphi) + assert jnp.all(jnp.isfinite(R)) + assert jnp.isfinite(field.B_mag(0.005, 0.2, 0.1)) diff --git a/tests/regression/test_second_order_reference.py b/tests/regression/test_second_order_reference.py new file mode 100644 index 0000000..2900f5c --- /dev/null +++ b/tests/regression/test_second_order_reference.py @@ -0,0 +1,88 @@ +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +# References were generated with landreman/pyQSC at audited commit +# cd75359ea47548d5db7ccb458c100085c04ba1bc. +@pytest.mark.parametrize( + "parameters, expected", + [ + ( + { + "rc": [1.0, 0.155, 0.0102], + "zs": [0.0, 0.154, 0.0111], + "nfp": 2, + "etabar": 0.64, + "B2c": -0.00322, + }, + { + "iota": -0.42047335182825857, + "beta_1s": 0.0, + "G2": 0.0, + "B20_mean": 0.16212431504754432, + "B20_residual": 0.13498174214090364, + "B20_variation": 0.3740338506893984, + "X20": [-0.12607391440629268, 1.277313706167806, -1.1433646049834585], + "Y20": [3.6484581072275835e-14, -0.061041743645613045, -0.28837977142377097], + "B20": [0.03617846642088235, 0.08753146934212608, 0.3504652075357061], + }, + ), + ( + { + "rc": [1.0, 0.09], + "zs": [0.0, -0.09], + "nfp": 2, + "etabar": 0.95, + "I2": 0.9, + "B2c": -0.7, + "p2": -600000.0, + }, + { + "iota": 0.9596981597369478, + "beta_1s": 3.0336182742616837, + "G2": -0.09758153451238105, + "B20_mean": 1.8129933830726364, + "B20_residual": 0.22096470852551833, + "B20_variation": 0.6235580990378657, + "X20": [1.0359701826468048, 0.9961323682664798, 1.8900857033253329], + "Y20": [-2.4299426345854533e-15, -0.9136862473826926, -0.30321915705664426], + "B20": [2.070309301450683, 1.8875840576321603, 1.4467512024128548], + }, + ), + ( + { + "rc": [1.0, 0.17, 0.01804, 0.001409, 5.877e-5], + "zs": [0.0, 0.1581, 0.01820, 0.001548, 7.772e-5], + "nfp": 4, + "etabar": 1.569, + "B2c": 0.1348, + }, + { + "iota": -1.1441369511851485, + "beta_1s": 0.0, + "G2": 0.0, + "B20_mean": 1.312575397053029, + "B20_residual": 0.032260200125823604, + "B20_variation": 0.11724965804270937, + "X20": [-0.18543205322365838, 0.7244473635400297, 0.8388295002982797], + "Y20": [-1.0252483956584066e-14, -0.07290690883420155, -0.012145324323969307], + "B20": [1.2242278807890326, 1.3246881250407083, 1.3414775388313407], + }, + ), + ], +) +def test_second_order_matches_upstream(parameters, expected): + solution = qsc.Qsc(nphi=31, order="r2", **parameters) + + for name in ("iota", "beta_1s", "G2", "B20_mean", "B20_residual", "B20_variation"): + np.testing.assert_allclose(getattr(solution, name), expected[name], rtol=3e-12, atol=3e-12) + indices = [0, 7, 15] + for name in ("X20", "Y20", "B20"): + np.testing.assert_allclose( + np.asarray(getattr(solution, name))[indices], + expected[name], + rtol=3e-11, + atol=3e-11, + ) diff --git a/tests/unit/test_axis.py b/tests/unit/test_axis.py new file mode 100644 index 0000000..cffb709 --- /dev/null +++ b/tests/unit/test_axis.py @@ -0,0 +1,78 @@ +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from pyqsc_jax import Axis +from pyqsc_jax.axis import evaluate_axis + + +def test_coefficients_are_normalized_and_round_trip(): + axis = Axis( + rc=[1.0, 0.2, -0.01], + rs=[0.0, 0.03], + zc=[0.1], + zs=[0.0, -0.15], + nfp=5, + ) + + assert axis.nfourier == 3 + np.testing.assert_array_equal(axis.rs, [0.0, 0.03, 0.0]) + np.testing.assert_array_equal(axis.zc, [0.1, 0.0, 0.0]) + np.testing.assert_array_equal(axis.zs, [0.0, -0.15, 0.0]) + np.testing.assert_allclose(axis.stellarator_symmetry_residual, 0.1) + rebuilt = Axis.from_dofs(axis.dofs, nfp=axis.nfp) + assert rebuilt.nfp == axis.nfp + np.testing.assert_array_equal(rebuilt.dofs, axis.dofs) + shifted = axis.with_dofs(axis.dofs.at[0].add(0.2)) + np.testing.assert_allclose(shifted.rc[0], 1.2) + + +def test_axis_input_validation(): + with pytest.raises(ValueError, match="positive integer"): + Axis(rc=[1.0], zs=[0.0], nfp=0) + with pytest.raises(ValueError, match="positive integer"): + Axis(rc=[1.0], zs=[0.0], nfp=True) + with pytest.raises(ValueError, match="one-dimensional"): + Axis(rc=[[1.0]], zs=[0.0]) + with pytest.raises(ValueError, match="At least one"): + Axis(rc=[], zs=[]) + with pytest.raises(ValueError, match="divisible by 4"): + Axis.from_dofs(jnp.arange(5.0), nfp=1) + + +def test_general_axis_and_analytic_derivatives(): + axis = Axis( + rc=[1.2, 0.3], + rs=[0.0, -0.11], + zc=[0.2, 0.04], + zs=[0.0, 0.17], + nfp=3, + ) + phi = jnp.array([0.0, 0.13, 0.51]) + samples = evaluate_axis(axis, phi) + angle = 3 * phi + + np.testing.assert_allclose(samples.R, 1.2 + 0.3 * jnp.cos(angle) - 0.11 * jnp.sin(angle)) + np.testing.assert_allclose(samples.Z, 0.2 + 0.04 * jnp.cos(angle) + 0.17 * jnp.sin(angle)) + np.testing.assert_allclose(samples.d_R_d_phi, -0.9 * jnp.sin(angle) - 0.33 * jnp.cos(angle)) + np.testing.assert_allclose(samples.d_Z_d_phi, -0.12 * jnp.sin(angle) + 0.51 * jnp.cos(angle)) + np.testing.assert_allclose(samples.d2_R_d_phi2, -2.7 * jnp.cos(angle) + 0.99 * jnp.sin(angle)) + np.testing.assert_allclose(samples.d3_Z_d_phi3, 1.08 * jnp.sin(angle) - 4.59 * jnp.cos(angle)) + + +def test_axis_is_jittable_vmappable_and_differentiable(): + phi = jnp.linspace(0.0, 0.9, 11) + axis = Axis.stellarator_symmetric(rc=[1.0, 0.04], zs=[0.0, -0.04], nfp=3) + eager = evaluate_axis(axis, phi) + compiled = jax.jit(evaluate_axis)(axis, phi) + np.testing.assert_allclose(compiled.R, eager.R, rtol=1e-14, atol=1e-14) + + offsets = jnp.array([-0.01, 0.0, 0.02]) + + def radius_at_zero(offset): + changed = Axis(rc=axis.rc.at[1].add(offset), zs=axis.zs, nfp=axis.nfp) + return evaluate_axis(changed, 0.0).R + + np.testing.assert_allclose(jax.vmap(radius_at_zero)(offsets), 1.04 + offsets) + np.testing.assert_allclose(jax.grad(radius_at_zero)(0.0), 1.0) diff --git a/tests/unit/test_configurations.py b/tests/unit/test_configurations.py new file mode 100644 index 0000000..2015642 --- /dev/null +++ b/tests/unit/test_configurations.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import numpy as np +import pytest + +import pyqsc_jax as qsc + + +def test_named_configurations_are_immutable_and_solve(): + assert qsc.available_configurations() == ( + "qa", + "qh", + "finite_pressure_current", + "b20_optimized_qa", + "database_example_3", + "database_qa_139524", + "database_low_b20_57409", + "b20_optimized_good", + "database_large_singularity_107579", + "plasma_dominant_channel", + "plasma_stellarator", + ) + qa = qsc.get_configuration("qa") + first_parameters = qa.parameters(nphi=15) + second_parameters = qa.parameters(nphi=31) + + assert first_parameters["nphi"] == 15 + assert second_parameters["nphi"] == 31 + solution = qa.solve(nphi=15, order="r1") + np.testing.assert_allclose(solution.inputs.etabar, qa.etabar) + + +def test_solve_configuration_preserves_topology_and_overrides(): + qh = qsc.solve_configuration("qh", nphi=31, order="r1") + finite = qsc.solve_configuration( + "finite_pressure_current", + nphi=15, + ) + + assert int(qh.helicity) == 1 + assert finite.second_order is not None + np.testing.assert_allclose(finite.inputs.I2, 0.9) + np.testing.assert_allclose(finite.inputs.p2, -600000.0) + + +@pytest.mark.physics +def test_database_showcase_configurations_are_traceable_and_screened(): + expected_sources = { + "database_example_3": 3, + "database_qa_139524": 139524, + "database_low_b20_57409": 57409, + "b20_optimized_good": 57409, + "database_large_singularity_107579": 107579, + "plasma_stellarator": 52521, + } + + for name, database_id in expected_sources.items(): + configuration = qsc.get_configuration(name) + solution = configuration.solve(nphi=61) + minimum_abs_iota = 0.3 if name == "database_qa_139524" else 0.4 + criteria = qsc.Criteria.from_curvo_2025( + minimum_abs_iota=minimum_abs_iota, + ) + + assert configuration.source_database_id == database_id + assert configuration.source_url == ( + f"https://stellarator.physics.wisc.edu/app/plot/{database_id}" + ) + assert criteria.evaluate(solution).passed + assert abs(float(solution.iota)) >= minimum_abs_iota + if name == "database_qa_139524": + assert int(solution.helicity) == 0 + assert solution.inputs.axis.nfp == 1 + assert float(solution.inputs.I2) == 0.0 + assert float(solution.inputs.p2) != 0.0 + assert float(np.sqrt(np.mean(np.asarray(solution.torsion) ** 2))) > 1.0 + if name == "plasma_stellarator": + assert float(solution.inputs.I2) == 0.0 + assert float(solution.inputs.p2) != 0.0 + assert float(np.sqrt(np.mean(np.asarray(solution.torsion) ** 2))) > 0.5 + + +def test_unknown_configuration_is_rejected(): + with pytest.raises(ValueError, match="Available configurations"): + qsc.get_configuration("missing") diff --git a/tests/unit/test_plotting.py b/tests/unit/test_plotting.py new file mode 100644 index 0000000..f03705f --- /dev/null +++ b/tests/unit/test_plotting.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import matplotlib +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.plotting import ( + field_split_frenet_components, + plot_axis, + plot_b20, + plot_field_jet_norms, + plot_field_split_components, + plot_surface_3d, + surface_coordinates, +) + +matplotlib.use("Agg") + + +def test_axis_and_b20_plotters_return_objects(): + solution = qsc.solve_configuration("qa", nphi=31) + figure_axis, axis = plot_axis(solution, samples=31, label="QA") + figure_b20, b20_axis = plot_b20(solution, label="QA") + reused_axis_figure, reused_axis = plot_axis(solution, ax=axis, samples=31) + reused_b20_figure, reused_b20_axis = plot_b20(solution, ax=b20_axis) + + assert axis.figure is figure_axis + assert b20_axis.figure is figure_b20 + assert reused_axis_figure is figure_axis + assert reused_axis is axis + assert reused_b20_figure is figure_b20 + assert reused_b20_axis is b20_axis + assert len(axis.lines) == 2 + assert len(b20_axis.lines) == 2 + + +def test_field_jet_norm_plotter_and_axes_guard(): + solution = qsc.solve_configuration( + "finite_pressure_current", + nphi=31, + ) + result = qsc.plasma_hessian_on_axis( + solution, + formal_radius=0.05, + ) + figure, axes = plot_field_jet_norms(result) + + assert axes.shape == (3,) + assert all(len(axis.lines) == 3 for axis in axes) + assert all(np.isfinite(line.get_ydata()).all() for axis in axes for line in axis.lines) + with pytest.raises(ValueError, match="three"): + plot_field_jet_norms(result, axes=axes[:2]) + + +def test_surface_and_angle_dependent_field_split_plotters(): + solution = qsc.solve_configuration("plasma_stellarator", nphi=31) + x, y, z = surface_coordinates(solution, radius=0.05, ntheta=12) + figure, axis = plot_surface_3d(solution, radius=0.05, ntheta=12) + reused_figure, reused_axis = plot_surface_3d( + solution, + radius=0.05, + ntheta=12, + ax=axis, + plot_axis_line=False, + ) + result = qsc.plasma_hessian_on_axis(solution, formal_radius=0.1) + components = field_split_frenet_components(result, solution) + component_figure, axes = plot_field_split_components(result, solution) + + expected_toroidal_points = solution.inputs.axis.nfp * solution.inputs.nphi + 1 + assert x.shape == y.shape == z.shape == (12, expected_toroidal_points) + assert axis.figure is figure + assert reused_figure is figure + assert reused_axis is axis + assert len(axis.collections) == 2 + assert components.shape == (3, 3, 31) + assert axes.shape == (3,) + assert component_figure is axes[0].figure + assert all(len(item.lines) == 3 for item in axes) + with pytest.raises(ValueError, match="three"): + plot_field_split_components(result, solution, axes=axes[:2]) + + +@pytest.mark.physics +@pytest.mark.parametrize( + ("configuration", "radius"), + ( + ("database_qa_139524", 0.03), + ("database_example_3", 0.075), + ("b20_optimized_good", 0.075), + ("database_large_singularity_107579", 0.15), + ), +) +def test_database_surface_coordinates_are_finite_and_smooth(configuration, radius): + solution = qsc.solve_configuration(configuration, nphi=121) + x, y, z = surface_coordinates(solution, radius=radius, ntheta=36) + points = np.stack((x, y, z), axis=-1) + toroidal_edges = np.linalg.norm(np.diff(points, axis=1), axis=-1) + + assert np.all(np.isfinite(points)) + assert np.max(toroidal_edges) < 0.25 + + +def test_plotter_guards(): + first_order = qsc.solve_configuration("qa", nphi=15, order="r1") + with pytest.raises(ValueError, match="r2"): + plot_b20(first_order) + with pytest.raises(ValueError, match="integer"): + plot_axis(first_order, samples=3) + with pytest.raises(TypeError, match="Axis"): + plot_axis(object()) + with pytest.raises(ValueError, match="positive"): + surface_coordinates(first_order, radius=0) + with pytest.raises(ValueError, match="integer"): + surface_coordinates(first_order, ntheta=3) diff --git a/tests/unit/test_solvers.py b/tests/unit/test_solvers.py new file mode 100644 index 0000000..287e7b4 --- /dev/null +++ b/tests/unit/test_solvers.py @@ -0,0 +1,122 @@ +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from pyqsc_jax.solvers import ( + RootSolveOptions, + dense_newton_root, + implicit_dense_linear_solve, + implicit_dense_root, +) + + +def test_dense_newton_converges_with_report(): + root, report = dense_newton_root(lambda x: x**2 - 2.0, jnp.asarray(1.0)) + + np.testing.assert_allclose(root, jnp.sqrt(2.0), rtol=2e-13) + assert bool(report.converged) + assert bool(report.finite) + assert not bool(report.stagnated) + assert 0 < int(report.iterations) < 20 + assert report.residual_norm <= report.tolerance + np.testing.assert_allclose(report.jacobian_condition_number, 1.0) + + +def test_backtracking_is_exercised(): + root, report = dense_newton_root(lambda x: x**3 - 1.0, jnp.asarray(0.1)) + + np.testing.assert_allclose(root, 1.0, rtol=2e-12) + assert bool(report.converged) + assert int(report.backtracking_steps) > 0 + + +def test_iteration_limit_returns_failure_report(): + options = RootSolveOptions(max_steps=0) + root, report = dense_newton_root(lambda x: x**2 - 2.0, jnp.asarray(1.0), options=options) + + np.testing.assert_allclose(root, 1.0) + assert not bool(report.converged) + assert bool(report.finite) + assert int(report.iterations) == 0 + + +def test_implicit_root_gradient_uses_converged_equation(): + def root(parameter): + solved, _ = implicit_dense_root( + lambda x: x**2 - parameter, + jnp.asarray(1.0), + ) + return solved + + parameter = 2.0 + np.testing.assert_allclose(root(parameter), jnp.sqrt(parameter), rtol=2e-13) + np.testing.assert_allclose( + jax.grad(root)(parameter), + 1 / (2 * jnp.sqrt(parameter)), + rtol=2e-12, + ) + np.testing.assert_allclose(jax.jit(root)(parameter), root(parameter), rtol=2e-13) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"atol": -1.0}, + {"rtol": -1.0}, + {"step_tolerance": -1.0}, + {"max_steps": -1}, + {"max_backtracking_steps": -1}, + ], +) +def test_root_options_validation(kwargs): + with pytest.raises(ValueError): + RootSolveOptions(**kwargs) + + +def test_implicit_dense_linear_solve_and_gradient(): + right_hand_side = jnp.asarray([1.0, -0.5]) + + def objective(parameter): + matrix = jnp.asarray([[parameter, 0.2], [-0.4, 1.7]]) + solution, _ = implicit_dense_linear_solve(matrix, right_hand_side) + return jnp.sum(solution**2) + + parameter = 2.0 + matrix = jnp.asarray([[parameter, 0.2], [-0.4, 1.7]]) + solution, report = implicit_dense_linear_solve(matrix, right_hand_side) + np.testing.assert_allclose(solution, jnp.linalg.solve(matrix, right_hand_side)) + assert bool(report.converged) + assert bool(report.finite) + assert bool(report.well_conditioned) + assert report.relative_residual_norm < 1e-14 + + step = 2e-5 + finite_difference = (objective(parameter + step) - objective(parameter - step)) / (2 * step) + np.testing.assert_allclose(jax.grad(objective)(parameter), finite_difference, rtol=2e-8) + + +@pytest.mark.parametrize( + "matrix, right_hand_side, message", + [ + (jnp.ones(2), jnp.ones(2), "square"), + (jnp.ones((2, 3)), jnp.ones(2), "square"), + (jnp.eye(2), jnp.ones((2, 1)), "right_hand_side"), + (jnp.eye(2), jnp.ones(3), "right_hand_side"), + ], +) +def test_linear_solve_shape_validation(matrix, right_hand_side, message): + with pytest.raises(ValueError, match=message): + implicit_dense_linear_solve(matrix, right_hand_side) + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"residual_tolerance": -1.0}, "residual_tolerance"), + ({"condition_limit": 0.0}, "condition_limit"), + ], +) +def test_linear_solve_policy_validation(kwargs, message): + with pytest.raises(ValueError, match=message): + implicit_dense_linear_solve(jnp.eye(2), jnp.ones(2), **kwargs) diff --git a/tests/unit/test_spectral.py b/tests/unit/test_spectral.py new file mode 100644 index 0000000..214cd27 --- /dev/null +++ b/tests/unit/test_spectral.py @@ -0,0 +1,69 @@ +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from pyqsc_jax.spectral import ( + differentiate, + differentiation_matrix, + fourier_coefficients, + fourier_interpolate, + periodic_grid, + periodic_integral, +) + + +@pytest.mark.parametrize("n", [8, 9, 31]) +def test_differentiation_matrix_is_spectrally_exact(n): + period = 3.7 + x = periodic_grid(n, period=period) + fundamental = 2 * jnp.pi / period + values = jnp.sin(2 * fundamental * x) + 0.2 * jnp.cos(3 * fundamental * x) + expected = 2 * fundamental * jnp.cos(2 * fundamental * x) - 0.6 * fundamental * jnp.sin( + 3 * fundamental * x + ) + + derivative = differentiation_matrix(n, period=period) @ values + np.testing.assert_allclose(derivative, expected, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose(differentiate(values, period=period), expected, rtol=2e-13) + + +@pytest.mark.parametrize("n", [8, 9]) +def test_fourier_interpolation_even_and_odd(n): + x_grid = periodic_grid(n) + samples = jnp.sin(2 * x_grid) + 0.3 * jnp.cos(3 * x_grid) + x = jnp.array([0.12, 1.23, 5.72]) + expected = jnp.sin(2 * x) + 0.3 * jnp.cos(3 * x) + + np.testing.assert_allclose(fourier_interpolate(samples, x), expected, rtol=2e-13, atol=2e-13) + np.testing.assert_allclose( + fourier_interpolate(samples, x_grid), samples, rtol=2e-13, atol=2e-15 + ) + + +def test_periodic_integral_and_coefficients(): + n = 17 + x = periodic_grid(n) + values = 2.3 + jnp.cos(3 * x) - 0.4 * jnp.sin(5 * x) + frequency, coefficients = fourier_coefficients(values) + + np.testing.assert_allclose(periodic_integral(values), 2.3 * 2 * jnp.pi, rtol=2e-14) + np.testing.assert_allclose(coefficients[frequency == 0], 2.3, rtol=2e-14) + np.testing.assert_allclose(jax.jit(differentiate)(values), differentiate(values), rtol=2e-14) + + +def test_periodic_grid_validation(): + with pytest.raises(ValueError, match="n >= 2"): + periodic_grid(1) + with pytest.raises(ValueError, match="n >= 2"): + differentiation_matrix(1) + with pytest.raises(ValueError, match="one-dimensional"): + fourier_interpolate(jnp.ones((2, 3)), 0.2) + + +def test_complex_fourier_interpolation(): + x_grid = periodic_grid(9) + samples = jnp.exp(2j * x_grid) + x = jnp.array([0.17, 2.4]) + + np.testing.assert_allclose(fourier_interpolate(samples, x), jnp.exp(2j * x), rtol=2e-13) diff --git a/tests/unit/test_vmec.py b/tests/unit/test_vmec.py new file mode 100644 index 0000000..5cd7628 --- /dev/null +++ b/tests/unit/test_vmec.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +from hashlib import sha256 +from time import perf_counter + +import jax +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax.near_axis import near_axis + + +def qa_solution(*, nphi: int = 61): + return qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=nphi, + order="r2", + ) + + +def test_vectorized_surface_matches_independent_legacy_root_solve(): + field = near_axis( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, + order="r2", + ) + radius = 0.01 + old_R, old_Z, old_phi0 = field.Frenet_to_cylindrical(radius, ntheta=8) + new_R, new_Z, new_phi0, residual = qsc.uniform_cylindrical_surface( + field.solution, + radius, + ntheta=8, + ) + + np.testing.assert_allclose(new_R, old_R, atol=4.0e-13) + np.testing.assert_allclose(new_Z, old_Z, atol=4.0e-13) + np.testing.assert_allclose(new_phi0, old_phi0, atol=4.0e-13) + assert float(residual) < 2.0e-15 + + +def test_vmec_boundary_is_accurate_and_fast_after_compilation(tmp_path): + solution = qa_solution() + arguments = {"ntheta": 40, "mpol": 12, "ntor": 14} + compiled = qsc.vmec_boundary(solution, 0.03, **arguments) + jax.block_until_ready(compiled.RBC) + + start = perf_counter() + boundary = qsc.vmec_boundary(solution, 0.03, **arguments) + jax.block_until_ready(boundary.RBC) + warm_seconds = perf_counter() - start + qsc.to_vmec(solution, tmp_path / "input.first", r=0.03, **arguments) + start = perf_counter() + qsc.to_vmec(solution, tmp_path / "input.warm", r=0.03, **arguments) + total_warm_seconds = perf_counter() - start + + assert float(boundary.maximum_toroidal_angle_residual) < 2.0e-15 + assert bool(boundary.toroidal_angle_converged) + assert float(boundary.toroidal_angle_tolerance) > 0 + assert float(boundary.maximum_R_reconstruction_error) < 5.0e-6 + assert float(boundary.maximum_Z_reconstruction_error) < 5.0e-6 + assert warm_seconds < 0.5 + assert total_warm_seconds < 0.5 + + +def test_to_vmec_is_deterministic_and_legacy_adapter_exposes_coefficients(tmp_path): + solution = qa_solution(nphi=121) + output = tmp_path / "input.qa_r0025" + repeated_output = tmp_path / "input.qa_r0025.repeated" + export = qsc.to_vmec( + solution, + output, + r=0.0025, + ntheta=32, + mpol=6, + ntor=6, + parameters={ + "ns_array": (31,), + "ftol_array": (1.0e-10,), + "niter_array": (3000,), + }, + ) + qsc.to_vmec( + solution, + repeated_output, + r=0.0025, + ntheta=32, + mpol=6, + ntor=6, + parameters={ + "ns_array": (31,), + "ftol_array": (1.0e-10,), + "niter_array": (3000,), + }, + ) + + assert sha256(output.read_bytes()).digest() == sha256(repeated_output.read_bytes()).digest() + assert export.phiedge == np.pi * 0.0025**2 + assert not export.lasym + assert export.boundary.RBC.shape == (13, 7) + assert "&INDATA" in output.read_text(encoding="utf-8") + + field = near_axis( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=31, + order="r2", + ) + legacy_export = field.to_vmec( + tmp_path / "input.legacy", + r=0.01, + params={"mpol": 4, "ntor": 4}, + ntheta=16, + ntorMax=4, + ) + assert legacy_export.path.exists() + assert field.RBC.shape == field.RBS.shape == field.ZBC.shape == field.ZBS.shape == (5, 9) + + +def test_asymmetric_axis_emits_all_four_vmec_coefficient_families(tmp_path): + solution = qsc.Qsc( + rc=[1.0, 0.04], + rs=[0.0, 0.005], + zc=[0.0, 0.01], + zs=[0.0, -0.04], + nfp=3, + etabar=-0.9, + nphi=31, + order="r2", + ) + export = qsc.to_vmec( + solution, + tmp_path / "input.asymmetric", + r=0.01, + ntheta=16, + mpol=4, + ntor=4, + ) + + assert export.lasym + assert float(np.max(np.abs(export.boundary.RBS))) > 1.0e-4 + assert float(np.max(np.abs(export.boundary.ZBC))) > 1.0e-4 + contents = export.path.read_text(encoding="utf-8") + assert "LASYM = T" in contents + assert "RBS(" in contents + assert "ZBC(" in contents + + +def test_first_and_third_order_surfaces_and_ntor_cap(tmp_path): + first_order = qsc.Qsc( + rc=[1.0, 0.045], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=15, + order="r1", + ) + first_boundary = qsc.vmec_boundary( + first_order, + 0.005, + ntheta=8, + mpol=3, + ntor=2, + ) + assert np.all(np.isfinite(first_boundary.R)) + + third_order = qsc.solve_configuration("qa", nphi=15, order="r3") + export = qsc.to_vmec( + third_order, + tmp_path / "input.r3", + r=0.005, + ntheta=10, + mpol=4, + ntor=4, + ntor_max=2, + parameters=qsc.VmecInputParameters( + ns_array=(15,), + ftol_array=(1.0e-10,), + niter_array=(2000,), + ), + ) + assert export.boundary.ntor == 2 + assert np.all(np.isfinite(export.boundary.Z)) + + +def test_finite_pressure_current_profiles_are_written(tmp_path): + solution = qsc.solve_configuration("plasma_dominant_channel", nphi=15) + export = qsc.to_vmec( + solution, + tmp_path / "input.plasma", + r=0.1, + ntheta=10, + mpol=4, + ntor=2, + ) + contents = export.path.read_text(encoding="utf-8") + + np.testing.assert_allclose(export.curtor, 210000.0) + np.testing.assert_allclose(export.pressure_axis, 1000.0) + assert "AM =" in contents + assert "CURTOR =" in contents + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"ns_array": (), "ftol_array": (), "niter_array": ()}, "nonempty"), + ( + { + "ns_array": (15,), + "ftol_array": (1.0e-10, 1.0e-11), + "niter_array": (2000,), + }, + "aligned", + ), + ({"delt": 0.0}, "positive"), + ({"nstep": 0}, "positive"), + ({"tcon0": 0.0}, "positive"), + ( + {"ns_array": (2,), "ftol_array": (1.0e-10,), "niter_array": (2000,)}, + "at least 3", + ), + ( + {"ns_array": (15,), "ftol_array": (0.0,), "niter_array": (2000,)}, + "tolerance", + ), + ( + {"ns_array": (15,), "ftol_array": (1.0e-10,), "niter_array": (0,)}, + "iteration", + ), + ], +) +def test_vmec_input_parameter_guards(kwargs, message): + with pytest.raises(ValueError, match=message): + qsc.VmecInputParameters(**kwargs) + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"r": 0.0}, "r must be positive"), + ({"r": float("nan")}, "r must be positive"), + ({"ntor_max": -1}, "ntor_max"), + ({"coefficient_tolerance": -1.0}, "coefficient_tolerance"), + ({"coefficient_tolerance": float("nan")}, "coefficient_tolerance"), + ({"toroidal_angle_tolerance": -1.0}, "toroidal_angle_tolerance"), + ({"toroidal_angle_tolerance": float("nan")}, "toroidal_angle_tolerance"), + ({"ntheta": 7, "mpol": 3}, "ntheta"), + ({"mpol": 0}, "positive"), + ({"ntor": -1}, "nonnegative"), + ({"newton_iterations": 0}, "positive"), + ({"ntheta": 8.0}, "integers"), + ({"ntor": 8}, "nphi"), + ], +) +def test_to_vmec_resolution_and_scalar_guards(tmp_path, kwargs, message): + options = { + "r": 0.01, + "ntheta": 8, + "mpol": 3, + "ntor": 2, + } + options.update(kwargs) + with pytest.raises(ValueError, match=message): + qsc.to_vmec( + qa_solution(nphi=15), + tmp_path / "input.invalid", + **options, + ) + + +def test_to_vmec_rejects_unconverged_angle_inversion_without_writing(tmp_path): + output = tmp_path / "input.unconverged" + with pytest.raises(RuntimeError, match="no VMEC input was written"): + qsc.to_vmec( + qa_solution(), + output, + r=0.03, + ntheta=40, + mpol=12, + ntor=14, + newton_iterations=1, + ) + assert not output.exists() + + +def test_to_vmec_rejects_unknown_control(tmp_path): + with pytest.raises(ValueError, match="Unknown VMEC input"): + qsc.to_vmec( + qa_solution(nphi=15), + tmp_path / "input.invalid", + r=0.01, + ntheta=8, + mpol=3, + ntor=2, + parameters={"missing": 1}, + ) diff --git a/tests/unit/test_vmex_interface.py b/tests/unit/test_vmex_interface.py new file mode 100644 index 0000000..7b73f60 --- /dev/null +++ b/tests/unit/test_vmex_interface.py @@ -0,0 +1,263 @@ +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import pyqsc_jax as qsc +from pyqsc_jax import vmex as bridge + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class FakeParams: + rbc: jax.Array + rbs: jax.Array + zbc: jax.Array + zbs: jax.Array + phiedge: jax.Array + pres_scale: jax.Array + curtor: jax.Array + am: jax.Array + ai: jax.Array + ac: jax.Array + + +class FakeVmecInput: + def __init__(self, **kwargs): + for name, value in kwargs.items(): + setattr(self, name, value) + + +class FakeImplicit: + @staticmethod + def params_from_input(inp, *, device=None): + del device + return FakeParams( + rbc=jnp.asarray(inp.rbc), + rbs=jnp.asarray(inp.rbs), + zbc=jnp.asarray(inp.zbc), + zbs=jnp.asarray(inp.zbs), + phiedge=jnp.asarray(inp.phiedge), + pres_scale=jnp.asarray(inp.pres_scale), + curtor=jnp.asarray(inp.curtor), + am=jnp.asarray(inp.am), + ai=jnp.zeros(21), + ac=jnp.asarray(inp.ac), + ) + + @staticmethod + def run(_inp, params, **_kwargs): + return SimpleNamespace( + state=params, + runtime=SimpleNamespace(), + wb=jnp.sum(params.rbc**2), + wp=params.am[0], + ) + + @staticmethod + def iota_profile(state, _runtime): + return jnp.linspace(0.4, 0.5, 7) + 1.0e-4 * jnp.sum(state.rbc) + + +class FakeQuasisymmetry: + def __init__(self, surfaces, helicity_m, helicity_n): + self.surfaces = jnp.asarray(surfaces) + self.helicity = helicity_m + helicity_n + + def profile_state(self, state, _runtime): + return self.surfaces**2 + self.helicity + 1.0e-5 * jnp.sum(state.rbc) + + +class FakeOptimize: + QuasisymmetryRatioResidual = FakeQuasisymmetry + + @staticmethod + def magnetic_well(state, _runtime): + return 0.1 * state.am[0] + 1.0e-3 * jnp.sum(state.rbc) + + @staticmethod + def aspect_ratio(state, _runtime): + return 5.0 + 1.0e-3 * jnp.sum(state.rbc) + + @staticmethod + def volume(state, _runtime): + return jnp.abs(state.phiedge) * 10.0 + + +@pytest.fixture +def fake_vmex(monkeypatch): + module = SimpleNamespace( + __version__="test", + VmecInput=FakeVmecInput, + implicit=FakeImplicit, + optimize=FakeOptimize, + ) + monkeypatch.setattr(bridge, "_import_vmex", lambda: module) + return module + + +def qa_solution(*, asymmetric=False): + return qsc.Qsc( + rc=[1.0, 0.045], + rs=[0.0, 0.005] if asymmetric else [0.0, 0.0], + zc=[0.0, 0.01] if asymmetric else [0.0, 0.0], + zs=[0.0, -0.045], + nfp=3, + etabar=-0.9, + nphi=15, + order="r2", + ) + + +def make_problem(fake_vmex, solution=None, **kwargs): + del fake_vmex + options = { + "r": 0.02, + "ntheta": 8, + "mpol": 3, + "ntor": 2, + "ns_array": (7,), + "ftol": 1.0e-7, + "max_iterations": 200, + "multigrid": False, + "qs_surfaces": (0.5, 1.0), + } + options.update(kwargs) + return qsc.to_vmex_problem(solution if solution is not None else qa_solution(), **options) + + +def test_problem_builds_without_disk_and_exposes_quantities(fake_vmex): + problem = make_problem(fake_vmex) + result = problem.solve() + quantities = result.quantities + + assert problem.vmex_version == "test" + assert problem.validated_commit == qsc.VMEX_VALIDATED_COMMIT + assert problem.adjoint_tol == 1.0e-11 + assert problem.input.mpol == 4 + assert problem.boundary.RBC.shape == (5, 4) + assert not problem.finite_beta + assert quantities.s.shape == quantities.iota.shape == (7,) + np.testing.assert_allclose(quantities.iota, -quantities.iota_vmec) + assert quantities.quasisymmetry.shape == (2,) + assert np.isfinite(float(quantities.magnetic_well)) + + +def test_finite_beta_and_traceable_parameter_remap(fake_vmex): + solution = qsc.solve_configuration("plasma_stellarator", nphi=15) + problem = make_problem(fake_vmex, solution=solution) + assert problem.finite_beta + + def objective(etabar): + candidate = qsc.solve( + axis=solution.inputs.axis, + etabar=etabar, + B0=solution.inputs.B0, + B2c=solution.inputs.B2c, + I2=solution.inputs.I2, + p2=solution.inputs.p2, + nphi=15, + order="r2", + ) + parameters = problem.parameters_for(candidate, radius=0.018) + return qsc.vmex_radial_quantities(problem, parameters).magnetic_well + + value, derivative = jax.value_and_grad(objective)(solution.inputs.etabar) + assert np.isfinite(float(value)) + assert np.isfinite(float(derivative)) + assert float(problem.quantities().thermal_energy) > 0 + + +def test_empty_qs_surface_list_supports_asymmetric_equilibrium(fake_vmex): + problem = make_problem(fake_vmex, solution=qa_solution(asymmetric=True), qs_surfaces=()) + assert problem.input.lasym + assert problem.quantities().quasisymmetry.shape == (0,) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"r": 0.0}, "positive"), + ({"qs_surfaces": (0.0,)}, "0 < s"), + ({"qs_surfaces": (0.75, 0.5)}, "increasing"), + ({"ns_array": ()}, "nonempty"), + ({"ns_array": (7, 7)}, "increasing"), + ({"ftol": 0.0}, "positive"), + ({"ftol_array": (1.0e-7, 1.0e-8)}, "one positive"), + ({"max_iterations": 0}, "positive integer"), + ({"adjoint_tol": 0.0}, "positive and finite"), + ({"toroidal_angle_tolerance": -1.0}, "nonnegative and finite"), + ({"newton_iterations": 0}, "positive"), + ({"ntheta": 7}, "ntheta"), + ({"helicity_m": 1.5}, "integer"), + ({"helicity_n": 1.5}, "integer"), + ], +) +def test_problem_validation(fake_vmex, kwargs, message): + with pytest.raises(ValueError, match=message): + make_problem(fake_vmex, **kwargs) + + +def test_asymmetric_quasisymmetry_guard(fake_vmex): + with pytest.raises(NotImplementedError, match="stellarator-symmetric"): + make_problem(fake_vmex, solution=qa_solution(asymmetric=True)) + + +def test_unconverged_vmex_boundary_is_rejected(fake_vmex): + with pytest.raises(RuntimeError, match="angle inversion did not converge"): + make_problem(fake_vmex, newton_iterations=1) + + +def test_parameter_remap_rejects_changed_field_period(fake_vmex): + problem = make_problem(fake_vmex) + changed = qsc.Qsc( + rc=[1.0, 0.02], + zs=[0.0, -0.02], + nfp=2, + etabar=-0.9, + nphi=15, + order="r2", + ) + with pytest.raises(ValueError, match="field periods"): + problem.parameters_for(changed) + + +def test_runtime_and_import_guards(fake_vmex, monkeypatch): + problem = make_problem(fake_vmex) + vmex = bridge._import_vmex() + monkeypatch.setattr( + vmex.implicit, + "run", + lambda *_args, **_kwargs: SimpleNamespace( + state=problem.parameters, + runtime=None, + wb=0.0, + wp=0.0, + ), + ) + with pytest.raises(RuntimeError, match="runtime"): + problem.solve() + + monkeypatch.undo() + monkeypatch.setattr( + bridge.importlib, + "import_module", + lambda _name: (_ for _ in ()).throw(ImportError("missing")), + ) + with pytest.raises(ImportError, match="requires VMEX"): + bridge._import_vmex() + + +def test_incomplete_vmex_api_is_rejected(monkeypatch): + monkeypatch.setattr( + bridge.importlib, + "import_module", + lambda _name: SimpleNamespace(VmecInput=FakeVmecInput), + ) + with pytest.raises(ImportError, match="implicit, optimize"): + bridge._import_vmex()