diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..687a2d5 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,159 @@ +# Tests + +```sh +pytest # everything +pytest -m "not slow" # without the end-to-end runs with a parallel forecast +pytest tests/assimilation/test_numerical_characterisation.py +``` + +The MiniRes tests are skipped unless the `minires` extra is installed (`pip install -e .[minires]`). +Every test starts in its own temporary directory (`conftest.py`), since PET writes its output to +the working directory. + +Most of the suite is unit and regression tests of one module or one fixed bug; their docstrings say +what each pins. This page is about the rest: the **cases**, i.e. the synthetic problems that the +methods are run on end to end, and which **methods** run on which case. + +## Layout + +| | | +|---|---| +| `assimilation/` | `pipt`: the update schemes, their analyses, localization, the ensemble, data reading. | +| `optimization/` | `popt`: the optimizers, their ensembles and gradients, bounds, constraints. | +| the top level | What both share: data structures, config parsing, simulators, the CLI, imports, logging. | + +## The cases + +### Assimilation + +#### Van der Pol + +The workhorse. The oscillator `x'' - μ(1 - x²)x' + x = 0` (`simulator/vanderpol.py`), with the +initial conditions `x1`, `x2` and the damping `μ` as the unknowns (truth `1, 0, 1`; prior means +`0.05, 0.05, 1.5`). `x1` is observed at 15 steps with noise σ = 0.1. The simulator also +integrates the sensitivities, so it provides adjoints. + +It comes in two sizes: + +- **1000 members** (`test_assimilation_pipeline.py`, marked `slow`): asserts *properties*, i.e. + that the data misfit falls below a threshold and that the posterior mean of `μ` is within 20 % + of the prior's error. Catches a method that is broken, not one that has changed. +- **100 members, seeded, serial** (`test_numerical_characterisation.py`): asserts *numbers*. The + posterior ensemble and the misfit trajectory are compared to `characterisation_reference.npz` + at `rtol = 1e-9`. Catches any change to the numerics, intended or not. If a change is intended, + regenerate with `python tests/assimilation/test_numerical_characterisation.py --regenerate` + and check the diff; a mismatch can also come from a new numpy or BLAS rather than from the code. + +The small case is reused (via `_write_synthetic_case` and `_write_config`) wherever a real run is +needed but its quality is not the point: restarts, seeding, the forecast backends, QA/QC, the +legacy scheme names, and the multilevel scheme (2 levels × 10 members). + +#### Statistical consistency (vs. known posterior), not just previous/pinned values + +`test_statistical_consistency.py`, a linear sanity check. The model is `d = G x + noise`, with a +Gaussian prior and Gaussian noise, so the correct posterior is known exactly. Every registered +scheme and analysis is run on it and compared with that posterior, on both the mean and the spread. + +The cases are multivariate but not spatial: 1, 2 or 10 unknowns, a random prior correlation, a +dense random `G` with fewer data than unknowns, and noise variances spanning three orders of +magnitude. The reference is the exact update of the prior *ensemble*'s own mean and covariance, so +prior sampling error does not count against a method. The error is averaged over 5 seeds of 100 +members each. Correct methods come within 0.14 posterior sd on the mean and 0.90–1.11 on the +variance ratio, and the tolerances (0.2, and 0.85–1.2) leave a margin above that. Known +deviations are listed, with their reasons, in `DEVIATIONS` (strict xfails, so a fix shows up) +and `EXCLUDED`. + +#### Linear 1D field + +`simulator.simple_models.lin_1d`: the state is a 1D field, 150 cells in a row, with a spherical +variogram of range 20, and is observed directly at every fifth cell. The model being the identity, the +case isolates the analysis from any nonlinearity. + +- `test_linear_model.py`: LM-EnRML, `full` analysis, 250 members; the posterior mean at six cells + is pinned. +- `test_enif.py`: one cell, prior N(0, 1), observed at 1 with variance 0.25, 400 members. The + posterior is known exactly, N(0.8, 0.2), so ES-MDA with the `enif` analysis is checked against + it, rather than against stored numbers. + +#### MiniRes + +A two-phase waterflood (`simulator/minires.py`, wrapping [MiniRes](https://github.com/patnr/MiniRes)), +the only case resembling a reservoir. + +- `test_minires_sim.py`: the wrapper on an 8 × 8 grid: the simulator protocol, pickling, the state + as log-permeability in the declared grid order, crashed members, and the adjoint. +- `test_benchmark_minires.py` (marked `slow`): runs every method of the + [benchmark](../benchmarks/minires/README.md) on one tiny twin experiment (coarse grid, 6 members) + and asserts that the scores are finite and the misfit did not increase. It checks that the + benchmark still runs, not how well the methods do; that is `benchmarks/minires/RESULTS.md`. + +### Optimization + +#### Rosenbrock and quadratic + +The Rosenbrock function (`scipy.optimize.rosen`), mostly in 2D from about `(-1, -1)` with bounds +`[-2, 2]`, and the quadratic in `popt/cost_functions`. Given the exact gradient and Hessian, the +optimizers must reach the minimum `(1, 1)` to `1e-4`, with and without the transform of the bounds. +Given only an ensemble of the controls (Gaussian or generalized), they estimate the gradient, and +the tolerance is looser; the hardest such case is a 100-dimensional Rosenbrock, from `-2` +everywhere, with 100 members. + +### Hand-built cases + +Many tests use no simulator at all: a fake ensemble or scheme, or frames built by hand, so that one +piece runs without the rest. These are local to their file; e.g. `test_compression.py` (seismic +vintages over a masked grid), `test_distance_loc.py` (a small grid with point observations), and +`test_scheme_base.py` (the loop with a fake ensemble). + +## The methods + +### Update schemes (`pipt`) + +A scheme is an algorithm (ES, ES-MDA, …) and an *analysis*, the linear algebra of its update +(`approx`, `full`, `subspace`, …). `R` = pinned by the characterisation reference, +`P` = property test at 1000 members, `M` = MiniRes benchmark, `L` = matches the exact posterior in +the linear sanity checks, `L✗` = a known deviation there in some or all dimensions. + +| Scheme | `approx` | `full` | `subspace` | `subspace2` | `margis` | `enif` | `hybrid` | +|---|---|---|---|---|---|---|---| +| ES | R M L | R L | L ¹ | | | | | +| EnKF | R L | L | L ¹ | | | | | +| ES-MDA | R P M L | R M L✗ | R M L✗ | R M L✗ | | M L✗ ² | ³ | +| LM-EnRML | R P M ⁴ | R L ⁵ | R L✗ | R L | | | | +| GN-EnRML | R P M ⁴ | R L | R L | R M L | R M L | | | + +1. Crashes on the Van der Pol case, so it is not in the characterisation reference. +2. Exact with 1 and 2 unknowns; about 0.2 sd off the mean with 10, at 100 members. Also + checked against its exact posterior in `test_enif.py`. +3. The multilevel scheme (`test_multilevel.py`): runs, keeps its levels, and lowers the misfit; + no reference numbers. +4. Not in the linear sanity checks: `approx` drops the prior increment term, so the iterative + schemes converge to a different point by design. +5. Also pinned on the linear 1D field (`test_linear_model.py`). + +Further: + +- `subspace2` is checked against an identity: it is `margis` with the error scale fixed at 1 + (`test_subspace2.py`). +- `co_lm_enrml` and `gn_enrml` are old names for LM-EnRML/`approx` and GN-EnRML/`subspace`, and must + give exactly their numbers (`test_legacy_scheme_names.py`). +- Localization: adaptive and distance-based localization are unit-tested + (`test_autoadaloc.py`, `test_distance_loc.py`) and run end to end only in the MiniRes benchmark. + +### Optimizers (`popt`) + +| Optimizer | Case | File | +|---|---|---| +| EnOpt | quadratic, from a Gaussian ensemble | `test_ensemble_optimization.py` | +| LineSearch: GD, BFGS, Newton-CG | Rosenbrock, exact derivatives; quadratic and 100-D Rosenbrock, ensemble gradient | `test_line_search.py`, `test_ensemble_optimization.py` | +| TrustRegion: iterative, CG-Steihaug, BFGS | Rosenbrock, exact derivatives | `test_trust_region.py` | +| GenOpt, with and without CMA | Rosenbrock, generalized ensemble | `test_genopt.py` | +| EPF (constraints) | the outer loop's convergence | `test_epf_convergence.py` | + +The ensemble gradients and Hessians themselves are checked in `test_ensembles.py`, robust +optimization over several models in `test_robust_evaluation.py`, and a crashed simulation in +`test_crashed_evaluation.py`. + +## Not covered + +- The GIES schemes (`pipt/update_schemes/gies`) and SmcOpt (`popt/optimization_methods/smcopt.py`). diff --git a/tests/assimilation/test_statistical_consistency.py b/tests/assimilation/test_statistical_consistency.py new file mode 100644 index 0000000..3ef099d --- /dev/null +++ b/tests/assimilation/test_statistical_consistency.py @@ -0,0 +1,190 @@ +"""Statistical consistency: every method against the known posterior, not against pinned numbers. + +A linear sanity check. The model is linear, `d = G x + noise`, and the prior and the noise +are Gaussian, so the posterior is known exactly. Every registered (scheme, analysis) +combination that claims to sample it is run on the same problem and compared with it, on +both its mean and its spread. The characterisation tests pin today's output, wrong or not; +the Van der Pol tests check that the misfit falls; neither says whether the answer is +right. + +The cases are multivariate but not spatial: one state vector of 1, 2 or 10 unknowns, with +a random prior correlation and a dense random `G`, observed through fewer data than +unknowns (except in 1D), with noise variances spanning three orders of magnitude. +Nothing has a position, so localization is out of scope, and EnIF is given a complete +graph (no assumed sparsity) rather than the nearest-neighbour one it builds from a grid. + +The reference is the exact update of the prior *ensemble*: its own mean and covariance, +through the Kalman gain. That is the right answer for those members, so sampling error in +the prior does not count against a method; what is left (the perturbed observations) is +averaged out over a few seeds. With 100 members and five seeds, the correct methods were +measured within 0.14 posterior standard deviations on the mean and 0.90--1.11 on the +variance ratio, over four independent sets of seeds; the tolerances leave a margin above that. + +A method passes if it does so in every dimension: a case that happens not to show a fault +does not make the method right there. Expected deviations are listed in ``DEVIATIONS``, +per method, each with its reason. They are strict xfails: when one is fixed, its test +fails until the entry is removed. Those in ``BORDERLINE`` are not, since they deviate by +about as much as the tolerance, and so pass on some seeds. +""" + +import networkx +import numpy as np +import pandas as pd +import pytest + +from misc.structures import PETDataFrame +from pipt.update_schemes.registry import available_schemes, get_scheme + +NE = 100 +SEEDS = range(5) + +#: Bounds on the seed-averaged error: of the mean, in posterior standard deviations, and +#: of the variance, as a ratio to the posterior variance. Per unknown, the worst counts. +MEAN_TOL = 0.2 +VAR_RATIO = (0.85, 1.2) + +#: Left out of the parametrization, and why. +EXCLUDED = { + # Aliases; test_legacy_scheme_names checks they are exactly what they name. + ("co_lm_enrml", "approx"): "alias", + ("gn_enrml", "subspace"): "alias", + # Multilevel ES-MDA needs a multilevel ensemble; test_multilevel runs it. + ("esmda", "hybrid"): "multilevel", + # `approx` drops the prior increment term (Chen & Oliver 2013). For ES, EnKF and ES-MDA + # that term is zero, but the iterative schemes converge to a different point: with + # 2000 members, 1.14 x the posterior variance in 1D and 0.4 sd off the mean in 10D. + ("lmenrml", "approx"): "approximation", + ("gnenrml", "approx"): "approximation", +} + +#: Expected deviations from the posterior, per (scheme, analysis). +DEVIATIONS = { + ("esmda", "full"): "Exact with one MDA step, not with four: 1.6 x the posterior variance in 1D", + **{ + ("esmda", analysis): "Exact with one MDA step, not with four: up to 3 x the posterior " + "variance in 1D; suspected to re-update the prior at every step" + for analysis in ("subspace", "subspace2") + }, + ("lmenrml", "subspace"): "Can stop at about the prior: the first step, at the default " + "lambda, barely moves the misfit, which reads as convergence", + ("esmda", "enif"): "About 0.2 sd off the mean in 10D at 100 members, though exact at 2000: " + "presumably its regularised estimators", +} + +#: Deviations about the size of the tolerance: non-strict, so a change in the noise does not fail them. +BORDERLINE = {("esmda", "enif")} + +DIMENSIONS = (1, 2, 10) + +COMBOS = [c for c in available_schemes() if c not in EXCLUDED] + + +class Linear: + """`d = G x`: one datum per index, all of type `value`.""" + + def __init__(self, G): + self.G = G + index = list(range(G.shape[0])) + self.input_dict = {"reporttype": "index", "reportpoint": index, "datatype": ["value"], "parallel": 1} + self.true_order = ["index", index] + self.datatype = ["value"] + + def run_fwd_sim(self, state, member_index, *args, **kwargs): + return [{"value": np.array([y])} for y in self.G @ state["x"]] + + +def problem(nx): + """The prior, the model, the noise and the observations, for `nx` unknowns.""" + rng = np.random.default_rng(nx) + if nx == 1: + G, R = np.array([[2.0]]), np.array([0.5]) + elif nx == 2: + G, R = np.array([[1.0, 1.0]]), np.array([0.1]) + else: + G, R = rng.standard_normal((4, nx)), np.array([0.01, 0.1, 1.0, 10.0]) + A = rng.standard_normal((nx, nx + 2)) + S = A @ A.T + sd = np.linspace(1, 3, nx) + C = sd[:, None] * S / np.sqrt(np.outer(np.diag(S), np.diag(S))) * sd[None] + m = np.linspace(-1, 1, nx) + truth = rng.multivariate_normal(m, C) + d = G @ truth + np.sqrt(R) * rng.standard_normal(len(R)) + return {"G": G, "R": R, "C": C, "m": m, "d": d} + + +def exact_update(X, G, R, d): + """The Kalman update of the ensemble's own mean and covariance.""" + m, C = X.mean(axis=1), np.atleast_2d(np.cov(X)) + K = C @ G.T @ np.linalg.inv(G @ C @ G.T + np.diag(R)) + return m + K @ (d - G @ m), C - K @ G @ C + + +def assimilate(scheme, analysis, P, seed): + """The posterior ensemble, and the prior ensemble it started from.""" + nx, nd = len(P["m"]), len(P["d"]) + prior = np.random.default_rng(1000 + seed).multivariate_normal(P["m"], P["C"], size=NE).T + np.savez("prior.npz", x=prior) + index = pd.Index(range(nd), name="index") + PETDataFrame(pd.DataFrame({"value": P["d"]}, index=index)).to_pickle("data.pkl") + PETDataFrame(pd.DataFrame({"value": [["abs", r] for r in P["R"]]}, index=index)).to_pickle("var.pkl") + + keys_da = { + "scheme": scheme, "analysis": analysis, "obsname": "index", "datatype": ["value"], + "truedataindex": list(range(nd)), + # EnKF assimilates one datum at a time; the others, all at once. + "assimindex": [[i] for i in range(nd)] if scheme == "enkf" else list(range(nd)), + "data": "data.pkl", "datavar": "var.pkl", "energy": 1.0, + "mda": {"tot_assim_steps": 4, "inflation_param": [4, 4, 4, 4]}, + "iteration": {"max_iter": 20, "energy": 1.0}, + "enif": {"parameter_graphs": {"x": networkx.complete_graph(nx)}}, + } + keys_en = { + "ne": NE, "state": ["x"], "importstate": "prior.npz", "seed": seed, "disable_tqdm": True, + # The variance only sets the state scaling, to which the update is invariant. + "prior_x": {"mean": P["m"].tolist(), "var": float(np.diag(P["C"]).mean()), "grid": [nx, 1]}, + } + run = get_scheme(scheme, analysis)(keys_da, keys_en, Linear(P["G"])) + run.run_assimilation() + return np.asarray(run.enX, dtype=float), prior + + +def deviation(scheme, analysis, nx): + """How the seed-averaged posterior misses the exact one in `nx` unknowns; `None` if it does not.""" + P = problem(nx) + mean_errors, var_ratios = [], [] + for seed in SEEDS: + posterior, prior = assimilate(scheme, analysis, P, seed) + mean, cov = exact_update(prior, P["G"], P["R"], P["d"]) + sd = np.sqrt(np.diag(cov)) + mean_errors.append((posterior.mean(axis=1) - mean) / sd) + var_ratios.append(posterior.var(axis=1, ddof=1) / sd**2) + mean_error = np.abs(np.mean(mean_errors, axis=0)).max() + var_ratio = np.mean(var_ratios, axis=0) + + misses = [] + if mean_error > MEAN_TOL: + misses.append(f"mean off by {mean_error:.2f} posterior sd") + if not VAR_RATIO[0] <= var_ratio.min() <= var_ratio.max() <= VAR_RATIO[1]: + misses.append(f"variance {var_ratio.min():.2f}..{var_ratio.max():.2f} x the posterior's") + return ", ".join(misses) or None + + +def cases(): + for scheme, analysis in COMBOS: + reason = DEVIATIONS.get((scheme, analysis)) + strict = (scheme, analysis) not in BORDERLINE + marks = [pytest.mark.xfail(reason=reason, strict=strict)] if reason else [] + yield pytest.param(scheme, analysis, marks=marks, id=f"{scheme}-{analysis}") + + +@pytest.mark.parametrize("scheme, analysis", list(cases())) +def test_the_posterior_is_the_exact_one(scheme, analysis): + misses = {nx: deviation(scheme, analysis, nx) for nx in DIMENSIONS} + misses = {nx: miss for nx, miss in misses.items() if miss} + assert not misses, "; ".join(f"{nx} unknowns: {miss}" for nx, miss in misses.items()) + + +def test_every_combination_is_checked_or_excluded_for_a_reason(): + assert set(EXCLUDED) <= set(available_schemes()) + assert set(DEVIATIONS) <= set(COMBOS) + assert BORDERLINE <= set(DEVIATIONS)