vmex 0.2.0__py3-none-any.whl

vmec_jax/__init__.py

@@ -0,0 +1,45 @@
+"""Compatibility shim: ``vmec_jax`` was renamed to :mod:`vmex`.
+
+Importing ``vmec_jax`` re-exports everything from :mod:`vmex` and emits a
+``DeprecationWarning``.  Submodule access (``vmec_jax.core.solver`` etc.) is
+forwarded by aliasing the already-imported ``vmex`` submodules into
+``sys.modules`` under the old names.  This shim is a one-release courtesy;
+update imports to ``vmex``.
+"""
+from __future__ import annotations
+
+import importlib
+import sys
+import warnings
+
+warnings.warn(
+    "the 'vmec_jax' package has been renamed to 'vmex'; import 'vmex' instead "
+    "(this compatibility shim will be removed in a future release)",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+import vmex as _vmex
+
+# Re-export the public API.
+from vmex import *  # noqa: F401,F403
+__version__ = _vmex.__version__
+if hasattr(_vmex, "__all__"):
+    __all__ = list(_vmex.__all__)
+
+# Forward submodule access: alias every imported vmex submodule under the old
+# top-level name so `import vmec_jax.core.X` and `from vmec_jax.core import X`
+# resolve to the identical vmex module object.
+for _name, _mod in list(sys.modules.items()):
+    if _name == "vmex" or _name.startswith("vmex."):
+        sys.modules["vmec_jax" + _name[len("vmex"):]] = _mod
+
+
+def __getattr__(name: str):
+    """Lazily import and forward any vmex submodule not yet loaded."""
+    try:
+        mod = importlib.import_module(f"vmex.{name}")
+    except ModuleNotFoundError as exc:  # pragma: no cover - mirror vmex
+        raise AttributeError(name) from exc
+    sys.modules[f"vmec_jax.{name}"] = mod
+    return mod

vmex/__init__.py

@@ -0,0 +1,131 @@
+"""vmex: a JAX implementation of VMEC2000 for fixed and free-boundary equilibria.
+
+Public API (lazily imported; ``import vmex as vj``):
+
+- :class:`~vmex.core.input.VmecInput` — INDATA / VMEC++-JSON input pytree
+- :func:`~vmex.core.solver.solve` — single-grid fixed-boundary solve
+- :func:`~vmex.core.multigrid.solve_multigrid` — NS_ARRAY ladder (runvmec.f)
+- :func:`~vmex.core.freeboundary.solve_free_boundary` — NESTOR free boundary
+- :func:`~vmex.core.wout.read_wout` / :func:`~vmex.core.wout.write_wout`
+  / :func:`~vmex.core.wout.wout_from_state` / :class:`~vmex.core.wout.WoutData`
+- :func:`~vmex.core.plotting.plot_wout` / :func:`~vmex.core.plotting.plot_boozmn`
+- :func:`~vmex.core.boozer.run_booz_xform` — Boozer transform (booz_xform_jax)
+- :func:`~vmex.core.mgrid.read_mgrid` / :func:`~vmex.core.mgrid.write_mgrid`
+  / :class:`~vmex.core.mgrid.MgridField` (external field is an mgrid or any
+  ``xyz->B`` callable; coils live in ESSOS, ``essos.coils.Coils``)
+- ``vmex.optimize`` — objectives + least-squares driver (module)
+- ``vmex.implicit`` — implicit differentiation of the equilibrium (module)
+- ``vmex.parallel`` — concurrent ensembles of independent solves (module)
+- ``vmex.errors`` — typed zero-crash exceptions (also exported directly)
+
+The ``vmec`` console entry point lives in :mod:`vmex.core.cli`.
+"""
+
+from importlib import import_module as _import_module
+from importlib.metadata import PackageNotFoundError as _PackageNotFoundError
+from importlib.metadata import version as _package_version
+import os as _os
+from pathlib import Path as _Path
+
+from ._compat import _default_compilation_cache_dir as _default_jax_cache_dir
+
+
+def _source_tree_version() -> str | None:
+    pyproject = _Path(__file__).resolve().parents[1] / "pyproject.toml"
+    if not pyproject.exists():
+        return None
+    in_project = False
+    for raw_line in pyproject.read_text(encoding="utf-8").splitlines():
+        line = raw_line.strip()
+        if line == "[project]":
+            in_project = True
+            continue
+        if in_project and line.startswith("["):
+            return None
+        if in_project and line.startswith("version"):
+            return line.split("=", 1)[1].strip().strip('"')
+    return None
+
+
+try:
+    __version__ = _source_tree_version() or _package_version("vmex")
+except _PackageNotFoundError:  # pragma: no cover - source tree without installed metadata.
+    __version__ = "0+unknown"
+
+# Suppress noisy C++ warnings from XLA/PjRt backend (e.g. repeated
+# "Assume version compatibility. PjRt-IFRT does not track XLA executable
+# versions." on persistent-cache hits). Must be set before *any* ``import
+# jax`` in the process. Uses setdefault so the user can still override via the
+# environment.
+_os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
+_os.environ.setdefault("ABSL_MIN_LOG_LEVEL", "2")
+_os.environ.setdefault("GLOG_minloglevel", "2")
+
+# Enable the JAX persistent XLA compilation cache in a machine-scoped
+# directory when requested by the backend/env policy in _compat. Accelerator
+# runs use the cache by default; CPU runs are opt-in to avoid XLA:CPU AOT
+# feature-mismatch warnings on shared or changing runtime environments.
+# ``core.solver._harden_compilation_cache`` re-applies this policy on every
+# solve path in case this module never ran (namespace-package shadowing).
+import jax as _jax
+
+_jax_cache_dir = _default_jax_cache_dir()
+if _jax_cache_dir is not None:
+    _os.makedirs(_jax_cache_dir, exist_ok=True)
+    _jax.config.update("jax_enable_compilation_cache", True)
+    _jax.config.update("jax_compilation_cache_dir", _jax_cache_dir)
+
+# Lazy public exports: name -> (module, attribute).  ``attribute=None``
+# exports the module itself.
+_LAZY_ATTRS: dict[str, tuple[str, str | None]] = {
+    # input
+    "VmecInput": (".core.input", "VmecInput"),
+    # solvers
+    "solve": (".core.solver", "solve"),
+    "solve_multigrid": (".core.multigrid", "solve_multigrid"),
+    "solve_free_boundary": (".core.freeboundary", "solve_free_boundary"),
+    # wout IO
+    "WoutData": (".core.wout", "WoutData"),
+    "read_wout": (".core.wout", "read_wout"),
+    "write_wout": (".core.wout", "write_wout"),
+    "wout_from_state": (".core.wout", "wout_from_state"),
+    # plotting + Boozer
+    "plot_wout": (".core.plotting", "plot_wout"),
+    "plot_boozmn": (".core.plotting", "plot_boozmn"),
+    "run_booz_xform": (".core.boozer", "run_booz_xform"),
+    # external fields
+    "MgridData": (".core.mgrid", "MgridData"),
+    "MgridField": (".core.mgrid", "MgridField"),
+    "read_mgrid": (".core.mgrid", "read_mgrid"),
+    "write_mgrid": (".core.mgrid", "write_mgrid"),
+    # errors
+    "VmecError": (".core.errors", "VmecError"),
+    "VmecInputError": (".core.errors", "VmecInputError"),
+    "VmecJacobianError": (".core.errors", "VmecJacobianError"),
+    "VmecConvergenceError": (".core.errors", "VmecConvergenceError"),
+    "MgridNotFoundError": (".core.errors", "MgridNotFoundError"),
+    # modules
+    "core": (".core", None),
+    "errors": (".core.errors", None),
+    "optimize": (".core.optimize", None),
+    "implicit": (".core.implicit", None),
+    "parallel": (".core.parallel", None),
+    "doctor": (".doctor", None),
+}
+
+__all__ = ["__version__", *sorted(_LAZY_ATTRS)]
+
+
+def __getattr__(name: str):
+    entry = _LAZY_ATTRS.get(name)
+    if entry is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attribute = entry
+    module = _import_module(module_name, __name__)
+    value = module if attribute is None else getattr(module, attribute)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(globals()) | set(_LAZY_ATTRS))

vmex/__main__.py

@@ -0,0 +1,14 @@
+"""Module entry point for `python -m vmex`."""
+
+import os as _os
+
+# Suppress noisy C++ warnings from XLA/PjRt before any JAX import.
+_os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
+_os.environ.setdefault("ABSL_MIN_LOG_LEVEL", "2")
+_os.environ.setdefault("GLOG_minloglevel", "2")
+
+from .core.cli import main
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

vmex/_compat.py

@@ -0,0 +1,310 @@
+"""JAX environment defaults + persistent compilation-cache policy.
+
+Historically this module was a full JAX/NumPy backend shim (``has_jax`` /
+``asarray`` / ``einsum`` / a no-op ``jit`` and a thread-local numpy mode).
+The core became JAX-only long ago and nothing imported that machinery any
+more, so it was deleted (Item I.8a).  What remains — and is
+actually used — is:
+
+- :func:`_configure_jax_environment` (run at import, i.e. before
+  ``vmex/__init__`` does ``import jax``): environment defaults that must
+  be set before JAX/XLA initializes — float64 (``JAX_ENABLE_X64``, VMEC
+  parity), synchronous CPU dispatch, quiet XLA/PjRt C++ logging, GPU
+  demand allocation, the machine-scoped persistent compilation-cache
+  directory, and the XLA:CPU fast-compile flags;
+- the compilation-cache policy helpers
+  :func:`_default_compilation_cache_dir` / :func:`_cache_machine_fingerprint`
+  / :func:`_configure_compilation_cache`, consumed by ``vmex/__init__``
+  and re-applied by ``core.solver._harden_compilation_cache`` on every solve
+  path (namespace-package shadowing guard).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+import hashlib
+from importlib import metadata as importlib_metadata
+import sys
+
+import os
+import platform
+
+
+def _env(name: str, default: str = "") -> str:
+    """Read ``VMEX_<name>``, falling back to the legacy ``VMEC_JAX_<name>``.
+
+    The package was renamed vmec_jax -> vmex; environment variables a user may
+    have set in their shell profile (the ``*_COMPILATION_CACHE*`` knobs in
+    particular) keep working under their old names for one release.
+    """
+    val = os.environ.get(f"VMEX_{name}")
+    if val is not None:
+        return val
+    return os.environ.get(f"VMEC_JAX_{name}", default)
+
+
+def _cache_machine_fingerprint() -> str:
+    """Return a short cache key for host-specific XLA CPU executables.
+
+    XLA CPU persistent-cache entries are native executables.  On shared home
+    directories, reusing an entry compiled on another CPU can trigger XLA AOT
+    loader errors or even illegal-instruction failures.  The fingerprint keeps
+    vmex's default cache portable by separating entries by OS, machine, and
+    CPU-feature/model signature.  Users who deliberately want a shared cache can
+    still set ``VMEX_COMPILATION_CACHE_DIR`` or ``JAX_COMPILATION_CACHE_DIR``.
+    """
+
+    parts = [
+        platform.system(),
+        platform.machine(),
+        platform.processor(),
+        f"python={sys.version_info.major}.{sys.version_info.minor}",
+    ]
+    for package in ("jax", "jaxlib"):
+        try:
+            parts.append(f"{package}={importlib_metadata.version(package)}")
+        except Exception:
+            pass
+    try:
+        if os.path.exists("/proc/cpuinfo"):
+            wanted = ("model name", "cpu family", "model", "stepping", "flags", "Features")
+            seen: set[str] = set()
+            with open("/proc/cpuinfo", encoding="utf-8", errors="ignore") as fh:
+                for line in fh:
+                    if ":" not in line:
+                        continue
+                    key, value = (part.strip() for part in line.split(":", 1))
+                    if key in wanted and key not in seen:
+                        parts.append(f"{key}={value}")
+                        seen.add(key)
+    except Exception:
+        pass
+    # macOS has no /proc/cpuinfo — capture the CPU brand + microarchitecture via
+    # sysctl so Intel/Apple-Silicon (and different chip generations) never share
+    # an XLA:CPU AOT cache entry.
+    if platform.system() == "Darwin":
+        try:
+            import subprocess
+            for key in ("machdep.cpu.brand_string", "hw.optional.arm.FEAT_SME",
+                        "hw.cpufamily"):
+                out = subprocess.run(["sysctl", "-n", key], capture_output=True,
+                                     text=True, timeout=2)
+                if out.returncode == 0 and out.stdout.strip():
+                    parts.append(f"{key}={out.stdout.strip()}")
+        except Exception:
+            pass
+    if not any(str(part).strip() for part in parts[:3]):
+        try:
+            parts.append(platform.node())
+        except Exception:
+            pass
+    digest = hashlib.sha256("|".join(parts).encode("utf-8", errors="ignore")).hexdigest()[:16]
+    system = platform.system().lower() or "unknown"
+    machine = platform.machine().lower() or "unknown"
+    return f"{system}-{machine}-{digest}"
+
+
+def _default_compilation_cache_dir() -> str | None:
+    """Return the configured JAX compilation-cache directory.
+
+    The persistent cache is enabled **by default on every backend** (CPU too)
+    so that repeated cold-process CLI/API runs reuse compiled kernels instead
+    of recompiling — a solovev CLI run drops 4.3 s -> 1.2 s on the second
+    invocation (R26c).  The XLA:CPU host-feature-mismatch hazard (AOT
+    executables tied to a specific instruction set, dangerous on shared home
+    filesystems) is handled by :func:`_cache_machine_fingerprint`, whose
+    per-machine suffix hashes the CPU model + feature flags (AVX2/AVX512/...),
+    so heterogeneous machines never share a cache entry.  Opt out with
+    ``VMEX_COMPILATION_CACHE=disabled`` (or ``VMEX_COMPILATION_CACHE_DIR=
+    disabled``); point it elsewhere with ``JAX_COMPILATION_CACHE_DIR=/path``.
+    """
+    # Already set by the user — respect it.
+    if "JAX_COMPILATION_CACHE_DIR" in os.environ:
+        val = os.environ["JAX_COMPILATION_CACHE_DIR"].strip()
+        if val.lower() in ("", "disabled", "0", "false", "no"):
+            return None
+        return val
+
+    # User can opt out via VMEX_COMPILATION_CACHE_DIR=disabled
+    vmec_val = _env("COMPILATION_CACHE_DIR").strip()
+    if vmec_val.lower() in ("disabled", "0", "false", "no"):
+        return None
+    if vmec_val:
+        return vmec_val
+
+    cache_flag = _env("COMPILATION_CACHE").strip().lower()
+    if cache_flag in ("disabled", "0", "false", "no", "off"):
+        return None
+
+    # Default cache location: ~/.cache/vmex/jax_cache/<machine-fingerprint>
+    # The host-specific suffix (CPU model + feature flags) prevents unsafe
+    # XLA:CPU AOT reuse on shared home filesystems where different machines see
+    # the same ~/.cache directory.
+    try:
+        import pathlib
+        return str(
+            pathlib.Path.home()
+            / ".cache"
+            / "vmex"
+            / "jax_cache"
+            / _cache_machine_fingerprint()
+        )
+    except Exception:
+        return None
+
+
+def _configure_compilation_cache(jax_module: Any, cache_dir: str | None) -> None:
+    """Apply vmex's persistent-cache defaults to an imported JAX module."""
+    if cache_dir is None:
+        return
+    try:
+        jax_module.config.update("jax_enable_compilation_cache", True)
+    except Exception:
+        pass
+    try:
+        jax_module.config.update("jax_compilation_cache_dir", cache_dir)
+    except Exception:
+        pass
+    try:
+        min_compile = _env("CACHE_MIN_COMPILE_TIME_SECS", "0")
+        jax_module.config.update("jax_persistent_cache_min_compile_time_secs", float(min_compile))
+    except Exception:
+        pass
+    try:
+        min_entry = _env("CACHE_MIN_ENTRY_SIZE_BYTES", "-1")
+        jax_module.config.update("jax_persistent_cache_min_entry_size_bytes", int(min_entry))
+    except Exception:
+        pass
+    try:
+        xla_caches = _env("PERSISTENT_CACHE_XLA_CACHES").strip()
+        if not xla_caches:
+            platform_name = os.environ.get("JAX_PLATFORM_NAME", "").strip().lower()
+            platforms = os.environ.get("JAX_PLATFORMS", "").strip().lower()
+            visible_cuda = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip().lower()
+            gpu_requested = (
+                platform_name in ("gpu", "cuda")
+                or any(part.strip() in ("gpu", "cuda") for part in platforms.split(","))
+                or visible_cuda not in ("", "-1", "none", "no")
+            )
+            xla_caches = "xla_gpu_per_fusion_autotune_cache_dir" if gpu_requested else "none"
+        if xla_caches.lower() not in ("", "none", "0", "false", "no", "off"):
+            jax_module.config.update("jax_persistent_cache_enable_xla_caches", xla_caches)
+    except Exception:
+        pass
+    try:
+        max_size = _env("COMPILATION_CACHE_MAX_SIZE")
+        if max_size:
+            jax_module.config.update("jax_compilation_cache_max_size", int(max_size))
+    except Exception:
+        pass
+    try:
+        explain = _env("EXPLAIN_CACHE_MISSES")
+        if explain.strip().lower() not in ("", "0", "false", "no"):
+            jax_module.config.update("jax_explain_cache_misses", True)
+    except Exception:
+        pass
+
+
+def _configure_jax_environment() -> None:
+    """Set JAX/XLA environment defaults, then import + configure JAX.
+
+    Runs once at ``vmex._compat`` import time — before
+    ``vmex/__init__`` (or anything else in the package) imports JAX — so
+    the env-var defaults reliably reach XLA backend initialization.  Every
+    default uses ``setdefault``: an explicit user environment always wins.
+    """
+    try:
+        # Enable x64 by default for VMEC parity unless the user opted out.
+        os.environ.setdefault("JAX_ENABLE_X64", "1")
+        # VMEC/JAX optimization callbacks immediately materialize most results
+        # on the host (SciPy residuals/Jacobians, history, wout writing).  On
+        # CPU, asynchronous dispatch can leave completed XLA/PjRt work and
+        # executable state queued across many exact-Jacobian callbacks in one
+        # long-lived process.  Default CPU dispatch to synchronous execution so
+        # memory is reclaimed at callback boundaries; users can still override
+        # this before import with JAX_CPU_ENABLE_ASYNC_DISPATCH=true.
+        os.environ.setdefault("JAX_CPU_ENABLE_ASYNC_DISPATCH", "false")
+        # Suppress noisy C++ warnings from XLA/PjRt backend (e.g.
+        # repeated "Assume version compatibility. PjRt-IFRT does not
+        # track XLA executable versions." on persistent-cache hits).
+        # These are harmless informational messages emitted by the XLA
+        # runtime logging stack. Level 0=INFO, 1=WARNING, 2=ERROR — we
+        # default to ERROR-only so that genuine errors still surface.
+        os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
+        os.environ.setdefault("ABSL_MIN_LOG_LEVEL", "2")
+        os.environ.setdefault("GLOG_minloglevel", "2")
+        # JAX's default GPU allocator preallocates most device memory.  That
+        # hurts vmex's exact-optimizer workload in practice: it prevents
+        # concurrent profiling/worker processes from starting and can make the
+        # accepted-point replay path much slower.  Default to demand allocation
+        # unless the user already set JAX's allocator env var or explicitly
+        # asks vmex to keep JAX's preallocation default.
+        _vmec_gpu_prealloc = _env("GPU_PREALLOCATE").strip().lower()
+        if (
+            "XLA_PYTHON_CLIENT_PREALLOCATE" not in os.environ
+            and _vmec_gpu_prealloc not in ("1", "true", "yes", "on")
+        ):
+            os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+
+        # Enable the JAX disk compilation cache in a machine-scoped directory.
+        # This avoids unsafe XLA:CPU AOT reuse across hosts while preserving
+        # repeated cold-process speedups on the same machine.
+        _cache_dir = _default_compilation_cache_dir()
+        if _cache_dir is not None:
+            os.environ.setdefault("JAX_COMPILATION_CACHE_DIR", _cache_dir)
+
+        # XLA:CPU compile-time flags.  The differentiable/optimization pipeline
+        # is COMPILE-dominated (the fused adjoint VJP + GMRES graph is ~21 s of a
+        # ~24 s cold ``value_and_grad``); XLA's default backend optimization
+        # level (3) spends most of that in expensive LLVM passes.  Level 1 plus
+        # disabling the expensive passes typically cuts compile wall-time
+        # ~1.3-2x, at the cost of slightly slower *warm* kernels -- a good trade
+        # for this compile-bound workload.  Applied on CPU only (LLVM codegen),
+        # never with fast-math (that would break float64 parity/determinism),
+        # skipped if the user set XLA_FLAGS, and opt-out via
+        # VMEX_FAST_COMPILE=0 (e.g. a very long single-process opt loop that
+        # amortizes compile over many warm calls prefers level 3).
+        _fast_compile = _env("FAST_COMPILE", "1").strip().lower()
+        _accel_req = os.environ.get("JAX_PLATFORM_NAME", "").strip().lower()
+        _accel_reqs = os.environ.get("JAX_PLATFORMS", "").strip().lower()
+        _cuda_vis = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
+        _on_accel = (
+            any(a in f"{_accel_req} {_accel_reqs}" for a in ("cuda", "gpu", "tpu", "rocm"))
+            or (_cuda_vis not in ("", "-1"))
+        )
+        if (
+            _fast_compile not in ("0", "false", "no", "off")
+            and "XLA_FLAGS" not in os.environ
+            and not _on_accel
+        ):
+            os.environ["XLA_FLAGS"] = (
+                "--xla_backend_optimization_level=1 "
+                "--xla_llvm_disable_expensive_passes=true"
+            )
+
+        import jax
+
+        try:
+            jax.config.update("jax_enable_x64", os.environ.get("JAX_ENABLE_X64", "0") == "1")
+        except Exception:
+            pass
+        try:
+            _cpu_async = os.environ.get("JAX_CPU_ENABLE_ASYNC_DISPATCH", "true")
+            jax.config.update(
+                "jax_cpu_enable_async_dispatch",
+                _cpu_async.strip().lower() not in ("0", "false", "no", "off"),
+            )
+        except Exception:
+            pass
+
+        # Wire up the compilation cache via jax.config too; the env-var path
+        # alone does not cover all JAX/JAXLIB versions and cache thresholds.
+        _configure_compilation_cache(jax, _cache_dir)
+    except Exception:
+        # Never block a vmex import over environment tuning (e.g. docs
+        # builds with a mocked JAX): core.solver enforces the hard
+        # requirements (x64, cache hardening) on every solve path anyway.
+        pass
+
+
+_configure_jax_environment()

vmex/core/__init__.py

@@ -0,0 +1,29 @@
+"""Clean-room core of vmex (§5) — will replace the legacy modules.
+
+Module map (each header docstring names its VMEC2000 counterpart):
+
+- ``errors``          typed zero-crash exceptions + werror table
+- ``printing``        VMEC2000-format console output (byte-exact)
+- ``input``           VmecInput: INDATA + VMEC++-JSON parsing, round-trip writers
+- ``profiles``        pressure/iota/current parameterizations (pure jnp)
+- ``fourier``         Resolution, ModeTable, trig tables (fixaray.f)
+- ``transforms``      totzsps/totzspa/tomnsps/tomnspa as batched matmuls
+- ``geometry``        real-space R/Z/lambda, half-mesh jacobian (jacobian.f)
+- ``fields``          metrics, B components, energies, tcon (bcovar.f)
+- ``forces``          MHD force kernels + spectral condensation (forces.f, alias.f)
+- ``residuals``       m=1 constraint, fsqr/fsqz/fsql, preconditioned lane (residue.f90)
+- ``preconditioner``  1D radial tridiagonal preconditioner (precondn.f, scalfor.f)
+- ``step``            Richardson stepping + restart control (evolve.f, restart.f)
+- ``setup``           radial profiles + initial guess (profil1d/3d.f, readin.f)
+- ``solver``          single-grid fixed-boundary solve loop (funct3d.f, eqsolve.f)
+- ``statephysics``    shared state-physics primitives (_field_chain, half-mesh iota/sampling)
+- ``implicit``        implicit differentiation of the equilibrium (custom VJP + adjoint GMRES)
+- ``stability``       differentiable ideal-MHD stability (infinite-n ballooning; COBRA port)
+- ``freeboundary_diff`` differentiable free-boundary residual via virtual casing (R15.3/R19)
+- ``device``          CPU/GPU placement policy (measured: benchmarks/gpu_baseline.json)
+
+Every module is validated by A/B equivalence tests against the legacy
+parity-proven implementation in ``tests/``; the solve loop is
+validated end-to-end against VMEC2000 golden runs
+(``tests/test_solver_end_to_end.py``).
+"""