dkx 2.3.0__tar.gz → 2.4.0__tar.gz

dkx-2.4.0/PKG-INFO

@@ -0,0 +1,181 @@
+Metadata-Version: 2.4
+Name: dkx
+Version: 2.4.0
+Summary: Neoclassical transport solver with CPU/GPU and differentiable JAX workflows
+Project-URL: Homepage, https://github.com/uwplasma/DKX
+Project-URL: Documentation, https://sfincs-jax.readthedocs.io/
+Project-URL: Repository, https://github.com/uwplasma/DKX
+Project-URL: Issues, https://github.com/uwplasma/DKX/issues
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: numpy>=2.1
+Requires-Dist: scipy
+Requires-Dist: h5py
+Requires-Dist: netCDF4
+Requires-Dist: matplotlib
+Requires-Dist: rich>=13.0
+Requires-Dist: solvax>=0.21.0
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-xdist; extra == "dev"
+Requires-Dist: pytest-split; extra == "dev"
+Requires-Dist: pypdf; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: myst-parser; extra == "docs"
+Requires-Dist: furo; extra == "docs"
+Dynamic: license-file
+
+# DKX
+
+[![PyPI](https://img.shields.io/pypi/v/dkx)](https://pypi.org/project/dkx/)
+[![CI](https://img.shields.io/github/actions/workflow/status/uwplasma/DKX/ci.yml?branch=main&label=ci)](https://github.com/uwplasma/DKX/actions/workflows/ci.yml)
+[![Docs](https://img.shields.io/readthedocs/sfincs-jax?label=docs)](https://sfincs-jax.readthedocs.io/en/latest/)
+[![License](https://img.shields.io/github/license/uwplasma/DKX)](LICENSE)
+
+**Differentiable neoclassical transport for stellarators and tokamaks, in JAX.**
+
+DKX solves the radially local, linearized drift-kinetic equation on a flux surface
+and returns particle and heat fluxes, parallel flows, bootstrap current, transport
+matrices and ambipolar radial electric fields. It implements the SFINCS Fortran v3
+model, including full Fokker–Planck collisions and Phi1, reads and writes SFINCS decks,
+runs on CPU or GPU, and every output is differentiable in every input.
+
+![W7-X standard configuration: |B| and parallel current density on the boundary, bootstrap current profile, ambipolar Er against Pablant et al. 2018](docs/_static/figures/readme/w7x_showcase.png)
+
+## Install
+
+```bash
+pip install dkx                    # CPU
+pip install -U "jax[cuda12]"       # add for NVIDIA GPUs
+```
+
+## Run
+
+```python
+import dkx
+
+case = dkx.Case.from_mapping({           # analytic tokamak, teaching grid: seconds, not converged
+    "schema": 1, "name": "tokamak", "run": {"workflow": "profile", "progress": False},
+    "geometry": {"format": "analytic", "file": "tokamak", "surfaces": [0.16, 0.25, 0.36]},
+    "species": [{"name": "deuterium", "charge": 1, "mass_amu": 2.014,
+                 "density_m3": [8.0e19, 7.0e19, 6.0e19], "temperature_keV": [1.0, 0.8, 0.6]}],
+    "physics": {"model": "full_local", "collisions": "pitch_angle_scattering", "magnetic_drifts": "dkes", "phi1": "off"},
+    "electric_field": {"mode": "prescribed", "value_kV_m": 0.0},
+    "resolution": {"theta": 9, "zeta": 1, "pitch": 8, "speed": 4}, "solver": {"method": "auto", "relative_tolerance": 1e-8},
+})
+result = dkx.run(case)
+print("solver route:", result.metadata["solver_route"])
+print("particle flux:", float(result.arrays["particle_flux_m2_s"][1, 0]))
+```
+
+From the shell: `dkx run case.toml`, `dkx converge case.toml` (check the resolution before
+trusting a number), `dkx wout_w7x.nc`, or `dkx input.namelist` for a SFINCS deck.
+
+## Gradients
+
+```python
+import jax, jax.numpy as jnp, dkx
+
+case = dkx.Case.from_file("examples/05_ambipolar_profile/case.toml")
+problem = dkx.prepare_er_scan(case, surface_index=1)          # geometry, grids, collisions once
+
+def bootstrap_current(er_kv_m):
+    scan = dkx.batched_er_scan(problem, er_kv_m, differentiable=True, retain_full_state=True)
+    return jnp.sum(scan.moments["FSABjHat"])
+
+j, dj_der = jax.jit(jax.value_and_grad(bootstrap_current))(jnp.array([-0.2, 0.0, 0.2]))
+```
+
+The derivative passes through the linear solve by the implicit function theorem, and
+every returned state has satisfied the original kinetic equation. Profiles and geometry
+differentiate the same way: [differentiability](docs/differentiability.rst).
+
+## Choose a workflow
+
+| Calculation | Start from |
+| --- | --- |
+| Tokamak profile with prescribed `E_r` | `examples/01_tokamak_profile` |
+| Stellarator from VMEC or Boozer files | `examples/02_vmec_stellarator`, `examples/03_boozer_stellarator` |
+| Monoenergetic transport scan | `examples/04_monoenergetic_scan` |
+| Ambipolar roots and branch selection | `examples/05_ambipolar_profile` |
+| Resolution study, gradient checks | `examples/06_convergence_certificate`, `examples/07_gradients` |
+| Geometry sensitivity and descent (analytic proxy) | `examples/08_vmex_optimization` |
+| Phi1 and impurities (expert path) | `examples/09_phi1_and_impurities` |
+
+## Why DKX
+
+| | DKX | SFINCS v3 | MONKES | yancc |
+| --- | :---: | :---: | :---: | :---: |
+| Full linearized Fokker–Planck, multispecies | ✅ | ✅ | ❌ | ✅ |
+| Analytic, VMEC, Boozer and `lasym` geometry | ✅ | ✅ | ✅ | ✅ |
+| Phi1 quasineutrality; Tangential magnetic drifts; `export_f` | ✅ | ✅ | ❌ | ❌ |
+| Ambipolar `E_r` root with retained branch evidence | ✅ | ✅ | ❌ | ❌ |
+| Transport matrices (RHSMode 2/3) and SFINCS deck/HDF5 I/O | ✅ | ✅ | ❌ | ❌ |
+| GPU, JIT-compiled scans, Krylov recycling | ✅ | ❌ | ❌ | ✅ |
+| Exact gradients of any output w.r.t. any input | ✅ | adjoint branches | ❌ | claimed |
+| Gradients verified against finite differences | ✅ | | | |
+
+## Verified
+
+![DKX against SFINCS Fortran v3, MONKES and YANCC: scaled differences on matched full Fokker-Planck decks, and Beidler-normalized monoenergetic coefficients](docs/_static/figures/readme/cross_code_validation.png)
+
+Against SFINCS Fortran v3 on 38 upstream decks with the same discretization, DKX
+agrees to solver tolerance: median 4e-6, full Fokker–Planck decks to 1e-8. Against the
+independent codes MONKES and YANCC, the four Beidler-normalized monoenergetic coefficients
+agree within 6 percent and `D33` within 0.1 percent on three configurations. Gradients
+agree with central finite differences over a step window on every shipped derivative
+example. Details, tolerances and scope: [validation matrix](docs/validation_matrix.rst).
+
+## Fast
+
+![Runtime and peak memory, DKX against SFINCS Fortran v3, on the 744k-unknown HSX PAS case](docs/_static/figures/readme/tier1_hsx_runtime_memory.png)
+
+`HSX_PASCollisions_DKESTrajectories`, RHSMode=1, **744,610 unknowns**, one machine, against the
+PETSc 3.23 / MUMPS 5.8.2 build of SFINCS v3. Warm is the second solve in a process, after XLA has compiled.
+
+| Configuration | Warm solve | Peak RSS |
+| --- | ---: | ---: |
+| DKX, `Nxi`-for-`x` ramp | **27.2 s** | **0.93 GB** |
+| DKX, uniform `Nxi` | 44.3 s | 1.16 GB |
+| DKX, RTX A4000 GPU | 45.0 s | — |
+| SFINCS Fortran v3, 1 rank | 463.6 s | 3.98 GB |
+| SFINCS Fortran v3, 2 ranks (its best) | 229.5 s | 2.86 GB |
+
+| Cold versus warm, M3 Max CPU | Unknowns | Cold | Warm |
+| --- | ---: | ---: | ---: |
+| HSX PAS reduced | 40,584 | 1.72 s | 0.12 s |
+| HSX PAS, `25x51x100x5` | 744,610 | 23.6 s | 20.0 s |
+
+That is **one measured 744k-unknown HSX PAS case**, chosen because DKX has an exact structured
+solver for it. Across all 38 upstream decks: structured route faster on 9 of 9; Krylov route faster
+on 7 of 23, six not completed. Every deck, hardware string and method: [performance](docs/performance.rst).
+
+![Measured parity envelopes of DKX against SFINCS Fortran v3](docs/_static/figures/readme/canonical_parity.png)
+
+## Documentation
+
+[Tutorial: first W7-X result](docs/usage.rst) ·
+[How-to: resolution, wout files, SFINCS decks](docs/case_files.rst) ·
+[Reference: schema, CLI, API](docs/api.rst) ·
+[Explanation: physics models, solver routes, limitations](docs/physics_models.rst)
+
+## Cite
+
+```bibtex
+@software{dkx,
+  author = {Jorge, Rogerio and contributors},
+  title  = {DKX: differentiable drift-kinetic neoclassical transport in JAX},
+  url    = {https://github.com/uwplasma/DKX},
+  year   = {2026}
+}
+```
+
+Please also cite SFINCS (Landreman, Smith, Mollén & Helander, Phys. Plasmas 21, 042503,
+2014) when you use its decks or model. Metadata: [CITATION.cff](CITATION.cff).
+
+[Examples](examples/) · [Research plan](plan.md) · [Contributing](docs/contributing.rst) ·
+[SOLVAX](https://github.com/uwplasma/SOLVAX) owns the reusable solver algorithms · [LICENSE](LICENSE)

dkx-2.4.0/README.md

@@ -0,0 +1,150 @@
+# DKX
+
+[![PyPI](https://img.shields.io/pypi/v/dkx)](https://pypi.org/project/dkx/)
+[![CI](https://img.shields.io/github/actions/workflow/status/uwplasma/DKX/ci.yml?branch=main&label=ci)](https://github.com/uwplasma/DKX/actions/workflows/ci.yml)
+[![Docs](https://img.shields.io/readthedocs/sfincs-jax?label=docs)](https://sfincs-jax.readthedocs.io/en/latest/)
+[![License](https://img.shields.io/github/license/uwplasma/DKX)](LICENSE)
+
+**Differentiable neoclassical transport for stellarators and tokamaks, in JAX.**
+
+DKX solves the radially local, linearized drift-kinetic equation on a flux surface
+and returns particle and heat fluxes, parallel flows, bootstrap current, transport
+matrices and ambipolar radial electric fields. It implements the SFINCS Fortran v3
+model, including full Fokker–Planck collisions and Phi1, reads and writes SFINCS decks,
+runs on CPU or GPU, and every output is differentiable in every input.
+
+![W7-X standard configuration: |B| and parallel current density on the boundary, bootstrap current profile, ambipolar Er against Pablant et al. 2018](docs/_static/figures/readme/w7x_showcase.png)
+
+## Install
+
+```bash
+pip install dkx                    # CPU
+pip install -U "jax[cuda12]"       # add for NVIDIA GPUs
+```
+
+## Run
+
+```python
+import dkx
+
+case = dkx.Case.from_mapping({           # analytic tokamak, teaching grid: seconds, not converged
+    "schema": 1, "name": "tokamak", "run": {"workflow": "profile", "progress": False},
+    "geometry": {"format": "analytic", "file": "tokamak", "surfaces": [0.16, 0.25, 0.36]},
+    "species": [{"name": "deuterium", "charge": 1, "mass_amu": 2.014,
+                 "density_m3": [8.0e19, 7.0e19, 6.0e19], "temperature_keV": [1.0, 0.8, 0.6]}],
+    "physics": {"model": "full_local", "collisions": "pitch_angle_scattering", "magnetic_drifts": "dkes", "phi1": "off"},
+    "electric_field": {"mode": "prescribed", "value_kV_m": 0.0},
+    "resolution": {"theta": 9, "zeta": 1, "pitch": 8, "speed": 4}, "solver": {"method": "auto", "relative_tolerance": 1e-8},
+})
+result = dkx.run(case)
+print("solver route:", result.metadata["solver_route"])
+print("particle flux:", float(result.arrays["particle_flux_m2_s"][1, 0]))
+```
+
+From the shell: `dkx run case.toml`, `dkx converge case.toml` (check the resolution before
+trusting a number), `dkx wout_w7x.nc`, or `dkx input.namelist` for a SFINCS deck.
+
+## Gradients
+
+```python
+import jax, jax.numpy as jnp, dkx
+
+case = dkx.Case.from_file("examples/05_ambipolar_profile/case.toml")
+problem = dkx.prepare_er_scan(case, surface_index=1)          # geometry, grids, collisions once
+
+def bootstrap_current(er_kv_m):
+    scan = dkx.batched_er_scan(problem, er_kv_m, differentiable=True, retain_full_state=True)
+    return jnp.sum(scan.moments["FSABjHat"])
+
+j, dj_der = jax.jit(jax.value_and_grad(bootstrap_current))(jnp.array([-0.2, 0.0, 0.2]))
+```
+
+The derivative passes through the linear solve by the implicit function theorem, and
+every returned state has satisfied the original kinetic equation. Profiles and geometry
+differentiate the same way: [differentiability](docs/differentiability.rst).
+
+## Choose a workflow
+
+| Calculation | Start from |
+| --- | --- |
+| Tokamak profile with prescribed `E_r` | `examples/01_tokamak_profile` |
+| Stellarator from VMEC or Boozer files | `examples/02_vmec_stellarator`, `examples/03_boozer_stellarator` |
+| Monoenergetic transport scan | `examples/04_monoenergetic_scan` |
+| Ambipolar roots and branch selection | `examples/05_ambipolar_profile` |
+| Resolution study, gradient checks | `examples/06_convergence_certificate`, `examples/07_gradients` |
+| Geometry sensitivity and descent (analytic proxy) | `examples/08_vmex_optimization` |
+| Phi1 and impurities (expert path) | `examples/09_phi1_and_impurities` |
+
+## Why DKX
+
+| | DKX | SFINCS v3 | MONKES | yancc |
+| --- | :---: | :---: | :---: | :---: |
+| Full linearized Fokker–Planck, multispecies | ✅ | ✅ | ❌ | ✅ |
+| Analytic, VMEC, Boozer and `lasym` geometry | ✅ | ✅ | ✅ | ✅ |
+| Phi1 quasineutrality; Tangential magnetic drifts; `export_f` | ✅ | ✅ | ❌ | ❌ |
+| Ambipolar `E_r` root with retained branch evidence | ✅ | ✅ | ❌ | ❌ |
+| Transport matrices (RHSMode 2/3) and SFINCS deck/HDF5 I/O | ✅ | ✅ | ❌ | ❌ |
+| GPU, JIT-compiled scans, Krylov recycling | ✅ | ❌ | ❌ | ✅ |
+| Exact gradients of any output w.r.t. any input | ✅ | adjoint branches | ❌ | claimed |
+| Gradients verified against finite differences | ✅ | | | |
+
+## Verified
+
+![DKX against SFINCS Fortran v3, MONKES and YANCC: scaled differences on matched full Fokker-Planck decks, and Beidler-normalized monoenergetic coefficients](docs/_static/figures/readme/cross_code_validation.png)
+
+Against SFINCS Fortran v3 on 38 upstream decks with the same discretization, DKX
+agrees to solver tolerance: median 4e-6, full Fokker–Planck decks to 1e-8. Against the
+independent codes MONKES and YANCC, the four Beidler-normalized monoenergetic coefficients
+agree within 6 percent and `D33` within 0.1 percent on three configurations. Gradients
+agree with central finite differences over a step window on every shipped derivative
+example. Details, tolerances and scope: [validation matrix](docs/validation_matrix.rst).
+
+## Fast
+
+![Runtime and peak memory, DKX against SFINCS Fortran v3, on the 744k-unknown HSX PAS case](docs/_static/figures/readme/tier1_hsx_runtime_memory.png)
+
+`HSX_PASCollisions_DKESTrajectories`, RHSMode=1, **744,610 unknowns**, one machine, against the
+PETSc 3.23 / MUMPS 5.8.2 build of SFINCS v3. Warm is the second solve in a process, after XLA has compiled.
+
+| Configuration | Warm solve | Peak RSS |
+| --- | ---: | ---: |
+| DKX, `Nxi`-for-`x` ramp | **27.2 s** | **0.93 GB** |
+| DKX, uniform `Nxi` | 44.3 s | 1.16 GB |
+| DKX, RTX A4000 GPU | 45.0 s | — |
+| SFINCS Fortran v3, 1 rank | 463.6 s | 3.98 GB |
+| SFINCS Fortran v3, 2 ranks (its best) | 229.5 s | 2.86 GB |
+
+| Cold versus warm, M3 Max CPU | Unknowns | Cold | Warm |
+| --- | ---: | ---: | ---: |
+| HSX PAS reduced | 40,584 | 1.72 s | 0.12 s |
+| HSX PAS, `25x51x100x5` | 744,610 | 23.6 s | 20.0 s |
+
+That is **one measured 744k-unknown HSX PAS case**, chosen because DKX has an exact structured
+solver for it. Across all 38 upstream decks: structured route faster on 9 of 9; Krylov route faster
+on 7 of 23, six not completed. Every deck, hardware string and method: [performance](docs/performance.rst).
+
+![Measured parity envelopes of DKX against SFINCS Fortran v3](docs/_static/figures/readme/canonical_parity.png)
+
+## Documentation
+
+[Tutorial: first W7-X result](docs/usage.rst) ·
+[How-to: resolution, wout files, SFINCS decks](docs/case_files.rst) ·
+[Reference: schema, CLI, API](docs/api.rst) ·
+[Explanation: physics models, solver routes, limitations](docs/physics_models.rst)
+
+## Cite
+
+```bibtex
+@software{dkx,
+  author = {Jorge, Rogerio and contributors},
+  title  = {DKX: differentiable drift-kinetic neoclassical transport in JAX},
+  url    = {https://github.com/uwplasma/DKX},
+  year   = {2026}
+}
+```
+
+Please also cite SFINCS (Landreman, Smith, Mollén & Helander, Phys. Plasmas 21, 042503,
+2014) when you use its decks or model. Metadata: [CITATION.cff](CITATION.cff).
+
+[Examples](examples/) · [Research plan](plan.md) · [Contributing](docs/contributing.rst) ·
+[SOLVAX](https://github.com/uwplasma/SOLVAX) owns the reusable solver algorithms · [LICENSE](LICENSE)

{dkx-2.3.0 → dkx-2.4.0}/pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dkx"
-version = "2.3.0"
+dynamic = ["version"]
 description = "Neoclassical transport solver with CPU/GPU and differentiable JAX workflows"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
   # numpy>=2.1 is not a new constraint -- it is the one jax already carries,
   # made explicit so the resolver says so.  Left implicit, `pip install dkx`
@@ -22,20 +22,25 @@ dependencies = [
   "h5py",
   "netCDF4",
   "matplotlib",
-  # solvax (PyPI) owns the structured linear-algebra solver tiers (block-Thomas
+  # Rich owns every user-facing table, progress bar, and error panel the CLI
+  # prints. It is a core dependency, not an extra: `dkx run` without it would
+  # have to carry a second, plainer output path, and two formatters drift.
+  "rich>=13.0",
+  # solvax (PyPI) owns the structured linear-algebra solver routes (block-Thomas
   # elimination, recycled GCROT, implicit differentiation) that every canonical
-  # solve uses; it is a core dependency, not optional.  0.13.0 is the floor:
-  # the tier-2 coarse preconditioner's reusable Schur-LU-only route needs
-  # block_thomas_factor_fn(..., store_offdiagonals=False), which is new there.
-  "solvax>=0.13.0",
+  # solve uses; it is a core dependency, not optional. 0.21.0 is the floor:
+  # its bordered preconditioner applies the coarse inverse once per Krylov
+  # iteration instead of twice (SOLVAX #101), and every published DKX Krylov
+  # timing since 2026-09-12 was measured with that behaviour. 0.19.0 added the
+  # selected-tail generated block solve and exact selected-head reuse used by the
+  # opt-in bounded Legendre-tail evidence.
+  "solvax>=0.21.0",
 ]
 
 [project.optional-dependencies]
-# Retained as a no-op alias so existing `pip install dkx[structured]`
-# invocations keep working; solvax is now a core dependency (see above).
-structured = [
-  "solvax>=0.13.0",
-]
+# `structured` is gone. It was an alias for a dependency that has been core
+# since solvax moved into the required list, so installing it changed nothing;
+# pip warns about the unknown extra and installs the same environment.
 dev = [
   "pytest",
   "pytest-cov",
@@ -47,7 +52,8 @@ dev = [
 ]
 docs = [
   "sphinx",
-  "sphinx-rtd-theme",
+  "myst-parser",
+  "furo",
 ]
 
 [project.urls]
@@ -60,16 +66,32 @@ Issues = "https://github.com/uwplasma/DKX/issues"
 dkx = "dkx.cli:main"
 
 [tool.setuptools.packages.find]
-where = ["."]
+where = ["src"]
 include = ["dkx*"]
 
+[tool.setuptools.package-dir]
+"" = "src"
+
 [tool.setuptools.package-data]
 dkx = ["validation/*.json"]
 
+[tool.setuptools.dynamic]
+version = {attr = "dkx._version.__version__"}
+
+# The coverage shards install the package non-editable, so they record paths
+# under site-packages, while the report job only checks the repository out.
+# Without this mapping `coverage combine` keeps the site-packages paths and the
+# report fails with "No source for code". Listing both tells coverage they are
+# the same code and rewrites the shard data onto the checkout.
+[tool.coverage.paths]
+source = [
+  "src/dkx",
+  "*/site-packages/dkx",
+]
+
 [tool.pytest.ini_options]
 addopts = ["--import-mode=prepend"]
 testpaths = ["tests"]
 markers = [
   "slow: heavy end-to-end example/optimization tests; excluded from the fast coverage shards via -m 'not slow' (still run in the full local suite).",
 ]
-

dkx-2.4.0/src/dkx/__init__.py

@@ -0,0 +1,184 @@
+"""Differentiable neoclassical transport solvers and SFINCS-style outputs in JAX.
+
+The public CLI and Python APIs are maintained as standalone research tools while
+retaining release-gated comparisons against SFINCS Fortran v3 for trust building.
+
+Importing this package is inert: it reads no environment, writes none, creates
+no directory, imports no JAX backend, and chooses no thread count. The runtime
+those things configure lives in :mod:`dkx.runtime` and is applied by
+:func:`dkx.runtime.configure`, which the CLI and every solve entry point call
+for you. See plan.md section 6.4.
+"""
+
+from __future__ import annotations
+
+from ._version import __version__
+from . import runtime
+from .runtime import configure, initialize_distributed_runtime_from_env
+
+from .api import (  # noqa: E402
+    BenchmarkReport,
+    GeometryState,
+    GridState,
+    OperatorState,
+    OutputSchema,
+    PreconditionerState,
+    SolveInputs,
+    SolverOptions,
+    SolverResult,
+    TransportResult,
+    prepare_er_scan,
+    batched_er_scan,
+    batched_surface_scan,
+    read_output,
+    run_ambipolar_brent,
+    run_monoenergetic_database,
+    write_output,
+)
+from .inputs import SfincsInput, load_sfincs_input  # noqa: E402
+from .config import (  # noqa: E402
+    Case,
+    CaseValidationError,
+    ConvergenceConfig,
+    ElectricFieldConfig,
+    GeometryConfig,
+    OutputConfig,
+    ParallelConfig,
+    PhysicsConfig,
+    ResolutionConfig,
+    RunConfig,
+    ScanAxis,
+    ScanConfig,
+    SolverConfig,
+    SpeciesConfig,
+    case_json_schema,
+)
+from .result import RESULT_SCHEMA_VERSION, Result  # noqa: E402
+
+# Heavy flagship entry points (they import the JAX solve stack) are exported
+# lazily via PEP 562 module __getattr__ so `import dkx` stays cheap.
+_LAZY_EXPORTS = {
+    "plot": ("dkx.plotting", "plot"),
+    "run_profile": ("dkx.run", "run_profile"),
+    "run_transport_matrix": ("dkx.run", "run_transport_matrix"),
+    "run_from_namelist": ("dkx.run", "run_from_namelist"),
+    "batched_solve": ("dkx.batch", "batched_solve"),
+    "monoenergetic_database": ("dkx.monoenergetic", "monoenergetic_database"),
+    "ambipolar_er": ("dkx.er", "ambipolar_er"),
+    "find_ambipolar_er": ("dkx.er", "find_ambipolar_er"),
+    "classical_impurity_flux": ("dkx.impurity", "classical_impurity_flux"),
+    "build_impurity_plasma": ("dkx.impurity", "build_impurity_plasma"),
+}
+
+
+def __getattr__(name: str):
+    if name == "run":
+        return _lazy_run_module()
+    try:
+        module_name, attr = _LAZY_EXPORTS[name]
+    except KeyError:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
+    import importlib  # noqa: PLC0415
+
+    value = getattr(importlib.import_module(module_name), attr)
+    globals()[name] = value  # cache: subsequent lookups skip __getattr__
+    return value
+
+
+def __dir__() -> list[str]:
+    # `annotations` is the __future__ import, not an export. Everything else
+    # public is either in __all__ or a submodule the user reached for.
+    return sorted((set(globals()) | set(_LAZY_EXPORTS)) - {"annotations"})
+
+
+# ``dkx/run.py`` is a module and ``dkx.run(case)`` is a call.  Rather than pick
+# one -- the function breaks ``dkx.run.run_profile`` and every monkeypatch that
+# targets it by path, the module breaks the call -- ``dkx/run.py`` makes itself
+# callable, so ``dkx.run`` is always the module and always invocable.  The name
+# is resolved here rather than listed in _LAZY_EXPORTS because what it yields
+# is the module itself, not an attribute of one.
+def _lazy_run_module():
+    import importlib  # noqa: PLC0415
+
+    module = importlib.import_module(f"{__name__}.run")
+    globals()["run"] = module
+    return module
+
+
+
+def require_float64() -> None:
+    """Raise unless JAX is in float64 mode.
+
+    :func:`dkx.runtime.configure` enables it, and every solve entry point calls
+    that; a caller who set ``DKX_NO_X64_SETUP`` or passed ``jax_x64=False`` has
+    taken the job on, and this is where they find out if they dropped it.  The
+    check is a dtype probe rather than a config read because the config can be
+    set and then overridden, and what matters is the dtype arrays actually get.
+    """
+    import jax.numpy as _jnp  # noqa: PLC0415
+
+    if _jnp.zeros(1).dtype != _jnp.float64:
+        raise RuntimeError(
+            "dkx requires JAX float64: the block eliminations and every parity "
+            "fixture depend on it, and single precision changes which results "
+            "are trustworthy rather than merely how accurate they are. "
+            "Enable it with dkx.runtime.configure() before the first array is "
+            "created, or jax.config.update('jax_enable_x64', True), or "
+            "JAX_ENABLE_X64=1 in the environment, or unset DKX_NO_X64_SETUP "
+            "and let dkx set it."
+        )
+
+__all__ = [
+    "configure",
+    "runtime",
+    "require_float64",
+    "BenchmarkReport",
+    "Case",
+    "CaseValidationError",
+    "ConvergenceConfig",
+    "ElectricFieldConfig",
+    "GeometryConfig",
+    "GeometryState",
+    "GridState",
+    "OperatorState",
+    "OutputSchema",
+    "OutputConfig",
+    "ParallelConfig",
+    "PhysicsConfig",
+    "PreconditionerState",
+    "RESULT_SCHEMA_VERSION",
+    "ResolutionConfig",
+    "Result",
+    "RunConfig",
+    "SfincsInput",
+    "SolveInputs",
+    "SolverOptions",
+    "SolverResult",
+    "ScanAxis",
+    "ScanConfig",
+    "SolverConfig",
+    "SpeciesConfig",
+    "TransportResult",
+    "__version__",
+    "ambipolar_er",
+    "prepare_er_scan",
+    "batched_er_scan",
+    "batched_surface_scan",
+    "case_json_schema",
+    "batched_solve",
+    "build_impurity_plasma",
+    "classical_impurity_flux",
+    "find_ambipolar_er",
+    "initialize_distributed_runtime_from_env",
+    "load_sfincs_input",
+    "monoenergetic_database",
+    "read_output",
+    "run_ambipolar_brent",
+    "run_from_namelist",
+    "run_monoenergetic_database",
+    "plot",
+    "run",
+    "run_profile",
+    "run_transport_matrix",
+    "write_output",
+]

dkx-2.4.0/src/dkx/_version.py

@@ -0,0 +1,3 @@
+"""The single source of truth for the DKX release version."""
+
+__version__ = "2.4.0"