Large refactor -- replace jitclasses and the relaxation solver with only functions and an explicit solver, pedagogical usability pass

.github/workflows/ci.yml

@@ -0,0 +1,42 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+concurrency:
+  group: ci-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    env:
+      MPLBACKEND: Agg
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.13"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Cache numba compilation
+        uses: actions/cache@v4
+        with:
+          path: |
+            openpopcon/__pycache__
+            ~/.cache/numba
+          key: numba-${{ runner.os }}-py${{ matrix.python-version }}-${{ hashFiles('openpopcon/core.py', 'openpopcon/lib/phys_lib.py', 'openpopcon/lib/openpopcon_util.py', 'pyproject.toml') }}
+
+      - name: Install project and test dependencies
+        run: uv sync --extra test --python ${{ matrix.python-version }}
+
+      - name: Run tests
+        run: uv run pytest

.gitignore

@@ -1,5 +1,11 @@
 .conda
 .conda/*
-__*
+__pycache__/
 .vscode
-outputs/*
+OpenPOPCON_outputs/
+.venv
+POPCON.png
+single_point.png
+dist/
+build/
+*.egg-info/

README.MD

@@ -5,26 +5,88 @@ OpenPOPCON is a tool for scoping Tokamak design and operation with 0-D fitted sc
 
 ## Installation
 
-We recommend the use of miniconda to manage the python environment. On a cluster, you may use `module load anaconda` to load anaconda. If you are running on your private machine, see [the anaconda project](https://docs.anaconda.com/miniconda/) for installation instructions.
+**With pip.** This installs the package and its dependencies, and works from anywhere afterwards:
 
-To install the required packages, run `conda env create -f environment.yml` in the root directory of this repository. This will create a new environment called `openpopcon` with all the required packages. From here, you can activate the environment with `conda activate openpopcon`. 
+```bash
+pip install git+https://github.com/hansec/OpenPOPCON.git
+```
+
+**From a clone, with uv.** If you have [uv](https://docs.astral.sh/uv/), nothing needs installing:
+
+```bash
+uv run example_run.py      # writes POPCON.png and single_point.png
+uv run --extra notebooks jupyter lab
+```
+
+**From a clone, with pip.** If you would rather manage the environment yourself, install the package in editable mode from the root of the repository:
+
+```bash
+pip install -e .              # add [notebooks] for jupyterlab, [netcdf4] for the h5netcdf backend
+```
+
+This pulls in the dependencies declared in `pyproject.toml` and makes `import openpopcon` work from anywhere.
 
 ## Usage
 
-Example notebooks are found under `resources/examples`. However, generally, the steps are:
+`example_run.py` is a complete working example; run it first to check your installation.
+
+Worked examples for several machines ship with the package. The installed copies are read only, so to get a set you can edit:
+
+```bash
+openpopcon examples ./my_scans
+cd my_scans/MANTA
+```
+
+Each example directory has its own settings file, a notebook, and where applicable a gEQDSK and a profiles file.
 
 ```python
 # 1. Import the POPCON class
 import openpopcon as op
 
-# 2. Load the settings files; make sure to create these, see the examples for a template
-pc = op.POPCON(settingsfile=settingsfile, plotsettingsfile=plotsettingsfile, scalinglawfile=scalinglawfile)
+# 2. Load the settings files; make sure to create these, see the examples for a template.
+#    scalinglawfile and plotsettingsfile both fall back to the copies shipped
+#    with the package if you leave them out.
+pc = op.POPCON(settingsfile=settingsfile, plotsettingsfile=plotsettingsfile)
 
 # 3. Run the code
 pc.run_POPCON()
 
 # 4. Plot the results
 pc.plot()
+
+# 5. Or look at one operating point in detail, here 60% of the Greenwald
+#    density at 9 keV volume-averaged Ti
+pc.single_point(0.6, 9.0)
+```
+
+`gfilename` and `profsfilename` inside a settings file are resolved relative to that settings file, so an example directory works wherever you copy it. The paths you pass to `POPCON(...)` are resolved relative to wherever you are running.
+
+Results written with `pc.write_output()` go to `./OpenPOPCON_outputs/`, or to whatever you pass as `directory=`.
+
+If a settings file has a mistake in it, OpenPOPCON reports every problem it finds at once, rather than stopping at the first:
+
+```
+ValueError: Found 2 problem(s) in .../POPCON_input_example.yaml:
+  - scalinglaw 'H98' is not defined. Available: H89, H98y2, H_NT23.
+  - Tmin_keV (11.0) must be less than Tmax_keV (6.0).
+```
+
+Grid points where power balance has no physical solution are reported in aggregate after a run and masked out of the plot.
+
+The trapped particle fraction that the neoclassical resistivity needs is taken from the gEQDSK when one is supplied. Without a gEQDSK it falls back to `f_t = sqrt(2*rho*a/R)`, which stays below 1 for the aspect ratios these examples use but would not for a spherical tokamak.
+
+## Development
+
+Run the tests with:
+
+```bash
+uv run --extra test pytest
+```
+
+The suite (in `tests/`) solves the MANTA example on a small grid and checks the results against a frozen golden file (`tests/data/golden_manta.json`), verifies physics/consistency invariants, and checks that settings validation accepts every shipped example while rejecting bad inputs. GitHub Actions runs it on every commit to an open pull request. If a change intentionally alters the computed results, regenerate the golden file and commit it alongside the change:
+
+```bash
+uv run python tests/generate_golden.py
 ```
 
 ## Bibliography
@@ -54,24 +116,37 @@ pc.plot()
 
 <img width="1668" alt="image" src="https://github.com/user-attachments/assets/a1234137-260f-4c5b-9124-c8312b3b9ef7" />
 
-OpenPOPCON employs a relaxation algorithm instead of solving power balance directly. This means it can often take quite long to converge on an answer compared to a code that would solve the problem explicitly. However, we have taken an approach with [numba](https://numba.pydata.org/) to speed up the code. This means that the code is compiled to machine code and can be run much faster than a traditional python code. To make full use of this feature, run OpenPOPCON from a Jupyter notebook. Every time you restart the Jupyter notebook, the code will be recompiled. The code structure is as follows:
+Power balance is solved in closed form. With the impurity fraction held fixed, every power integral at a given (n, T) is fixed, and the only self-consistency is that the confinement time is a power law in the heating power, `tauE = K * P_heat^alpha`. That makes the balance a single scalar equation with an exact root:
+
+```
+P_heat = (W_tot / K)^(1/(1+alpha))
+P_aux  = P_heat - P_OH - P_fus_heat + P_brem + P_synch
+```
+
+so there is no iteration, no convergence tolerance, and no points that fail to converge. Points with no physical solution, where impurity radiation exceeds the confinement loss or the balance has no real root, are set to `P_aux = 99999`, counted in a summary after the run, and masked out of the plot.
+
+The numerics are compiled to machine code with [numba](https://numba.pydata.org/). The compiled functions are cached to disk, so the compilation cost is paid the first time the code runs on a machine and not on every session; a second run in a fresh interpreter starts in a second or two. The code structure is as follows:
 
 ```python
+# module-level, @nb.njit(cache=True), operating on an immutable State namedtuple
+def get_profile(s, rho, profid)      # J, ne, ni, Ti, Te, q, ... by integer id
+def volume_integral(s, rho, func)    # integrates dV-like
+def P_aux_impfrac(s, n_e_20, T_i_keV)
+    "Closed-form power balance. Returns (P_aux, flag)."
+
+# module-level drivers; numba.prange parallelizes like an OMP for collapse(2)
+def solve_nT[_par](state, ...)       # P_aux over the whole grid, then
+def populate_outputs[_par](state, ...)  # every other diagnostic from it
+
 class POPCON:
     """The general class the user interacts with."""
-    self.algorithms: POPCON_algorithms: jitted 'class'
-    self.settings: POPCON_settings: 'class'
-    self.plotsettings: POPCON_plotsettings: 'class'
-    self.output: POPCON_data: jitted 'class'
-
-    @njit(parallel=True)
+    self.algorithms: POPCON_algorithms   # sets up State, calls the functions above
+    self.settings: POPCON_settings
+    self.plotsettings: POPCON_plotsettings
+    self.output: xarray.Dataset          # output arrays at power balance
+
     def run_POPCON(self):
-        """calls a function from self.algorithms and creates a new 
-        openPOPCON_results object. numba.prange allows for 
-        parallelization like an OMP for collapse(2)."""
-        for i in numba.prange(n_n):
-            for j in numba.prange(n_T):
-                self.openPOPCON_params.solve(n, T)
+        "Sets up and solves the whole n,T grid"
     def single_point(self, n, T):
         "Runs a single n,T and plots the solution profiles"
     def plot(self):
@@ -80,20 +155,4 @@ class POPCON:
         "Reads in a previous solution"
     def write_output(self):
         "Writes the current solution to a folder or zip archive"
-
-@jitclass
-class POPCON_algorithms:
-    """A compiled class that does all numerical calculations."""
-
-    def P_aux_relax_impfrac:
-        "Solves the equations for a given n,T, holding impfrac const."
-
-    def get_value(self):
-        "Helper functions to get all physical parameters"
-
-@jitclass    
-class POPCON_data:
-    """output arrays of physical parameters at power balance"""
-
-
 ```

example_run.py

@@ -1,10 +1,40 @@
-from src import openpopcon
+"""
+Minimal OpenPOPCON run, using the MANTA example that ships with the package.
+From a clone of the repository:
 
+    uv run example_run.py
 
-settingsfile = "./resources/examples/MANTA/POPCON_input_example.yaml"
-plotsettingsfile = "./resources/plotsettings.yml"
-scalinglawfile = "./resources/scalinglaws.yml"
+or, once installed with pip, from anywhere:
 
-pc = openpopcon.POPCON(settingsfile=settingsfile, plotsettingsfile=plotsettingsfile, scalinglawfile=scalinglawfile)
-pc.single_popcon(plot=True)
-pc.single_point(-1,0.6,12,True,show=True)
+    python example_run.py
+
+To start changing the machine, get an editable copy of the examples with
+`openpopcon examples ./my_scans` and point settingsfile at one of those.
+"""
+
+import os
+
+import matplotlib.pyplot as plt
+
+import openpopcon as op
+
+manta = op.example_dir("MANTA")
+settingsfile = os.path.join(manta, "POPCON_input_example.yaml")
+plotsettingsfile = os.path.join(manta, "plotsettings.yml")
+
+# scalinglawfile defaults to the copy shipped with the package
+pc = op.POPCON(settingsfile=settingsfile, plotsettingsfile=plotsettingsfile)
+
+# solve power balance over the whole density/temperature grid
+pc.run_POPCON()
+
+# the POPCON itself
+fig, ax = pc.plot(show=False)
+fig.savefig("POPCON.png", dpi=150, bbox_inches="tight")
+print("Wrote POPCON.png")
+
+# profiles and a power breakdown at one operating point,
+# here 60% of the Greenwald density and 9 keV volume-averaged Ti
+pc.single_point(0.6, 9.0, plot=True, show=False)
+plt.savefig("single_point.png", dpi=150, bbox_inches="tight")
+print("Wrote single_point.png")

openpopcon/__init__.py

@@ -0,0 +1,41 @@
+"""
+OpenPOPCON: a 0-D tokamak scoping tool.
+
+    import openpopcon as op
+    pc = op.POPCON(settingsfile='POPCON_input_example.yaml')
+    pc.run_POPCON()
+    pc.plot()
+
+Copies of the worked examples can be written somewhere editable with
+
+    openpopcon examples ./my_scans
+
+Names not listed here live in openpopcon.core, which is importable if you need
+the physics functions or the State namedtuple directly.
+"""
+
+from .core import (
+    POPCON,
+    POPCON_scan,
+    POPCON_settings,
+    POPCON_plotsettings,
+    POPCON_algorithms,
+    POPCON_data_spec,
+    build_dataset,
+)
+from .lib.openpopcon_util import example_dir, list_examples
+
+__version__ = "2.0.0"
+
+__all__ = [
+    "POPCON",
+    "POPCON_scan",
+    "POPCON_settings",
+    "POPCON_plotsettings",
+    "POPCON_algorithms",
+    "POPCON_data_spec",
+    "build_dataset",
+    "example_dir",
+    "list_examples",
+    "__version__",
+]