Restructure PIPT and POPT around explicit schemes, analyses and optimizers

.github/workflows/tests.yml

@@ -4,7 +4,6 @@
 # - Compute test coverage and submit to coveralls.io
 # - Also config for macOS and/or Windows
 # - Also config for conda
-# - Lint
 
 name: CI tests
 
@@ -19,6 +18,19 @@ on:
       - '*'
 
 jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.12"
+    - name: Install ruff
+      run: python -m pip install ruff
+    - name: Run ruff
+      run: ruff check src tests
+
   bundled:
 
     runs-on: ubuntu-latest
@@ -28,15 +40,15 @@ jobs:
         python-version: ["3.10", 3.11, 3.12]
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip setuptools
-        python -m pip install -e .
+        python -m pip install -e ".[dev]"
     - name: Launch tests
       run: |
-        pytest
+        pytest --cov=src --cov-report=term

.gitignore

@@ -147,3 +147,6 @@ docs-generated/
 
 # mergetoolfiles
 *.orig
+
+# Stray test artifact
+PET-test-log

README.md

@@ -16,7 +16,7 @@ at NORCE Norwegian Research Centre AS.
 Before installing ensure you have python3 pre-requisites. On a Debian system run:
 
 ```
-sudo upt-get update
+sudo apt-get update
 sudo apt-get install python3
 sudo apt-get install python3-pip
 sudo apt-get install python3-venv
@@ -41,12 +41,6 @@ python3 -m venv venv-PET
 source venv-PET/bin/activate
 ```
 
-Some additional features might be not part of your default installation and need to be set in the Python (virtual) environment manually:
-
-```
-python3 -m pip install wheel
-```
-
 If you do not install PET inside a virtual environment,
 you may have to include the `--user` option in the following
 (to install to your local Python site packages, usually located in `~/.local`).
@@ -61,6 +55,76 @@ python3 -m pip install -e .
 - The `-e` option installs PET such that changes to it take effect immediately
   (without re-installation).
 
+To also install the tools needed for running tests and linting locally:
+
+```sh
+python3 -m pip install -e ".[dev]"
+```
+
+## Documentation
+
+The [configuration reference](docs/configuration.md) lists every key of the
+`dataassim`, `ensemble`, `optim` and `simulator` sections; the
+[architecture page](docs/architecture.md) explains how a run is put together
+and where a new scheme, analysis, localization, optimizer or simulator goes.
+
+## Command-line interface
+
+Installing PET also installs a `pet` command for working with config files:
+
+```sh
+pet validate my_config.toml   # check a config file for missing/invalid keys
+pet convert my_case.pipt      # convert a legacy .pipt/.popt file to .toml (or --to yaml)
+pet migrate my_config.toml    # update a config file to the current schema
+pet version                   # print the installed PET version
+```
+
+### Config schema change: `daalg` becomes `scheme`
+
+The analysis flavour is a parameter of an algorithm, not a separate algorithm,
+so the two-element `daalg` key has been replaced by a single `scheme` key:
+
+```toml
+# before                              # after
+[dataassim]                           [dataassim]
+daalg = ["esmda", "esmda"]            scheme = "esmda"
+analysis = "approx"                   analysis = "approx"
+```
+
+`pet migrate` performs this rewrite in place, keeping the original as
+`<config>.bak`. Use `--dry-run` to preview. Loading a config that still uses
+`daalg` raises an error pointing at the command. For a legacy `.pipt`/`.popt`
+file, convert first and then migrate:
+
+```sh
+pet convert my_case.pipt && pet migrate my_case.toml
+```
+
+The same change is reflected in the Python API, where one constructor per
+algorithm now takes the flavour as an argument:
+
+```python
+from pipt import ESMDA, available_schemes
+
+scheme = ESMDA(cfg_da, cfg_en, sim)   # flavour comes from the config's `analysis`
+result = scheme.run_assimilation()    # the scheme owns its iteration loop
+
+available_schemes()   # every valid (scheme, analysis) pair
+```
+
+`analysis=` overrides the config when passed. `ESMDA.assimilate(cfg_da, cfg_en,
+sim)` is the one-line form for when the scheme object is not needed afterwards;
+it returns the same `AssimilationResult`, whose `x` is the posterior ensemble.
+
+The eighteen per-flavour classes this used to produce (`esmda_approx`,
+`lmenrml_full`, ...) are gone: each was a one-line subclass pinning the
+flavour a constructor argument already expresses. Use `ESMDA(..., analysis=
+"approx")` and friends instead.
+
+Running a data-assimilation or optimization job itself is still done from a
+Python driver script that wires up your forward simulator/cost function -- see
+the tutorials below.
+
 ## Examples
 
 PET needs to be set up with a configuration file. See the example [repository](https://github.com/Python-Ensemble-Toolbox/Examples) for inspiration.
@@ -75,8 +139,8 @@ Some basic plotting functionality is provided [here](https://github.com/Python-E
 
 ## Tutorials
 
-- A PIPT tutorial is found [here](https://python-ensemble-toolbox.github.io/PET/tutorials/pipt/tutorial_pipt)
-- A POPT tutorial is found [here](https://python-ensemble-toolbox.github.io/PET/tutorials/popt/tutorial_popt)
+- A PIPT tutorial is found [here](https://python-ensemble-toolbox.github.io/PET/tutorials/pipt/TinyBox/tutorial_pipt)
+- A POPT tutorial is found [here](https://python-ensemble-toolbox.github.io/PET/tutorials/popt/5Spot/tutorial_popt)
 
 ## Suggested readings:
 

docs/architecture.md

@@ -0,0 +1,134 @@
+# Architecture
+
+PET is three layers. `ensemble` is the foundation: the base ensemble that runs
+a forward simulator over the members, the checkpoint mixin, the loggers, the
+`ForwardSimulator` protocol. `pipt` (data assimilation) and `popt`
+(optimisation) build on it and never import each other; `ensemble` imports
+neither. `misc` holds the data structures and the observed-data reader,
+`input_output` the configuration boundary, `simulator` the analytical models
+and the wrappers around external simulators.
+
+## A run
+
+1. **Configuration.** `input_output.read_config.read(file)` returns the
+   problem, simulator and ensemble sections as plain dictionaries in one
+   canonical form (`input_output.config.normalize`); `validate` says what a run
+   would fail on. A dictionary built in a script gets the same treatment when
+   the ensemble is constructed. See the [configuration reference](configuration.md).
+2. **Ensemble.** `pipt.ensembles.AssimilationEnsemble(keys_da, keys_en, sim)`
+   draws or loads the prior (`(nx, ne)` array; its variable rows are a
+   `StateLayout`), reads the observations, fixes the **data layout** and
+   builds the observation vector and variance in that order, and sets up
+   localization and compression. It owns the random stream (`seed`).
+3. **Scheme.** `ESMDA(keys_da, keys_en, sim, analysis="approx")` and the
+   other schemes build the ensemble (or take one passed as `ensemble=`), bind
+   an analysis object, and `run_assimilation()`: forecast the prior, then call
+   `update_step()` until a criterion fires or `maxiter` updates are done.
+   `ESMDA.assimilate(...)` is the one-line form.
+4. **Result.** An `AssimilationResult` (`x`, `data_misfit`,
+   `prior_data_misfit`, `nit`, `success`, `message`, `why_stop`), a
+   `dict` subclass with attribute access like SciPy's `OptimizeResult`.
+
+## The scheme contract
+
+`pipt.update_schemes.core.AssimilationScheme` owns the loop, the convergence
+checks (`misfit_tol`, `step_tol`, plus the scheme's own `check_convergence`),
+the run table, restart, QA/QC and saving. A scheme supplies:
+
+- `update_step() -> StepReport(accepted, state, misfit, why_stop)`: one
+  iteration, retries included (LM-EnRML re-damps inside it). The loop commits
+  the returned state and misfit; a scheme never assigns them itself.
+- `score()`: the per-member data misfit of the current forecast.
+- `log_columns()`: its columns of the run table.
+- Hooks: `after_prior_forecast`, `after_analysis`, `after_forecast`,
+  `after_accepted_iteration`, `after_loop`.
+- `COMPATIBLE_ANALYSES`: the analysis flavours it accepts.
+- `RESTART_ATTRIBUTES`: the attributes a checkpoint must carry for it.
+
+Registration is one line: `register_scheme(name, analysis, cls)` in
+`pipt.update_schemes.registry` (`available_schemes()` lists every pair). The
+iterative family shares `IterativeEnRML`, where LM-EnRML and GN-EnRML differ
+only in ten small hooks around their control parameter.
+
+## Analyses
+
+An analysis (`pipt.update_schemes.analysis`) is a class with
+`update(enX, enY, enE, **kwargs) -> AnalysisResult`, returning exactly one of
+a state-space `step`, a weight-space `w_step` or a `W_step`. The scheme turns
+it into a proposal with `propose_state(result, step_scale)`. The flavours are
+`approx`, `full`, `subspace`, `margis` and the multilevel `hybrid`;
+`register_analysis` adds one. Analyses read what they need from the scheme:
+`state_scaling`, `scale_data`, `proj`, `cov_data`, `trunc_energy`, `lam`.
+
+## Data on the analysis path
+
+Everything the analyses see is a matrix in one fixed row order:
+
+- `DataLayout` (`misc.structures`): the order of the data vector, computed
+  once from the observed frame -- label-major, then data type, empty cells
+  skipped. The ensemble exposes `obs_vector` and `obs_variance` built from it.
+- `PredictedData`: the `(nd, ne)` forecast filled directly from each member's
+  simulator output through the layout, scaled as the observations were, with
+  compressed vintages reduced on the way in. `pred_data.matrix` is what the
+  schemes read; `pred_data.to_frame()` is the frame view.
+- Adjoints: an `(nd, nx, ne)` array in the same rows, when the simulator
+  computes them.
+- `StateLayout`: the `{variable: (start, stop)}` rows of the state array; the
+  ensemble's `state_layout` converts to and from dictionaries, builds member
+  inputs for the simulator, and clips to the prior's limits.
+
+`PETDataFrame` remains as the table observed data arrive in and results are
+saved as; it is built from the matrices on demand, never on the analysis path.
+
+## Forecast
+
+`BaseEnsemble.calc_prediction(enX)` runs one level: member inputs
+(`_simulator_input`), a backend (`_run_members`: in sequence, a local process
+pool, or the wrapper's HPC queue), crash replacement, adjoint splitting, and
+the raw outputs kept as `member_outputs`. `ForecastMixin.forecast` then fills
+`pred_data`, corrects multilevel levels, applies the `scale` option and saves
+what was asked for. Outlier replacement (`OutlierMixin`) reorders the raw
+outputs and the state together.
+
+A simulator is anything satisfying `ensemble.protocols.ForwardSimulator`: an
+`input_dict` and `run_fwd_sim(state, member_index)` returning one dict per
+report point (or a DataFrame), `False` on failure, or `(output, adjoint)`.
+`simulator/vanderpol.py` is the smallest complete example.
+
+## Restart, random numbers, logging
+
+- One checkpoint per run, on `RestartMixin` (`ensemble.checkpoint`): the
+  loop's bookkeeping, the scheme's `RESTART_ATTRIBUTES`, and the ensemble's
+  state, prior, forecast, scaling and random stream, so a resumed run
+  continues the interrupted one exactly. Driven by `restart`, `restartsave`,
+  `restart_file`.
+- Every draw comes from `ensemble.rng`: a private `RandomState` when the
+  config gives a `seed`, otherwise NumPy's global stream as before.
+  `misc.sampling.gen_real` is the Gaussian sampler.
+- One named logger per log file (`ensemble.logger.PetLogger`); `NullLogger`
+  stands in when logging is off.
+
+## popt
+
+`popt.optimization_methods.OptimizerBase` has the same shape as the scheme
+base: it owns the loop, the starting evaluation, the callback, result
+recording, the log and the function, state and projected-gradient checks. An
+optimizer implements `update_step()`, committing an improving point with
+`_commit_step(x, f, jac=..., hess=...)` and returning a `StepReport(accepted,
+message)`, and `log_columns()`. `EnOpt`, `LineSearch`, `TrustRegion` and
+`SmcOpt` are exported from `popt`, as are the ensembles that estimate
+gradients and Hessians (`GaussianEnsemble`, `GeneralizedEnsemble`).
+
+## Where a new piece goes
+
+| Adding | Do |
+| --- | --- |
+| a scheme | subclass `AssimilationScheme` (or `IterativeEnRML`), implement `update_step`/`score`/`log_columns`, declare `COMPATIBLE_ANALYSES`, `register_scheme` |
+| an analysis | subclass `AnalysisBase`, implement `update` returning `AnalysisResult`, `register_analysis` |
+| a localization | a builder taking `info` (and `rng`, `data`, ...), `register_localization` |
+| an optimizer | subclass `OptimizerBase`, implement `update_step` with `_commit_step`, `log_columns` |
+| a simulator | a class with `input_dict` and `run_fwd_sim`; see the protocol's docstring for the optional hooks |
+| a config key | read it from the normalised section; add it to `KNOWN_DATAASSIM`/`KNOWN_ENSEMBLE` in `input_output.config` and to the [configuration reference](configuration.md) |
+
+The two notebooks under *Extending PIPT* in the tutorials walk through the
+first two.