vmec-jax 0.0.1__tar.gz

vmec_jax-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 UW Plasma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

vmec_jax-0.0.1/PKG-INFO

@@ -0,0 +1,244 @@
+Metadata-Version: 2.4
+Name: vmec-jax
+Version: 0.0.1
+Summary: End-to-end differentiable JAX implementation of VMEC2000 for fixed and free-boundary equilibria.
+Author: vmec_jax contributors
+License-Expression: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: netCDF4
+Requires-Dist: tomli; python_version < "3.11"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Requires-Dist: jaxlib; extra == "jax"
+Provides-Extra: netcdf
+Requires-Dist: netCDF4; extra == "netcdf"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: furo; extra == "docs"
+Provides-Extra: plots
+Requires-Dist: matplotlib; extra == "plots"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: types-setuptools; extra == "dev"
+Dynamic: license-file
+
+# vmec-jax
+
+End-to-end differentiable JAX implementation of **VMEC2000** for fixed-boundary
+and free-boundary ideal-MHD equilibria.
+
+## Showcase (single-grid)
+
+All figures below use the same **single-grid** run settings: `NS_ARRAY=151`, `NITER_ARRAY=5000`, `FTOL_ARRAY=1e-14`, `NSTEP=500`.
+
+<table>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_cross_sections.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_cross_sections.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> cross-section (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> cross-section (VMEC2000 vs vmec_jax)</td>
+  </tr>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_iota.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_iota.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> iota (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> iota (VMEC2000 vs vmec_jax)</td>
+  </tr>
+</table>
+
+<p align="center">
+  <img src="docs/_static/figures/readme_fsq_trace_single_grid.png" width="860" />
+</p>
+
+<p align="center">
+  <img src="docs/_static/figures/readme_runtime_compare.png" width="860" />
+</p>
+
+**Cold vs warm runtime**: the *cold* bar includes XLA JIT compilation on the first call (one-time cost per process); the *warm* bar is the steady-state solve time for all subsequent calls in the same process, with the compiled kernels already in-memory. VMEC2000 is a pre-compiled Fortran binary and therefore has no compilation overhead — it is always effectively "cold". The warm vmec_jax time is the fair comparison for repeated solves (e.g., in an optimization loop). Starting from v0.2, vmec_jax automatically caches compiled XLA kernels to disk (`~/.cache/vmec_jax/jax_cache`), so that *cold* runs in a fresh process on the same machine after the first invocation benefit from the on-disk cache and approach warm-run speed.
+
+## More visuals (single-grid)
+
+<table>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_3d.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_3d.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> 3D LCFS (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> 3D LCFS (VMEC2000 vs vmec_jax)</td>
+  </tr>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_bmag_surface.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_bmag_surface.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> |B| on LCFS (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> |B| on LCFS (VMEC2000 vs vmec_jax)</td>
+  </tr>
+</table>
+
+## What it is
+
+- VMEC2000-parity solver for fixed-boundary and free-boundary equilibria.
+- Supports axisymmetric and non-axisymmetric configurations, with `lasym=False` and `lasym=True` for stellarator symmetry/asymmetry and up-down symmetry/asymmetry.
+- Default CLI path is `vmec_jax input.name`.
+- `wout_*.nc` outputs, iteration diagnostics, and manifest-based parity sweeps are built around VMEC2000-compatible workflows.
+- JAX-native kernels for geometry, transforms, and residual assembly.
+- Differentiable optimization workflows are available through the Python API and bundled examples.
+
+## Quickstart
+
+Install and run the showcase:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install -e .
+python examples/showcase_axisym_input_to_wout.py --suite
+```
+
+If you want a release-style non-editable install instead, use:
+
+```bash
+python -m pip install .
+```
+
+If you want the bundled reference outputs and mgrid files, fetch the assets once:
+
+```bash
+python tools/fetch_assets.py
+```
+
+Lightweight clone (keeps full history, downloads blobs lazily):
+
+```bash
+git clone --filter=blob:none https://github.com/uwplasma/vmec_jax
+```
+
+Note: the repo history was rewritten on 2026-03-16 to remove large assets from
+all commits. If you cloned before that date, please re-clone (or prune and
+reset) to get the smaller history.
+
+CLI (VMEC2000-style executable):
+
+```bash
+vmec_jax examples/data/input.circular_tokamak
+```
+
+Sanity check (verifies the console script is wired to the right interpreter):
+
+```bash
+vmec_jax --help
+```
+
+If the `vmec_jax` command is not found or raises `ModuleNotFoundError`, make sure
+you installed with the same interpreter and use the module entrypoint:
+
+```bash
+python -m pip install -e .
+python -m vmec_jax examples/data/input.circular_tokamak
+```
+
+For fixed-boundary inputs, the default CLI path now uses the optimized
+controller: it tries the fast final-grid scan route first, then escalates to
+staged continuation and strict parity finishing only when the input structure
+and residual history require it. Pass `--parity` to force the conservative
+VMEC2000 loop. Pass `--solver-mode accelerated` to request the optimized track
+explicitly.
+
+Python driver comparison (reference track vs optimized CLI-style track):
+
+```bash
+python examples/fixed_boundary_driver_tracks.py \
+  examples/data/input.circular_tokamak \
+  --quiet --json
+```
+
+Run tests:
+
+```bash
+pytest -q
+```
+
+Full test suite (requires netCDF assets):
+
+```bash
+python tools/fetch_assets.py
+RUN_FULL=1 pytest -q
+```
+
+Advanced optimization examples live in `examples/optimization/`. They are
+intended as deeper workflow templates rather than README quickstarts, so use
+the fixed-boundary driver example above as the validated copy/paste entry point
+and then adapt the optimization scripts for your target objective. The simplest
+starting point is:
+
+```bash
+python examples/optimization/target_iota_aspect_volume.py --opt-steps 2
+```
+
+That example keeps the boundary parameterization small (`max |m|,|n| <= 1`),
+targets equilibrium volume, aspect ratio, and mean iota, and defaults to the
+bundled current-driven `cth_like_fixed_bdy` case so the iota channel is active.
+
+## Performance vs parity
+
+- Default runs aim for VMEC2000-compatible behavior while selecting the fastest stable path for the input.
+- Use `--parity` (or `performance_mode=False` in Python) to force the conservative VMEC2000 loop.
+- Use `--solver-mode accelerated` to force the optimized fixed-boundary controller explicitly.
+
+Details, profiling guidance, and parity methodology:
+
+- `docs/performance.rst`
+- `docs/validation.rst`
+- `tools/diagnostics/parity_manifest.toml` + `tools/diagnostics/parity_sweep_manifest.py`
+
+## VMEC++ notes
+
+The current runtime benchmark compares vmec_jax against VMEC2000. VMEC++ is not included in this benchmark.
+
+When VMEC++ is available, it can be added to the runtime plot via `--cpu-summary` entries with `backend=vmecpp`. Some inputs are not supported or do not converge under the same single-grid settings:
+
+VMEC++ unsupported inputs (`lasym=True`):
+
+- `LandremanSenguptaPlunk_section5p3_low_res`
+- `basic_non_stellsym_pressure`
+- `cth_like_free_bdy_lasym_small`
+- `up_down_asymmetric_tokamak`
+
+VMEC++ known non-convergence on these `lasym=False` cases under the same single-grid settings:
+
+- `DIII-D_lasym_false`
+- `LandremanPaul2021_QA_reactorScale_lowres`
+- `LandremanPaul2021_QH_reactorScale_lowres`
+- `LandremanSengupta2019_section5.4_B2_A80`
+- `cth_like_fixed_bdy`
+
+## CLI output and `NSTEP`
+
+The VMEC-style iteration loop prints every `NSTEP` iterations. Larger `NSTEP` means fewer print callbacks and faster runs.
+
+To disable live printing, set:
+
+```bash
+export VMEC_JAX_SCAN_PRINT=0
+```
+
+Quiet runs (`--quiet` or `verbose=False`) default the scan path to a minimal
+history mode (only `fsqr/fsqz/fsql` and `w_history` are kept) to reduce
+host/device traffic. You can override this with:
+
+```bash
+export VMEC_JAX_SCAN_MINIMAL=0  # keep full scan diagnostics even when quiet
+```

vmec_jax-0.0.1/README.md

@@ -0,0 +1,213 @@
+# vmec-jax
+
+End-to-end differentiable JAX implementation of **VMEC2000** for fixed-boundary
+and free-boundary ideal-MHD equilibria.
+
+## Showcase (single-grid)
+
+All figures below use the same **single-grid** run settings: `NS_ARRAY=151`, `NITER_ARRAY=5000`, `FTOL_ARRAY=1e-14`, `NSTEP=500`.
+
+<table>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_cross_sections.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_cross_sections.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> cross-section (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> cross-section (VMEC2000 vs vmec_jax)</td>
+  </tr>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_iota.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_iota.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> iota (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> iota (VMEC2000 vs vmec_jax)</td>
+  </tr>
+</table>
+
+<p align="center">
+  <img src="docs/_static/figures/readme_fsq_trace_single_grid.png" width="860" />
+</p>
+
+<p align="center">
+  <img src="docs/_static/figures/readme_runtime_compare.png" width="860" />
+</p>
+
+**Cold vs warm runtime**: the *cold* bar includes XLA JIT compilation on the first call (one-time cost per process); the *warm* bar is the steady-state solve time for all subsequent calls in the same process, with the compiled kernels already in-memory. VMEC2000 is a pre-compiled Fortran binary and therefore has no compilation overhead — it is always effectively "cold". The warm vmec_jax time is the fair comparison for repeated solves (e.g., in an optimization loop). Starting from v0.2, vmec_jax automatically caches compiled XLA kernels to disk (`~/.cache/vmec_jax/jax_cache`), so that *cold* runs in a fresh process on the same machine after the first invocation benefit from the on-disk cache and approach warm-run speed.
+
+## More visuals (single-grid)
+
+<table>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_3d.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_3d.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> 3D LCFS (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> 3D LCFS (VMEC2000 vs vmec_jax)</td>
+  </tr>
+  <tr>
+    <td><img src="docs/_static/figures/axisym_compare_bmag_surface.png" width="420" /></td>
+    <td><img src="docs/_static/figures/qa_compare_bmag_surface.png" width="420" /></td>
+  </tr>
+  <tr>
+    <td align="center"><code>ITERModel</code> |B| on LCFS (VMEC2000 vs vmec_jax)</td>
+    <td align="center"><code>LandremanPaul2021_QA_lowres</code> |B| on LCFS (VMEC2000 vs vmec_jax)</td>
+  </tr>
+</table>
+
+## What it is
+
+- VMEC2000-parity solver for fixed-boundary and free-boundary equilibria.
+- Supports axisymmetric and non-axisymmetric configurations, with `lasym=False` and `lasym=True` for stellarator symmetry/asymmetry and up-down symmetry/asymmetry.
+- Default CLI path is `vmec_jax input.name`.
+- `wout_*.nc` outputs, iteration diagnostics, and manifest-based parity sweeps are built around VMEC2000-compatible workflows.
+- JAX-native kernels for geometry, transforms, and residual assembly.
+- Differentiable optimization workflows are available through the Python API and bundled examples.
+
+## Quickstart
+
+Install and run the showcase:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install -e .
+python examples/showcase_axisym_input_to_wout.py --suite
+```
+
+If you want a release-style non-editable install instead, use:
+
+```bash
+python -m pip install .
+```
+
+If you want the bundled reference outputs and mgrid files, fetch the assets once:
+
+```bash
+python tools/fetch_assets.py
+```
+
+Lightweight clone (keeps full history, downloads blobs lazily):
+
+```bash
+git clone --filter=blob:none https://github.com/uwplasma/vmec_jax
+```
+
+Note: the repo history was rewritten on 2026-03-16 to remove large assets from
+all commits. If you cloned before that date, please re-clone (or prune and
+reset) to get the smaller history.
+
+CLI (VMEC2000-style executable):
+
+```bash
+vmec_jax examples/data/input.circular_tokamak
+```
+
+Sanity check (verifies the console script is wired to the right interpreter):
+
+```bash
+vmec_jax --help
+```
+
+If the `vmec_jax` command is not found or raises `ModuleNotFoundError`, make sure
+you installed with the same interpreter and use the module entrypoint:
+
+```bash
+python -m pip install -e .
+python -m vmec_jax examples/data/input.circular_tokamak
+```
+
+For fixed-boundary inputs, the default CLI path now uses the optimized
+controller: it tries the fast final-grid scan route first, then escalates to
+staged continuation and strict parity finishing only when the input structure
+and residual history require it. Pass `--parity` to force the conservative
+VMEC2000 loop. Pass `--solver-mode accelerated` to request the optimized track
+explicitly.
+
+Python driver comparison (reference track vs optimized CLI-style track):
+
+```bash
+python examples/fixed_boundary_driver_tracks.py \
+  examples/data/input.circular_tokamak \
+  --quiet --json
+```
+
+Run tests:
+
+```bash
+pytest -q
+```
+
+Full test suite (requires netCDF assets):
+
+```bash
+python tools/fetch_assets.py
+RUN_FULL=1 pytest -q
+```
+
+Advanced optimization examples live in `examples/optimization/`. They are
+intended as deeper workflow templates rather than README quickstarts, so use
+the fixed-boundary driver example above as the validated copy/paste entry point
+and then adapt the optimization scripts for your target objective. The simplest
+starting point is:
+
+```bash
+python examples/optimization/target_iota_aspect_volume.py --opt-steps 2
+```
+
+That example keeps the boundary parameterization small (`max |m|,|n| <= 1`),
+targets equilibrium volume, aspect ratio, and mean iota, and defaults to the
+bundled current-driven `cth_like_fixed_bdy` case so the iota channel is active.
+
+## Performance vs parity
+
+- Default runs aim for VMEC2000-compatible behavior while selecting the fastest stable path for the input.
+- Use `--parity` (or `performance_mode=False` in Python) to force the conservative VMEC2000 loop.
+- Use `--solver-mode accelerated` to force the optimized fixed-boundary controller explicitly.
+
+Details, profiling guidance, and parity methodology:
+
+- `docs/performance.rst`
+- `docs/validation.rst`
+- `tools/diagnostics/parity_manifest.toml` + `tools/diagnostics/parity_sweep_manifest.py`
+
+## VMEC++ notes
+
+The current runtime benchmark compares vmec_jax against VMEC2000. VMEC++ is not included in this benchmark.
+
+When VMEC++ is available, it can be added to the runtime plot via `--cpu-summary` entries with `backend=vmecpp`. Some inputs are not supported or do not converge under the same single-grid settings:
+
+VMEC++ unsupported inputs (`lasym=True`):
+
+- `LandremanSenguptaPlunk_section5p3_low_res`
+- `basic_non_stellsym_pressure`
+- `cth_like_free_bdy_lasym_small`
+- `up_down_asymmetric_tokamak`
+
+VMEC++ known non-convergence on these `lasym=False` cases under the same single-grid settings:
+
+- `DIII-D_lasym_false`
+- `LandremanPaul2021_QA_reactorScale_lowres`
+- `LandremanPaul2021_QH_reactorScale_lowres`
+- `LandremanSengupta2019_section5.4_B2_A80`
+- `cth_like_fixed_bdy`
+
+## CLI output and `NSTEP`
+
+The VMEC-style iteration loop prints every `NSTEP` iterations. Larger `NSTEP` means fewer print callbacks and faster runs.
+
+To disable live printing, set:
+
+```bash
+export VMEC_JAX_SCAN_PRINT=0
+```
+
+Quiet runs (`--quiet` or `verbose=False`) default the scan path to a minimal
+history mode (only `fsqr/fsqz/fsql` and `w_history` are kept) to reduce
+host/device traffic. You can override this with:
+
+```bash
+export VMEC_JAX_SCAN_MINIMAL=0  # keep full scan diagnostics even when quiet
+```

vmec_jax-0.0.1/pyproject.toml

@@ -0,0 +1,108 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vmec-jax"
+version = "0.0.1"
+description = "End-to-end differentiable JAX implementation of VMEC2000 for fixed and free-boundary equilibria."
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [{name = "vmec_jax contributors"}]
+license = "MIT"
+
+dependencies = [
+  "numpy",
+  "jax",
+  "jaxlib",
+  "netCDF4",
+  "tomli; python_version < '3.11'",
+]
+
+[project.scripts]
+vmec_jax = "vmec_jax.cli:main"
+xvmec_jax = "vmec_jax.cli:main"
+
+[project.optional-dependencies]
+jax = ["jax", "jaxlib"]
+netcdf = ["netCDF4"]
+docs = [
+  "sphinx",
+  "furo",
+]
+plots = ["matplotlib"]
+dev = ["pytest", "ruff", "mypy", "types-setuptools"]
+
+[tool.setuptools.packages.find]
+where = ["."]
+exclude = ["tests*", "docs*", "examples*", "tools*"]
+
+[tool.pytest.ini_options]
+markers = [
+  "vmec2000: integration tests requiring a local VMEC2000 build (+ mpi4py/netCDF4)",
+  "full: tests that require large netCDF assets (run tools/fetch_assets.py)",
+]
+
+[tool.ruff]
+line-length = 120
+
+[tool.ruff.lint]
+# E402: module-level import not at top-of-file.  This codebase uses deferred
+#        imports inside functions deliberately to avoid triggering JAX/GPU
+#        initialisation at import time.  Silencing E402 globally is intentional.
+# F821: undefined name in annotation.  Ruff occasionally reports false positives
+#        inside very long functions (closures over outer-scope locals) and for
+#        TYPE_CHECKING-only imports used in annotations.  Silence globally and
+#        rely on mypy for real undefined-name checks.
+# F841: local variable is assigned to but never used.  Many VMEC routines assign
+#        intermediate diagnostics for debugging that are not always consumed by
+#        the caller.  mypy --warn-unused-variables or a per-symbol noqa is a
+#        better gate here.
+# E501: line too long.  Handled by the line-length setting above; long lines in
+#        generated / ported Fortran-style code are acceptable.
+ignore = ["E402", "F821", "F841", "E501"]
+
+# Allow single-statement blocks on the same line as the colon (E701) and
+# semicolons as statement separators (E702) in tightly coupled numerical kernels.
+# These are common in ported Fortran code and do not impair readability.
+# E731: lambda-assignment — allow where lambdas are used as concise callbacks.
+# E741: ambiguous variable name — l/O/I are used as loop indices in physics code.
+extend-ignore = ["E701", "E702", "E731", "E741"]
+
+[tool.mypy]
+# vmec_jax is a JAX/NumPy numerical physics codebase ported from Fortran.
+# Many type errors are pre-existing false positives stemming from:
+#   - JAX's dynamically-typed array operations (jax.Array is opaque to mypy)
+#   - Very long functions with nested closures (mypy loses type context)
+#   - TYPE_CHECKING-guarded imports used in annotations
+#   - Ported Fortran patterns (lots of Any-typed intermediate arrays)
+# We configure mypy conservatively: catch obvious bugs but avoid blocking
+# development with spurious false-positive noise.
+python_version = "3.10"
+ignore_missing_imports = true
+no_strict_optional = true
+warn_return_any = false
+warn_unused_ignores = false
+# Disable the most noisy checks that produce mostly false positives here:
+disable_error_code = ["attr-defined", "operator", "valid-type", "name-defined", "call-arg", "return-value", "arg-type", "assignment", "index", "union-attr", "misc", "type-arg", "no-any-return", "override"]
+
+[[tool.mypy.overrides]]
+module = [
+  "vmec_jax.solve",
+  "vmec_jax.driver",
+  "vmec_jax.boundary",
+  "vmec_jax.vmec_forces",
+  "vmec_jax.geom",
+  "vmec_jax.wout",
+  "vmec_jax.vmec_parity",
+  "vmec_jax.fourier",
+  "vmec_jax.static",
+  "vmec_jax.implicit",
+  "vmec_jax.diagnostics",
+  "vmec_jax.free_boundary",
+  "vmec_jax.init_guess",
+  "vmec_jax.booz_input",
+  "vmec_jax.vmec_residue",
+  "vmec_jax.vmec_tomnsp",
+]
+ignore_errors = true

vmec_jax-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

vmec_jax-0.0.1/tests/test_bcovar_lambda_axis_closure.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import numpy as np
+
+from vmec_jax.vmec_bcovar import _apply_vmec_lambda_axis_closure
+
+
+def test_lambda_axis_closure_copies_m0_npos_modes():
+    lsin = np.array(
+        [
+            [1.0, 2.0, 3.0, 4.0],
+            [10.0, 20.0, 30.0, 40.0],
+            [100.0, 200.0, 300.0, 400.0],
+        ],
+        dtype=float,
+    )
+    m_modes = np.array([0, 0, 1, 0], dtype=int)
+    n_modes = np.array([0, 1, 1, 2], dtype=int)
+
+    out = np.asarray(
+        _apply_vmec_lambda_axis_closure(
+            Lsin=lsin,
+            m_modes=m_modes,
+            n_modes=n_modes,
+            lthreed=True,
+            ntor=2,
+        )
+    )
+
+    np.testing.assert_allclose(out[0, 0], lsin[0, 0])  # n=0 unchanged
+    np.testing.assert_allclose(out[0, 1], lsin[1, 1])  # m=0,n>0 copied
+    np.testing.assert_allclose(out[0, 2], lsin[0, 2])  # m!=0 unchanged
+    np.testing.assert_allclose(out[0, 3], lsin[1, 3])  # m=0,n>0 copied
+    np.testing.assert_allclose(out[1:], lsin[1:])  # interior unchanged
+
+
+def test_lambda_axis_closure_disabled_for_axisymmetric_or_ntor_zero():
+    lsin = np.array([[1.0, 2.0], [10.0, 20.0]], dtype=float)
+    m_modes = np.array([0, 0], dtype=int)
+    n_modes = np.array([0, 1], dtype=int)
+
+    out_axis = np.asarray(
+        _apply_vmec_lambda_axis_closure(
+            Lsin=lsin,
+            m_modes=m_modes,
+            n_modes=n_modes,
+            lthreed=False,
+            ntor=1,
+        )
+    )
+    out_ntor0 = np.asarray(
+        _apply_vmec_lambda_axis_closure(
+            Lsin=lsin,
+            m_modes=m_modes,
+            n_modes=n_modes,
+            lthreed=True,
+            ntor=0,
+        )
+    )
+    np.testing.assert_allclose(out_axis, lsin)
+    np.testing.assert_allclose(out_ntor0, lsin)
+
+
+def test_lambda_axis_closure_noop_when_ns_one():
+    lsin = np.array([[1.0, 2.0, 3.0]], dtype=float)
+    m_modes = np.array([0, 0, 1], dtype=int)
+    n_modes = np.array([0, 2, 1], dtype=int)
+
+    out = np.asarray(
+        _apply_vmec_lambda_axis_closure(
+            Lsin=lsin,
+            m_modes=m_modes,
+            n_modes=n_modes,
+            lthreed=True,
+            ntor=2,
+        )
+    )
+    np.testing.assert_allclose(out, lsin)