dkx 2.2.0__tar.gz → 2.3.0__tar.gz

{dkx-2.2.0 → dkx-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dkx
-Version: 2.2.0
+Version: 2.3.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/
@@ -10,7 +10,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: jax
-Requires-Dist: numpy
+Requires-Dist: numpy>=2.1
 Requires-Dist: scipy
 Requires-Dist: h5py
 Requires-Dist: netCDF4
@@ -23,6 +23,7 @@ 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: sphinx-rtd-theme; extra == "docs"
@@ -92,45 +93,41 @@ throughout: 37.3 s on W7-X and 45.8 s on a finite-β precise-QA equilibrium, the
 latter split 16.4 s for the 21-point scan, 6.0 s for the `|B|` solve, and 22.5 s
 for the radial scan's 50 drift-kinetic solves.
 
-The same solve from Python (mirrors
-[`examples/getting_started/run_tokamak.py`](examples/getting_started/run_tokamak.py),
-which also builds the namelist from dicts and plots the result):
+The same solve from Python. No input file, no argument parsing — the
+parameters are the ones a deck carries, passed as keywords
+([`examples/1_basics/run_tokamak.py`](examples/1_basics/run_tokamak.py)):
 
 ```python
-from pathlib import Path
-from dkx.run import run_profile
-
-deck = Path("input.namelist")
-deck.write_text("""\
-&geometryParameters
-  geometryScheme = 1  ! circular tokamak: BHat = 1 + 0.1 cos(theta)
-  inputRadialCoordinate = 3
-  rN_wish = 0.3
-  B0OverBBar = 1.0  GHat = 1.0  IHat = 0.0  iota = 1.31
-  epsilon_t = 0.1  epsilon_h = 0.0  psiAHat = 0.045  aHat = 0.1
-/
-&speciesParameters
-  Zs = 1  mHats = 1.0  nHats = 1.0  THats = 0.5
-  dNHatdrHats = -6.0  dTHatdrHats = -3.0
-/
-&physicsParameters
-  Delta = 4.5694d-3  alpha = 1.0  nu_n = 8.4774d-3
-  Er = 0.0  collisionOperator = 1  ! pitch-angle scattering
-/
-&resolutionParameters
-  Ntheta = 15  Nzeta = 1  Nxi = 8  NL = 4  Nx = 6
-  solverTolerance = 1d-10
-/
-""")
-
-run = run_profile(deck, solve_method="auto", out_path=Path("sfincsOutput.h5"))
+import dkx
+
+run = dkx.run(
+    geometryScheme=1,          # circular tokamak: BHat = 1 + 0.1 cos(theta)
+    inputRadialCoordinate=3, rN_wish=0.3,
+    B0OverBBar=1.0, GHat=1.0, IHat=0.0, iota=1.31,
+    epsilon_t=0.1, epsilon_h=0.0, psiAHat=0.045, aHat=0.1,
+    Zs=[1.0], mHats=[1.0], nHats=[1.0], THats=[0.5],
+    dNHatdrHats=[-6.0], dTHatdrHats=[-3.0],
+    Ntheta=15, Nzeta=1, Nxi=8, NL=4, Nx=6,
+    Delta=4.5694e-3, alpha=1.0, nu_n=8.4774e-3,
+    Er=0.0, collisionOperator=1,   # 1 = pitch-angle scattering, 0 = Fokker-Planck
+)
 print("particle flux:", float(run.moments["particleFlux_vm_psiHat"][0]))
 print("bootstrap current <j.B>:", float(run.moments["FSABjHat"]))
 ```
 
-`run_profile` prints the Fortran-parity console flow, writes `sfincsOutput.h5`/`.nc`
-keyed by the SFINCS output names, and returns the state, solver statistics, and
-all moments in memory.
+`dkx.run` reads `RHSMode` from the case and dispatches, so the same call covers
+a profile-gradient run and a transport matrix. It also takes an
+`input.namelist` path — `dkx.run("input.namelist")` — and overrides on top of
+one, which is what a convergence scan is:
+
+```python
+dkx.run("input.namelist", Ntheta=25)     # same deck, finer grid
+dkx.run(..., out="sfincsOutput.h5")      # write a file; the suffix picks the format
+```
+
+Results come back in memory: `run.moments` is keyed by the SFINCS output names,
+and `run.solve_result` carries the residual, the iteration count, and whether it
+converged.
 
 ## Parity with SFINCS Fortran v3
 
@@ -166,7 +163,7 @@ flow corrections, an extended-collisionality Sugama operator, monoenergetic
 database mode, batched GPU scans, a bounce-averaged 1/ν surrogate) — lives in
 [docs/feature_matrix.rst](docs/feature_matrix.rst).
 
-*Reproduce with the drivers in [`examples/parity/`](examples/parity/).*
+*Reproduce with the drivers in [`tools/parity/`](tools/parity/).*
 
 ### Interpreting a difference
 
@@ -269,7 +266,7 @@ matched-deck SFINCS Fortran v3 cross-check points at solver precision. The
 quasi-helically symmetric HSX case shows the suppressed 1/ν branch that W7-X and
 TJ-II retain.
 
-*Reproduce with `python examples/paper_benchmarks/monoenergetic_icnts_w7x.py`
+*Reproduce with `python tools/paper_benchmarks/monoenergetic_icnts_w7x.py`
 (and the `_tjii` / `_hsx` companions).*
 
 ## Low-collisionality bootstrap convergence
@@ -283,7 +280,7 @@ asymptote — the 1/ν-regime offset does not decay without orbit precession —
 a small finite `EStar` detaches below `nuPrime ~ 1e-3` and flattens back toward
 the asymptote, the E×B-precession dip.
 
-*Reproduce with `python examples/paper_benchmarks/shaing_callen_convergence.py`.*
+*Reproduce with `python tools/paper_benchmarks/shaing_callen_convergence.py`.*
 
 ## Ambipolar Er and electron roots
 
@@ -296,8 +293,8 @@ sign, and follows the physical branch by radial continuity to reproduce the
 electron-root → ion-root crossover near `ρ ~ 0.6`, matching the reference `E_r`
 within the digitization uncertainty (~1.5 kV/m mean difference).
 
-*Reproduce with `python examples/paper_benchmarks/w7x_ambipolar_er.py` and
-`python examples/paper_benchmarks/electron_root_optimization.py`.*
+*Reproduce with `python tools/paper_benchmarks/w7x_ambipolar_er.py` and
+`python tools/paper_benchmarks/electron_root_optimization.py`.*
 
 ## Impurity transport
 
@@ -310,7 +307,7 @@ diagnostic recovers the exact `-Z` density-peaking coefficient and the classical
 1/2 collisional screening coefficient, with an autodiff ion-temperature-gradient
 derivative verified against finite differences.
 
-*Reproduce with `python examples/paper_benchmarks/impurity_transport.py`.*
+*Reproduce with `python tools/paper_benchmarks/impurity_transport.py`.*
 
 ## Kinetic-in-the-loop bootstrap
 
@@ -324,7 +321,7 @@ error the kinetic-in-the-loop iteration removes by construction — and one
 `jax.value_and_grad` differentiates the total bootstrap current through the
 equilibrium → Boozer → kinetic chain.
 
-*Reproduce with `python examples/paper_benchmarks/bootstrap_consistency_kinetic_loop.py`
+*Reproduce with `python tools/paper_benchmarks/bootstrap_consistency_kinetic_loop.py`
 (needs the optional `vmex` + `booz_xform_jax` companions).*
 
 To put that inside an optimizer instead of after it, `dkx.bootstrap.KineticBootstrapCurrent`

{dkx-2.2.0 → dkx-2.3.0}/README.md

@@ -62,45 +62,41 @@ throughout: 37.3 s on W7-X and 45.8 s on a finite-β precise-QA equilibrium, the
 latter split 16.4 s for the 21-point scan, 6.0 s for the `|B|` solve, and 22.5 s
 for the radial scan's 50 drift-kinetic solves.
 
-The same solve from Python (mirrors
-[`examples/getting_started/run_tokamak.py`](examples/getting_started/run_tokamak.py),
-which also builds the namelist from dicts and plots the result):
+The same solve from Python. No input file, no argument parsing — the
+parameters are the ones a deck carries, passed as keywords
+([`examples/1_basics/run_tokamak.py`](examples/1_basics/run_tokamak.py)):
 
 ```python
-from pathlib import Path
-from dkx.run import run_profile
-
-deck = Path("input.namelist")
-deck.write_text("""\
-&geometryParameters
-  geometryScheme = 1  ! circular tokamak: BHat = 1 + 0.1 cos(theta)
-  inputRadialCoordinate = 3
-  rN_wish = 0.3
-  B0OverBBar = 1.0  GHat = 1.0  IHat = 0.0  iota = 1.31
-  epsilon_t = 0.1  epsilon_h = 0.0  psiAHat = 0.045  aHat = 0.1
-/
-&speciesParameters
-  Zs = 1  mHats = 1.0  nHats = 1.0  THats = 0.5
-  dNHatdrHats = -6.0  dTHatdrHats = -3.0
-/
-&physicsParameters
-  Delta = 4.5694d-3  alpha = 1.0  nu_n = 8.4774d-3
-  Er = 0.0  collisionOperator = 1  ! pitch-angle scattering
-/
-&resolutionParameters
-  Ntheta = 15  Nzeta = 1  Nxi = 8  NL = 4  Nx = 6
-  solverTolerance = 1d-10
-/
-""")
-
-run = run_profile(deck, solve_method="auto", out_path=Path("sfincsOutput.h5"))
+import dkx
+
+run = dkx.run(
+    geometryScheme=1,          # circular tokamak: BHat = 1 + 0.1 cos(theta)
+    inputRadialCoordinate=3, rN_wish=0.3,
+    B0OverBBar=1.0, GHat=1.0, IHat=0.0, iota=1.31,
+    epsilon_t=0.1, epsilon_h=0.0, psiAHat=0.045, aHat=0.1,
+    Zs=[1.0], mHats=[1.0], nHats=[1.0], THats=[0.5],
+    dNHatdrHats=[-6.0], dTHatdrHats=[-3.0],
+    Ntheta=15, Nzeta=1, Nxi=8, NL=4, Nx=6,
+    Delta=4.5694e-3, alpha=1.0, nu_n=8.4774e-3,
+    Er=0.0, collisionOperator=1,   # 1 = pitch-angle scattering, 0 = Fokker-Planck
+)
 print("particle flux:", float(run.moments["particleFlux_vm_psiHat"][0]))
 print("bootstrap current <j.B>:", float(run.moments["FSABjHat"]))
 ```
 
-`run_profile` prints the Fortran-parity console flow, writes `sfincsOutput.h5`/`.nc`
-keyed by the SFINCS output names, and returns the state, solver statistics, and
-all moments in memory.
+`dkx.run` reads `RHSMode` from the case and dispatches, so the same call covers
+a profile-gradient run and a transport matrix. It also takes an
+`input.namelist` path — `dkx.run("input.namelist")` — and overrides on top of
+one, which is what a convergence scan is:
+
+```python
+dkx.run("input.namelist", Ntheta=25)     # same deck, finer grid
+dkx.run(..., out="sfincsOutput.h5")      # write a file; the suffix picks the format
+```
+
+Results come back in memory: `run.moments` is keyed by the SFINCS output names,
+and `run.solve_result` carries the residual, the iteration count, and whether it
+converged.
 
 ## Parity with SFINCS Fortran v3
 
@@ -136,7 +132,7 @@ flow corrections, an extended-collisionality Sugama operator, monoenergetic
 database mode, batched GPU scans, a bounce-averaged 1/ν surrogate) — lives in
 [docs/feature_matrix.rst](docs/feature_matrix.rst).
 
-*Reproduce with the drivers in [`examples/parity/`](examples/parity/).*
+*Reproduce with the drivers in [`tools/parity/`](tools/parity/).*
 
 ### Interpreting a difference
 
@@ -239,7 +235,7 @@ matched-deck SFINCS Fortran v3 cross-check points at solver precision. The
 quasi-helically symmetric HSX case shows the suppressed 1/ν branch that W7-X and
 TJ-II retain.
 
-*Reproduce with `python examples/paper_benchmarks/monoenergetic_icnts_w7x.py`
+*Reproduce with `python tools/paper_benchmarks/monoenergetic_icnts_w7x.py`
 (and the `_tjii` / `_hsx` companions).*
 
 ## Low-collisionality bootstrap convergence
@@ -253,7 +249,7 @@ asymptote — the 1/ν-regime offset does not decay without orbit precession —
 a small finite `EStar` detaches below `nuPrime ~ 1e-3` and flattens back toward
 the asymptote, the E×B-precession dip.
 
-*Reproduce with `python examples/paper_benchmarks/shaing_callen_convergence.py`.*
+*Reproduce with `python tools/paper_benchmarks/shaing_callen_convergence.py`.*
 
 ## Ambipolar Er and electron roots
 
@@ -266,8 +262,8 @@ sign, and follows the physical branch by radial continuity to reproduce the
 electron-root → ion-root crossover near `ρ ~ 0.6`, matching the reference `E_r`
 within the digitization uncertainty (~1.5 kV/m mean difference).
 
-*Reproduce with `python examples/paper_benchmarks/w7x_ambipolar_er.py` and
-`python examples/paper_benchmarks/electron_root_optimization.py`.*
+*Reproduce with `python tools/paper_benchmarks/w7x_ambipolar_er.py` and
+`python tools/paper_benchmarks/electron_root_optimization.py`.*
 
 ## Impurity transport
 
@@ -280,7 +276,7 @@ diagnostic recovers the exact `-Z` density-peaking coefficient and the classical
 1/2 collisional screening coefficient, with an autodiff ion-temperature-gradient
 derivative verified against finite differences.
 
-*Reproduce with `python examples/paper_benchmarks/impurity_transport.py`.*
+*Reproduce with `python tools/paper_benchmarks/impurity_transport.py`.*
 
 ## Kinetic-in-the-loop bootstrap
 
@@ -294,7 +290,7 @@ error the kinetic-in-the-loop iteration removes by construction — and one
 `jax.value_and_grad` differentiates the total bootstrap current through the
 equilibrium → Boozer → kinetic chain.
 
-*Reproduce with `python examples/paper_benchmarks/bootstrap_consistency_kinetic_loop.py`
+*Reproduce with `python tools/paper_benchmarks/bootstrap_consistency_kinetic_loop.py`
 (needs the optional `vmex` + `booz_xform_jax` companions).*
 
 To put that inside an optimizer instead of after it, `dkx.bootstrap.KineticBootstrapCurrent`

{dkx-2.2.0 → dkx-2.3.0}/dkx/__init__.py

@@ -10,12 +10,49 @@ from __future__ import annotations
 # CLI invocations unless the user explicitly disables it. This improves cold-start
 # performance without requiring environment configuration.
 import os
+import sys as _sys
+import types as _types
 import tempfile
 
 # Suppress low-value XLA/PjRt C++ warning chatter by default. Users can still
 # override this before importing dkx if they need backend debug logs.
 os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
 
+
+def _check_numpy() -> None:
+    """Fail with the actual problem instead of a NameError from numpy internals.
+
+    ``jax`` requires ``numpy>=2.1``.  Installing dkx into an environment that
+    holds a conda-managed numpy 1.x -- the Anaconda base environment is the
+    common case -- resolves without complaint and then dies on the first
+    ``import jax`` with::
+
+        File numpy/core/getlimits.py, line 606, in smallest_normal
+        NameError: name 'isnan' is not defined
+
+    raised while ``ml_dtypes`` probes ``bfloat16``'s ``finfo``.  That traceback
+    names neither dkx, nor jax, nor numpy's version as the cause, so it costs a
+    user real time.  Checking here is cheap: numpy is imported either way, and
+    this runs before any jax import.
+    """
+    import numpy  # noqa: PLC0415
+
+    major = int(numpy.__version__.split(".", 1)[0])
+    if major < 2:
+        raise ImportError(
+            f"dkx needs numpy>=2.1 (jax's own requirement); found {numpy.__version__}. "
+            "numpy 1.x crashes inside ml_dtypes at 'import jax'. Fix the "
+            "environment with:\n"
+            "    pip install -U 'numpy>=2.1' jax\n"
+            "or, better, install into a clean environment rather than the "
+            "Anaconda base one:\n"
+            "    conda create -n dkx python=3.11 -y && conda activate dkx\n"
+            "    pip install dkx"
+        )
+
+
+_check_numpy()
+
 _distributed_runtime_initialized = False
 
 
@@ -132,19 +169,28 @@ if _disable_cache not in {"1", "true", "yes", "on"}:
             except OSError:
                 return False
 
+        # Versioned because the cache is bounded now (see the cap below) and
+        # JAX's LRU keeps a sidecar "-atime" file per entry.  A directory
+        # filled in before the cap existed has none, so every eviction pass
+        # tries to touch a file that was never written and warns -- measured,
+        # 12 warnings on a single small solve against a legacy cache, and 0
+        # against a fresh one.  Starting a new directory is what makes the
+        # bound work; the old one simply stops being written to.
+        _CACHE_DIR_NAME = "jax_compilation_cache"
+
         cache_override = os.environ.get("DKX_COMPILATION_CACHE_DIR", "").strip()
         if cache_override:
             default_cache_dir = cache_override
         else:
             xdg_cache = os.environ.get("XDG_CACHE_HOME", "").strip()
             if xdg_cache:
-                default_cache_dir = os.path.join(xdg_cache, "dkx", "jax_compilation_cache")
+                default_cache_dir = os.path.join(xdg_cache, "dkx", _CACHE_DIR_NAME)
             else:
-                default_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "dkx", "jax_compilation_cache")
+                default_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "dkx", _CACHE_DIR_NAME)
         try:
             os.makedirs(default_cache_dir, exist_ok=True)
         except OSError:
-            default_cache_dir = os.path.join(tempfile.gettempdir(), "dkx", "jax_compilation_cache")
+            default_cache_dir = os.path.join(tempfile.gettempdir(), "dkx", _CACHE_DIR_NAME)
             try:
                 os.makedirs(default_cache_dir, exist_ok=True)
             except OSError:
@@ -153,7 +199,7 @@ if _disable_cache not in {"1", "true", "yes", "on"}:
             # Some environments (CI sandboxes, read-only homes) can create the directory but
             # cannot write compilation entries. Fall back to a tempdir cache to avoid noisy
             # warnings and degraded cold-start performance.
-            default_cache_dir = os.path.join(tempfile.gettempdir(), "dkx", "jax_compilation_cache")
+            default_cache_dir = os.path.join(tempfile.gettempdir(), "dkx", _CACHE_DIR_NAME)
             try:
                 os.makedirs(default_cache_dir, exist_ok=True)
             except OSError:
@@ -207,6 +253,20 @@ try:
             )
         except ValueError:
             pass
+        # The cache is deliberately NOT size-capped.  Setting
+        # jax_compilation_cache_max_size turns on jax's LRU, which writes a
+        # sidecar "-atime" file per entry and does size bookkeeping on every
+        # write.  With the zero thresholds above -- which exist so even small
+        # kernels are cached -- that is measurably expensive: on
+        # tests/test_monoenergetic_database.py, 32 s and 1792 files capped
+        # against 20 s and 896 uncapped, a 60% penalty on every run.  CI proved
+        # it at scale, nine of ten coverage shards crossing a 10-minute timeout
+        # they had been finishing in four to eight.
+        #
+        # So the cache grows without bound.  That is a disk-space cost -- it
+        # reached 782 MB over 64k entries on a development machine -- and the
+        # remedy is to delete the directory, which loses nothing but compile
+        # time.  Paying 60% on every run to avoid it is the worse trade.
 except Exception:
     # Keep import lightweight for tooling that inspects the package without JAX.
     pass
@@ -233,6 +293,7 @@ from .inputs import SfincsInput, load_sfincs_input  # 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"),
@@ -246,6 +307,8 @@ _LAZY_EXPORTS = {
 
 
 def __getattr__(name: str):
+    if name == "run":
+        return _lazy_run_module()
     try:
         module_name, attr = _LAZY_EXPORTS[name]
     except KeyError:
@@ -261,6 +324,20 @@ def __dir__() -> list[str]:
     return sorted(set(globals()) | set(_LAZY_EXPORTS))
 
 
+# ``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.
@@ -309,9 +386,11 @@ __all__ = [
     "run_ambipolar_brent",
     "run_from_namelist",
     "run_monoenergetic_database",
+    "plot",
+    "run",
     "run_profile",
     "run_transport_matrix",
     "write_output",
 ]
 
-__version__ = "2.2.0"
+__version__ = "2.3.0"

{dkx-2.2.0 → dkx-2.3.0}/dkx/plotting.py

@@ -54,6 +54,29 @@ def _add_heatmap(ax, data: dict[str, object], key: str, *, title: str | None = N
     return True
 
 
+
+def _theta_zeta_oriented(b_hat: np.ndarray, data: dict[str, object]) -> np.ndarray:
+    """Return ``BHat`` as ``(theta, zeta)`` whatever order it was stored in.
+
+    The output file stores it ``(zeta, theta)`` (Fortran layout) while the
+    operator hands it over ``(theta, zeta)``, and ``imshow`` puts axis 0 on the
+    vertical.  Labelling without checking is how every axisymmetric figure came
+    out with its variation on the axis marked zeta: the tokamak's Nzeta=1 strip
+    was drawn horizontally and called zeta when it was theta.
+
+    Orientation is decided by matching the array's dimensions against the run's
+    own ``Ntheta``/``Nzeta`` rather than by assuming an order.  When the two are
+    equal there is nothing to disambiguate, and the stored layout is kept.
+    """
+    b_hat = np.asarray(b_hat)
+    if b_hat.ndim != 2:
+        return b_hat
+    n_theta = int(np.asarray(data.get("Ntheta", 0)).reshape(-1)[0]) if "Ntheta" in data else 0
+    n_zeta = int(np.asarray(data.get("Nzeta", 0)).reshape(-1)[0]) if "Nzeta" in data else 0
+    if n_theta and n_zeta and n_theta != n_zeta and b_hat.shape == (n_zeta, n_theta):
+        return b_hat.T
+    return b_hat
+
 def _add_profile(ax, data: dict[str, object], x: np.ndarray, key: str, *, title: str | None = None) -> bool:
     if key not in data:
         return False
@@ -220,8 +243,9 @@ def plot_sfincs_output_summary(
         axes[0].set_ylabel("FSABFlow_vs_x")
         axes[0].grid(True, alpha=0.25)
     else:
-        theta = np.asarray(data.get("theta", np.arange(b_hat.shape[0]))).ravel()
-        axes[0].plot(theta, b_hat[:, 0], lw=1.8)
+        oriented = _theta_zeta_oriented(b_hat, data)
+        theta = np.asarray(data.get("theta", np.arange(oriented.shape[0]))).ravel()
+        axes[0].plot(theta, oriented[:, 0], lw=1.8)
         axes[0].set_title("BHat(theta, zeta=0)")
         axes[0].set_xlabel("theta")
         axes[0].set_ylabel("BHat")
@@ -252,15 +276,118 @@ def plot_sfincs_output_summary(
         )
         axes[1].set_title("Run summary")
 
-    im = axes[2].imshow(b_hat, aspect="auto", origin="lower")
-    axes[2].set_title("BHat(theta, zeta)")
-    axes[2].set_xlabel("zeta index")
-    axes[2].set_ylabel("theta index")
-    axes[2].set_xticks([0, max(0, b_hat.shape[-1] - 1)])
-    axes[2].set_xticklabels(["0", f"{float(zeta[-1]):.2f}"])
+    oriented = _theta_zeta_oriented(b_hat, data)
+    im = axes[2].imshow(oriented, aspect="auto", origin="lower")
+    axes[2].set_title(r"$\hat B(\theta, \zeta)$")
+    axes[2].set_xlabel(r"$\zeta$ index")
+    axes[2].set_ylabel(r"$\theta$ index")
+    if zeta.size and oriented.shape[-1] == zeta.size:
+        axes[2].set_xticks([0, max(0, oriented.shape[-1] - 1)])
+        axes[2].set_xticklabels(["0", f"{float(zeta[-1]):.2f}"])
     fig.colorbar(im, ax=axes[2], fraction=0.046, pad=0.04)
 
-    fig.suptitle(f"SFINCS output summary: {Path(input_h5).name}", y=1.03)
+    # y=1.03 put the title outside the figure, where constrained_layout does
+    # not reserve room for it, so it landed on the subplot titles.  Let the
+    # layout engine place it.
+    fig.suptitle(f"SFINCS output summary: {Path(input_h5).name}")
     fig.savefig(output_png, dpi=200, bbox_inches="tight")
     plt.close(fig)
     return output_png.resolve()
+
+
+
+def _shown(written: Path) -> Path:
+    """Display an already-written figure, for IDE and notebook users.
+
+    The panel builders save and close their own figure, so there is nothing
+    left to ``plt.show()``.  Re-reading the PNG and showing it is what puts it
+    in Spyder's plots pane, and it costs one imread.
+    """
+    try:
+        image = plt.imread(written)
+    except Exception:  # pragma: no cover - unreadable/exotic format
+        return written
+    figure = plt.figure(figsize=(13.5, 8.0))
+    axis = figure.add_axes((0.0, 0.0, 1.0, 1.0))
+    axis.imshow(image)
+    axis.axis("off")
+    # block=False so a script run from a terminal is not held open by a
+    # window nobody is there to close; Spyder and Jupyter render regardless.
+    plt.show(block=False)
+    return written
+
+
+def plot(source, out="dkx_panels.png", *, style="panels", show=False):
+    """Plot a run or an output file.  One call, one figure, returns its path.
+
+    ``source`` is whatever you already have:
+
+    - the object :func:`dkx.run` returned, when the run was given ``out=``;
+    - a path to an ``sfincsOutput`` file, DKX's or Fortran SFINCS's --- the
+      layout is the same, so both work;
+    - a directory containing one.
+
+    The figure is the same six-panel one ``dkx --plot`` writes: the
+    monoenergetic coefficients, ``|B|`` on the surface, the bootstrap current
+    and the species fluxes.  ``out`` picks the format by suffix.
+
+    ``style="summary"`` selects the compact three-panel page instead.  Only
+    that style expands with a ``.pdf`` suffix, into a four-page diagnostics
+    book; ``style="panels"`` writes its one figure whatever the suffix.
+
+    For a figure this does not draw, read the numbers off ``run.moments`` and
+    use matplotlib directly --- ``examples/1_basics/plot_custom.py`` shows
+    that, and it is not a fallback so much as the normal way to make a figure
+    for a paper.
+
+    Args:
+        source: a run object, an output-file path, or a directory.
+        out: destination path; the suffix selects png or pdf.  A ``.pdf``
+            is multi-page only under ``style="summary"``.
+        style: ``"panels"`` (default) for the ``dkx --plot`` figure, or
+            ``"summary"`` for the compact three-panel page.
+        show: also display the figure.  For Spyder, Jupyter or any IDE where
+            you want it in the plots pane rather than only on disk; a no-op
+            under a headless backend.
+
+    Returns:
+        The :class:`~pathlib.Path` actually written.
+    """
+    out = Path(out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    path = getattr(source, "output_path", None) or (
+        source if isinstance(source, (str, Path)) else None)
+    if path is not None:
+        path = Path(path)
+        if path.is_dir():
+            candidates = sorted(
+                p for p in path.iterdir()
+                if p.suffix in {".h5", ".nc", ".npz"} and "sfincsOutput" in p.name
+            )
+            if not candidates:
+                raise FileNotFoundError(f"no sfincsOutput file in {path}")
+            path = candidates[0]
+        if style == "summary":
+            written = plot_sfincs_output_summary(input_h5=path, output_png=out)
+            return _shown(written) if show else written
+        if style != "panels":
+            raise ValueError(f"style must be 'panels' or 'summary'; got {style!r}")
+        from dkx.representative import plot_output_file  # noqa: PLC0415
+
+        written = plot_output_file(path, out)
+        return _shown(written) if show else written
+
+    if hasattr(source, "moments"):
+        raise ValueError(
+            "this run has no output file to plot: pass out=... to dkx.run so the "
+            "solve writes one, e.g.\n"
+            '    run = dkx.run(case, out="sfincsOutput.h5")\n'
+            "    dkx.plot(run)\n"
+            "Plotting cannot reuse the in-memory result because the panels read "
+            "the output file's full dataset, not just the moments."
+        )
+    raise TypeError(
+        "plot() takes the object dkx.run() returned, a path to an sfincsOutput "
+        f"file, or a directory holding one; got {type(source).__name__}"
+    )