From e77086abe7f3fb80e5ea5143b9f074c001500e75 Mon Sep 17 00:00:00 2001 From: Kristian Fossum Date: Tue, 15 Sep 2026 06:33:02 +0000 Subject: [PATCH 1/2] feat!: add EnIF and EnIF-MDA analyses Use ERT's graphite-maps estimators with PET's MDA lifecycle and graph-aware parameter updates. Include EnIF in the standard installation, document configuration, and test numerical parity and assimilation. BREAKING CHANGE: PET requires Python 3.12 through 3.14 to support the required graphite-maps dependency. --- .github/workflows/tests.yml | 2 +- README.md | 4 + docs/tutorials/README.md | 1 + docs/tutorials/enif.md | 118 ++++++++++++++ pyproject.toml | 5 +- src/pipt/update_schemes/enif.py | 256 +++++++++++++++++++++++++++++++ tests/test_enif.py | 263 ++++++++++++++++++++++++++++++++ 7 files changed, 646 insertions(+), 3 deletions(-) create mode 100644 docs/tutorials/enif.md create mode 100644 src/pipt/update_schemes/enif.py create mode 100644 tests/test_enif.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 3a0cf634..fa907d8a 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -25,7 +25,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", 3.11, 3.12] + python-version: ["3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v2 diff --git a/README.md b/README.md index fcc6c21a..57786f86 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,9 @@ at NORCE Norwegian Research Centre AS. ## Installation +PET requires Python 3.12 through 3.14. The standard installation includes EnIF +and EnIF-MDA with their dependencies. + Before installing ensure you have python3 pre-requisites. On a Debian system run: ``` @@ -77,6 +80,7 @@ Some basic plotting functionality is provided [here](https://github.com/Python-E - 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) +- [EnIF and EnIF-MDA](docs/tutorials/enif.md): installation, analysis settings and parameter graphs. ## Suggested readings: diff --git a/docs/tutorials/README.md b/docs/tutorials/README.md index 415f3f75..c44c04aa 100644 --- a/docs/tutorials/README.md +++ b/docs/tutorials/README.md @@ -4,3 +4,4 @@ Here are some tutorials. - [`tutorial_pipt.ipynb`](pipt/tutorial_pipt): Tutorial for running PIPT - [`tutorial_pipt.ipynb`](popt/tutorial_popt): Tutorial for running POPT +- [EnIF and EnIF-MDA](enif.md): Information-filter analyses and parameter graphs diff --git a/docs/tutorials/enif.md b/docs/tutorials/enif.md new file mode 100644 index 00000000..54809dd0 --- /dev/null +++ b/docs/tutorials/enif.md @@ -0,0 +1,118 @@ +# Ensemble information filter (EnIF) + +PET provides the original, single-update EnIF and an EnIF-MDA variant. Both +use the sparse regression and precision estimation from ERT's +[`_enif_update.py`](https://github.com/equinor/ert/blob/main/src/ert/analysis/_enif_update.py) +through `graphite-maps`. + +## Installation + +Install PET from your checkout, including EnIF and its dependencies: + +```sh +python -m pip install -e . +``` + +PET requires Python 3.12 through 3.14, matching its `graphite-maps` dependency. + +## Select the analysis + +Keep your existing ensemble, observation and simulator settings. For one EnIF +update, set these entries in `dataassim`: + +```yaml +daalg: [enif, enif] +analysis: full +``` + +For EnIF-MDA, use: + +```yaml +daalg: [enif, enif] +analysis: mda +mda: + tot_assim_steps: 3 + inflation_param: [2, 4, 4] +``` + +The corresponding classes are `enif_full` and `enif_mda` in +`pipt.update_schemes.enif`. PET's `pipt_init.init_da` loads them through the +existing configuration interface. + +Both variants assimilate all selected `assimindex` entries together. Standard +EnIF performs one update with inflation 1; it does not need `mda` settings. +EnIF-MDA reruns the simulator and refits the regression and state precision +after each update. + +MDA requires positive, finite inflation factors satisfying +`sum(1 / alpha) = 1`. If you omit `inflation_param`, PET uses +`tot_assim_steps` for each factor. A scalar factor repeats across the schedule. +The schedule retains its original indexing on restart. + +## Parameter graphs + +EnIF estimates a separate prior precision block for each state in `idX`. +By default: + +- A state with `grid` metadata in `prior_` uses nearest-neighbour + connectivity. PET's prior parser converts `grid` to `nx`, `ny` and `nz`. +- A state without grid metadata uses independent graph nodes. + +The regular-grid ordering matches PET's layered prior generator: +`row = z * nx * ny + x * ny + y` (y varies fastest). For imported ensembles +with a different ordering, reduced active-cell arrays or irregular geometry, +provide a graph whose node numbers match the imported parameter rows. + +You can configure graphs and neighbourhood sizes under `enif`: + +```yaml +enif: + parameter_graphs: + perm: perm_graph.npz + neighbourhood_expansion: 2 + neighbor_propagation_order: 15 +``` + +Write a graph file as a symmetric sparse adjacency array with +`scipy.sparse.save_npz`. For example, for five parameters arranged in a chain: + +```python +import networkx as nx +from scipy import sparse + +graph = nx.path_graph(5) +sparse.save_npz('perm_graph.npz', nx.to_scipy_sparse_array(graph, format='csc')) +``` + +Python configurations can also supply NetworkX graphs or SciPy sparse +adjacency arrays directly in `parameter_graphs`. Use local node numbers +`0` through `number_of_parameter_rows - 1` for each state. Graph weights do +not affect the fit; EnIF uses connectivity. + +EnIF excludes rows containing non-finite values and rows with zero ensemble +spread from estimation. It removes their graph nodes without connecting +neighbours across the resulting gaps. The analysis gives those rows a zero +increment, then applies PET's configured state limits. + +## Update and diagnostics + +EnIF uses PET's perturbed observations, random-number stream, state clipping, +forecast loop and misfit reporting. It scales observation covariance by the +current MDA factor once. It also estimates the unexplained response variance, +as in ERT. For a correlated observation covariance, it whitens the observations, +forecasts and perturbations before fitting the response map. + +The EnIF-specific update and helpers live in `pipt/update_schemes/enif.py`. +The scheme inherits PET's `esmdaMixIn` lifecycle and returns an additive +`step`, following the existing update-method interface. It uses the direct +sparse solver, matching ERT's non-iterative transport setting. + +For `analysisdebug`, the scheme exposes `H`, `Prec_u`, `Prec_eps` and +`Prec_posterior`. These matrices use standardized, retained state rows; +`enif_active_rows` maps them back to the full state. With correlated observation +errors, `H` and `Prec_eps` use whitened observation coordinates. + +This scheme requires at least two ensemble members and positive observation +variances. It does not support PET's covariance localization, local analysis, +multilevel ensembles or `emp_cov` sample input. Use parameter graphs to specify +spatial dependence. diff --git a/pyproject.toml b/pyproject.toml index 0dff01f7..e4b2829c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ maintainers = [ ] license = { file = "LICENSE.txt" } readme = "README.md" -requires-python = ">=3.8" +requires-python = ">=3.12,<3.15" dependencies = [ "numpy", "scipy", @@ -26,6 +26,7 @@ dependencies = [ "PyWavelets", "psutil", "geostat @ git+https://github.com/Python-Ensemble-Toolbox/Geostatistics@main", + "graphite-maps>=0.0.11,<0.1", "pytest", "pandas", "p_tqdm", @@ -59,4 +60,4 @@ Homepage = "https://github.com/Python-Ensemble-Toolbox/PET" package-dir = {"" = "src"} [tool.setuptools.packages.find] -where = ["src"] \ No newline at end of file +where = ["src"] diff --git a/src/pipt/update_schemes/enif.py b/src/pipt/update_schemes/enif.py new file mode 100644 index 00000000..aac15b99 --- /dev/null +++ b/src/pipt/update_schemes/enif.py @@ -0,0 +1,256 @@ +"""Ensemble information filter (EnIF) and multiple data assimilation schemes. + +The sparse regression and precision estimators follow ERT's EnIF analysis, +using ``graphite-maps``. PET supplies the perturbed observations and manages +forecasting, state limits, convergence diagnostics and acceptance of updates. +""" + +from os import PathLike + +import networkx as nx +import numpy as np +from graphite_maps.enif import EnIF +from graphite_maps.linear_regression import linear_boost_ic_regression +from graphite_maps.precision_estimation import fit_precision_cholesky_approximate +from scipy import linalg, sparse +from sklearn.preprocessing import StandardScaler + +from pipt.update_schemes.esmda import esmdaMixIn + + +class enif_update: + """Graph-informed information-space update, with PET's ``update`` interface.""" + + def update(self, enX, enY, enE, **kwargs): + """Compute ``self.step`` from the current ensemble and perturbed observations. + + Parameters + ---------- + enX : ndarray + State ensemble, shape (number of parameters, number of members). + enY : ndarray + Forecast ensemble, shape (number of observations, number of members). + enE : ndarray + Perturbed observations with covariance ``alpha * cov_data`` and + the same shape as ``enY``. These are used without adding more noise. + + Notes + ----- + Each parameter group has its own precision block. Parameters containing + non-finite values, and parameters with no ensemble spread, are held fixed. + The regression and prior precision are refitted at every MDA step. + """ + if enX.ndim != 2 or enX.shape[1] < 2: + raise ValueError('EnIF requires at least two ensemble members.') + if enY.ndim != 2 or enY.shape[1] != enX.shape[1] or enE.shape != enY.shape: + raise ValueError('EnIF state, forecast and observation ensembles have incompatible shapes.') + if enY.shape[0] == 0 or self.vecObs.shape != (enY.shape[0],): + raise ValueError('EnIF requires observations matching the forecast rows.') + if not all(np.all(np.isfinite(value)) for value in (enY, enE, self.vecObs)): + raise ValueError('EnIF observations and forecasts must be finite.') + + finite = np.all(np.isfinite(enX), axis=1) + if not finite.any(): + raise ValueError('No finite parameter rows available for EnIF.') + active = finite.copy() + active[finite] = np.ptp(enX[finite], axis=1) > 0 + self.step = np.zeros(enX.shape, dtype=float) + self.enif_active_rows = np.flatnonzero(active) + if not active.any(): + return + + scaler = StandardScaler() + U = scaler.fit_transform(enX[active].T) + Y, E, d, self.Prec_eps = self._observation_precision(enY, enE) + self.H = linear_boost_ic_regression(U=U, Y=Y.T) + + # Keep precision blocks in the same row order as the augmented state. + blocks = [] + for name, (start, stop) in sorted(self.idX.items(), key=lambda item: item[1][0]): + local_active = active[start:stop] + if not local_active.any(): + continue + graph = self._parameter_graph(name, stop - start) + graph = graph.subgraph(np.flatnonzero(local_active)) + graph = nx.convert_node_labels_to_integers(graph, ordering='sorted') + local_scaler = StandardScaler() + local_U = local_scaler.fit_transform(enX[start:stop][local_active].T) + blocks.append(fit_precision_cholesky_approximate( + local_U, + graph, + neighbourhood_expansion=self.enif_options.get('neighbourhood_expansion', 2), + use_tqdm=not self.disable_tqdm, + )) + self.Prec_u = sparse.csc_array(sparse.block_diag(blocks, format='csc')) + + gtmap = EnIF(Prec_u=self.Prec_u, Prec_eps=self.Prec_eps, H=self.H) + self.update_indices = gtmap.get_update_indices( + neighbor_propagation_order=self.enif_options.get('neighbor_propagation_order', 15), + ) + canonical = gtmap.pushforward_to_canonical(U) + residuals = gtmap.response_residual(U, Y.T) + # ERT transport draws noise internally. Use PET's existing perturbations + # instead: d - (residuals + d - E) == E - residuals. + canonical = gtmap.update_canonical( + canonical=canonical, + residual_noisy=residuals + d - E.T, + d=d, + ) + updated = gtmap.pullback_from_canonical( + updated_canonical=canonical, + update_indices=self.update_indices, + U_prior=U, + iterative=False, + ) + self.Prec_posterior = gtmap.Prec_u + self.step[active] = scaler.inverse_transform(updated).T - enX[active] + + def _parameter_graph(self, name, size): + """Load a group graph or build nearest-neighbour connectivity from its grid. + + Graph nodes are local parameter rows, numbered ``0 .. size-1``. Regular + grids use y-fastest ordering, then x, then z, matching PET's layered + prior ensembles. Without grid metadata, parameters are independent. + """ + graph = self.enif_options.get('parameter_graphs', {}).get(name) + if graph is not None: + if isinstance(graph, (str, PathLike)): + graph = sparse.load_npz(graph) + if sparse.issparse(graph): + if graph.shape != (size, size) or (graph != graph.T).nnz: + raise ValueError(f'EnIF graph for {name} must be a symmetric ({size}, {size}) adjacency.') + graph = nx.from_scipy_sparse_array(graph) + if not isinstance(graph, nx.Graph) or graph.is_directed() or graph.is_multigraph(): + raise ValueError(f'EnIF graph for {name} must be an undirected simple graph.') + if set(graph.nodes) != set(range(size)): + raise ValueError(f'EnIF graph for {name} must have nodes 0 through {size - 1}.') + return graph.copy() + + info = self.prior_info[name] + if not all(key in info for key in ('nx', 'ny')): + return nx.empty_graph(size) + shape = (int(info.get('nz', 1)), int(info['nx']), int(info['ny'])) + if min(shape) < 1 or np.prod(shape) != size: + raise ValueError( + f'EnIF grid for {name} has {np.prod(shape)} cells but {size} parameter rows. ' + 'Provide parameter_graphs for a reduced or irregular grid.' + ) + cells = np.arange(size).reshape(shape) + graph = nx.empty_graph(size) + for axis in range(3): + left = [slice(None)] * 3 + right = [slice(None)] * 3 + left[axis] = slice(None, -1) + right[axis] = slice(1, None) + graph.add_edges_from(zip(cells[tuple(left)].ravel(), cells[tuple(right)].ravel())) + return graph + + def _observation_precision(self, enY, enE): + """Inflate observation covariance once; whiten correlated observation errors.""" + covariance = np.asarray(self.cov_data, dtype=float) + alpha = self.alpha[self.iteration - 1] + nd = enY.shape[0] + if not np.all(np.isfinite(covariance)): + raise ValueError('EnIF observation covariance must be finite.') + if covariance.ndim == 2: + if covariance.shape != (nd, nd) or not np.allclose(covariance, covariance.T): + raise ValueError('EnIF observation covariance must be square and symmetric.') + if np.count_nonzero(covariance - np.diag(covariance.diagonal())): + chol = linalg.cholesky(covariance, lower=True) + Y = linalg.solve_triangular(chol, enY, lower=True) + E = linalg.solve_triangular(chol, enE, lower=True) + d = linalg.solve_triangular(chol, self.vecObs, lower=True) + precision = sparse.diags_array(np.full(nd, 1.0 / alpha), format='csc') + return Y, E, d, precision + covariance = covariance.diagonal() + if covariance.shape != (nd,) or np.any(covariance <= 0): + raise ValueError('EnIF requires one strictly positive observation variance per forecast row.') + precision = sparse.diags_array(1.0 / (alpha * covariance), format='csc') + return enY, enE, self.vecObs, precision + + +class enifMixIn(esmdaMixIn): + """Use PET's MDA lifecycle with EnIF-specific settings and inflation validation. + + Parameters + ---------- + keys_da : dict + Standard PET assimilation settings. Optional ``enif`` dictionary: + + - ``parameter_graphs``: maps state names to NetworkX graphs, sparse + adjacency arrays, or files written with ``scipy.sparse.save_npz``. + - ``neighbourhood_expansion``: precision fitting graph hops (default 2). + - ``neighbor_propagation_order``: update propagation hops (default 15). + + Covariance localization and local analysis cannot be combined with + this scheme; spatial dependence is specified by the parameter graphs. + keys_en : dict + Standard PET ensemble settings, including ``disable_tqdm``. + sim : object + PET forward simulator. + """ + + def __init__(self, keys_da, keys_en, sim): + for key in ('localization', 'localanalysis', 'multilevel'): + if key in keys_da or key in keys_en: + raise ValueError(f'EnIF does not support {key}.') + if keys_da.get('emp_cov') == 'yes': + raise ValueError('EnIF requires observation variances, not emp_cov samples.') + super().__init__(keys_da, keys_en, sim) + self.enif_options = self.keys_da.get('enif', {}) + if not isinstance(self.enif_options, dict): + raise ValueError('ENIF settings must be a dictionary.') + for key, minimum in (('neighbourhood_expansion', 1), ('neighbor_propagation_order', 0)): + value = self.enif_options.get(key, minimum) + if not isinstance(value, (int, np.integer)) or isinstance(value, bool) or value < minimum: + raise ValueError(f'EnIF {key} must be an integer >= {minimum}.') + graphs = self.enif_options.get('parameter_graphs', {}) + if not isinstance(graphs, dict) or set(graphs) - set(self.idX): + raise ValueError('EnIF parameter_graphs must map known state names to graphs.') + + def _mda_options(self): + """Accept PET's dictionary and legacy list forms for MDA settings.""" + if 'mda' not in self.keys_da: + raise ValueError('EnIF-MDA requires MDA settings with tot_assim_steps.') + options = self.keys_da['mda'] + try: + return dict(options) + except (TypeError, ValueError): + return dict([options]) + + def _ext_assim_steps(self): + """Keep the full schedule on restart; PET resumes using ``iteration``.""" + steps = self._mda_options().get('tot_assim_steps') + if not isinstance(steps, (int, float, np.integer)) or isinstance(steps, bool): + raise ValueError('MDA tot_assim_steps must be a positive integer.') + if not np.isfinite(steps) or steps < 1 or int(steps) != steps: + raise ValueError('MDA tot_assim_steps must be a positive integer.') + return list(range(int(steps))) + + def _ext_inflation_param(self): + """Validate a positive MDA schedule whose reciprocal factors sum to one.""" + count = len(self._ext_assim_steps()) + alpha = np.asarray(self._mda_options().get('inflation_param', count), dtype=float) + if alpha.ndim == 0: + alpha = np.full(count, alpha.item()) + if alpha.shape != (count,) or not np.all(np.isfinite(alpha)) or np.any(alpha <= 0): + raise ValueError('MDA requires one finite positive inflation factor per assimilation step.') + if not np.isclose(np.sum(1.0 / alpha), 1.0, rtol=1e-12, atol=1e-12): + raise ValueError('The inverse MDA inflation factors must sum to one.') + return alpha.tolist() + + +class enif_full(enifMixIn, enif_update): + """Original, single-update EnIF: ``daalg=['enif', 'enif'], analysis='full'``.""" + + def _ext_assim_steps(self): + return [0] + + def _ext_inflation_param(self): + return [1.0] + + +class enif_mda(enifMixIn, enif_update): + """EnIF-MDA: ``daalg=['enif', 'enif'], analysis='mda'`` with PET's ``mda`` settings.""" + + pass diff --git a/tests/test_enif.py b/tests/test_enif.py new file mode 100644 index 00000000..4d6c842f --- /dev/null +++ b/tests/test_enif.py @@ -0,0 +1,263 @@ +"""Numerical and PET lifecycle tests for the EnIF schemes.""" + +from copy import deepcopy + +import networkx as nx +import numpy as np +import pytest +from graphite_maps.enif import EnIF +from graphite_maps.linear_regression import linear_boost_ic_regression +from graphite_maps.precision_estimation import fit_precision_cholesky_approximate +from scipy import sparse +from sklearn.preprocessing import StandardScaler + +from pipt.loop.assimilation import Assimilate +from pipt.pipt_init import init_da +from pipt.update_schemes.enif import enif_full, enif_mda, enif_update +from simulator.simple_models import lin_1d + + +class LinearModel(lin_1d): + """Return independent response buffers for PET's serial simulation loop.""" + + def run_fwd_sim(self, state, member_i): + return deepcopy(super().run_fwd_sim(state, member_i)) + + +@pytest.fixture(autouse=True) +def preserve_random_state(): + state = np.random.get_state() + yield + np.random.set_state(state) + + +@pytest.fixture +def update_scheme(): + scheme = enif_update() + scheme.idX = {'field': (0, 6)} + scheme.prior_info = {'field': {'nx': 3, 'ny': 2, 'nz': 1}} + scheme.enif_options = {} + scheme.disable_tqdm = True + scheme.alpha = [1.0] + scheme.iteration = 1 + scheme.vecObs = np.array([1.2, -0.3]) + scheme.cov_data = np.array([0.2, 0.5]) + return scheme + + +@pytest.mark.parametrize('alpha', [1.0, 4.0]) +def test_matches_ert_transport(update_scheme, alpha): + """Compare the PET step with ERT's fit-and-transport recipe, member by member.""" + rng = np.random.default_rng(13) + X = rng.normal(size=(6, 80)) * np.arange(1, 7)[:, None] + 5 + Y = np.vstack((X[0] + 0.3 * X[1] ** 2, X[4] - X[5])) + graph = nx.grid_2d_graph(3, 2) + graph = nx.convert_node_labels_to_integers(graph) + scaler = StandardScaler() + U = scaler.fit_transform(X.T) + H = linear_boost_ic_regression(U=U, Y=Y.T) + precision = fit_precision_cholesky_approximate(U, graph, use_tqdm=False) + reference = EnIF( + Prec_u=precision, + Prec_eps=sparse.diags_array(1 / (alpha * update_scheme.cov_data), format='csc'), + H=H, + ) + noise = reference.generate_observation_noise(X.shape[1], seed=19) + expected = reference.transport( + U, Y.T, update_scheme.vecObs, + update_indices=reference.get_update_indices(neighbor_propagation_order=15), + iterative=False, seed=19, + ) + expected = scaler.inverse_transform(expected).T + + update_scheme.alpha = [alpha] + E = update_scheme.vecObs[:, None] - noise.T + update_scheme.update(X, Y, E) + + np.testing.assert_allclose(X + update_scheme.step, expected, rtol=1e-11, atol=1e-11) + np.testing.assert_allclose(update_scheme.Prec_posterior.toarray(), reference.Prec_u.toarray()) + + +def test_parameter_grid_order_and_custom_graphs(update_scheme, tmp_path): + update_scheme.prior_info['field']['nz'] = 2 + graph = update_scheme._parameter_graph('field', 12) + assert set(graph.neighbors(0)) == {1, 2, 6} + assert set(graph.neighbors(5)) == {3, 4, 11} + assert graph.number_of_edges() == 20 + + custom = nx.path_graph(6) + filename = tmp_path / 'graph.npz' + sparse.save_npz(filename, nx.to_scipy_sparse_array(custom)) + update_scheme.enif_options['parameter_graphs'] = {'field': filename} + assert set(update_scheme._parameter_graph('field', 6).edges) == set(custom.edges) + + update_scheme.enif_options['parameter_graphs']['field'] = nx.empty_graph(6) + assert update_scheme._parameter_graph('field', 6).number_of_edges() == 0 + update_scheme.enif_options['parameter_graphs']['field'] = nx.path_graph(5) + with pytest.raises(ValueError, match='nodes 0 through 5'): + update_scheme._parameter_graph('field', 6) + update_scheme.enif_options = {} + with pytest.raises(ValueError, match='12 cells but 6 parameter rows'): + update_scheme._parameter_graph('field', 6) + + +def test_masks_and_group_precision(update_scheme): + rng = np.random.default_rng(5) + X = rng.normal(size=(6, 60)) + X[1] = np.nan + X[3, 0] = np.nan + X[4] = 2.0 + Y = np.vstack((X[0], X[5])) + update_scheme.idX = {'other': (5, 6), 'field': (0, 5)} + update_scheme.prior_info = {'field': {'nx': 5, 'ny': 1}, 'other': {}} + update_scheme.update(X, Y, np.tile(update_scheme.vecObs[:, None], (1, X.shape[1]))) + + np.testing.assert_array_equal(update_scheme.enif_active_rows, [0, 2, 5]) + np.testing.assert_array_equal(update_scheme.step[[1, 3, 4]], 0) + assert np.isfinite(update_scheme.step).all() + assert np.linalg.norm(update_scheme.step[[0, 5]]) > 0 + # Removing inactive nodes must not bridge across the hole in the field. + assert update_scheme.Prec_u[0, 1] == 0 + assert update_scheme.Prec_u[:2, 2:].nnz == 0 + + +def test_correlated_observations(update_scheme): + """Whitened EnIF agrees with a Gaussian update using the full covariance.""" + rng = np.random.default_rng(17) + X = rng.normal(size=(6, 400)) + X -= X.mean(axis=1, keepdims=True) + X /= X.std(axis=1, keepdims=True) + Y = np.vstack((X[0], X[1])) + covariance = np.array([[0.4, 0.2], [0.2, 0.6]]) + update_scheme.cov_data = covariance + update_scheme.prior_info = {'field': {}} + E = np.tile(update_scheme.vecObs[:, None], (1, X.shape[1])) + update_scheme.update(X, Y, E) + + expected_mean = np.linalg.solve(np.eye(2) + covariance, update_scheme.vecObs) + np.testing.assert_allclose((X + update_scheme.step).mean(axis=1)[:2], expected_mean, atol=0.025) + + +@pytest.mark.parametrize('covariance', [np.array([0.0, 1.0]), np.array([-1.0, 1.0]), + np.array([np.nan, 1.0]), np.ones(3), + np.array([[1.0, 0.2], [0.0, 1.0]])]) +def test_invalid_observation_covariance(update_scheme, covariance): + update_scheme.cov_data = covariance + with pytest.raises(ValueError): + update_scheme._observation_precision(np.ones((2, 20)), np.ones((2, 20))) + + +def test_constant_and_nonfinite_parameters(update_scheme): + X = np.ones((6, 20)) + Y = np.ones((2, 20)) + update_scheme.update(X, Y, Y) + np.testing.assert_array_equal(update_scheme.step, 0) + with pytest.raises(ValueError, match='No finite parameter rows'): + update_scheme.update(X * np.nan, Y, Y) + with pytest.raises(ValueError, match='at least two ensemble members'): + update_scheme.update(X[:, :1], Y[:, :1], Y[:, :1]) + with pytest.raises(ValueError, match='must be finite'): + update_scheme.update(X, Y * np.nan, Y) + + +@pytest.fixture +def pet_inputs(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + rng = np.random.default_rng(42) + np.savez('prior.npz', field=rng.normal(size=(1, 400))) + keys_da = { + 'daalg': ['enif', 'enif'], 'analysis': 'full', + 'obsname': 'index', 'truedataindex': [0], 'assimindex': [0], + 'datatype': ['value'], 'truedata': [1.0], 'datavar': ['abs', 0.25], + } + keys_en = { + 'ne': 400, 'state': ['field'], 'staticvar': ['field'], + 'prior_field': {'mean': 0.0, 'var': 1.0}, + 'importstaticvar': 'prior.npz', 'disable_tqdm': True, + } + sim = LinearModel({'reporttype': 'index', 'reportpoint': [0], 'datatype': ['value']}) + return keys_da, keys_en, sim + + +@pytest.mark.parametrize('analysis, mda, steps', [ + ('full', None, 1), + ('mda', {'tot_assim_steps': 1}, 1), + ('mda', {'tot_assim_steps': 3}, 3), + ('mda', {'tot_assim_steps': 3, 'inflation_param': [2, 4, 4]}, 3), +]) +def test_pet_assimilation_loop(pet_inputs, analysis, mda, steps): + keys_da, keys_en, sim = pet_inputs + keys_da['analysis'] = analysis + if mda is not None: + keys_da['mda'] = mda + np.random.seed(21) + ensemble = init_da(keys_da, keys_en, sim) + prior = ensemble.enX.copy() + assimilation = Assimilate(ensemble) + assimilation.run() + + assert isinstance(ensemble, enif_full if analysis == 'full' else enif_mda) + assert ensemble.iteration == steps + 1 + assert ensemble.enX_temp is None + assert ensemble.data_misfit < ensemble.prior_data_misfit + np.testing.assert_array_equal(ensemble.prior_enX, prior) + # N(0, 1) prior observed at 1 with variance 0.25 has N(0.8, 0.2) posterior. + np.testing.assert_allclose(ensemble.enX.mean(), 0.8, atol=0.07) + np.testing.assert_allclose(ensemble.enX.var(), 0.2, atol=0.05) + posterior = np.load('SaveOutputs/posterior_state_estimate.npz')['field'] + np.testing.assert_array_equal(posterior, ensemble.enX) + np.testing.assert_allclose(ensemble.pred_data[0]['value'], ensemble.enX) + + +def test_one_step_mda_is_original_enif(pet_inputs): + results = [] + for analysis in ('full', 'mda'): + keys_da, keys_en, sim = deepcopy(pet_inputs) + keys_da.update(analysis=analysis, mda={'tot_assim_steps': 1}) + np.random.seed(22) + ensemble = init_da(keys_da, keys_en, sim) + ensemble.iteration = 1 + ensemble.pred_data = [{'value': ensemble.enX.copy()}] + ensemble.calc_analysis() + results.append(ensemble.enX_temp.copy()) + np.testing.assert_array_equal(*results) + + +def test_state_limits(pet_inputs): + keys_da, keys_en, sim = pet_inputs + keys_en['prior_field']['limits'] = [-0.1, 0.1] + ensemble = init_da(keys_da, keys_en, sim) + ensemble.iteration = 1 + ensemble.pred_data = [{'value': ensemble.enX.copy()}] + ensemble.calc_analysis() + assert np.min(ensemble.enX_temp) >= -0.1 + assert np.max(ensemble.enX_temp) <= 0.1 + + +@pytest.mark.parametrize('options', [ + {'tot_assim_steps': 0}, {'tot_assim_steps': 1.5}, + {'tot_assim_steps': 2, 'inflation_param': [2]}, + {'tot_assim_steps': 2, 'inflation_param': [1, 1]}, + {'tot_assim_steps': 2, 'inflation_param': [0, 2]}, + {'tot_assim_steps': 2, 'inflation_param': [np.inf, 1]}, + {'tot_assim_steps': 2, 'inflation_param': [-1, 0.5]}, +]) +def test_invalid_mda_schedule(options): + ensemble = enif_mda.__new__(enif_mda) + ensemble.keys_da = {'mda': options} + with pytest.raises(ValueError): + ensemble._ext_inflation_param() + + +@pytest.mark.parametrize('options', [ + {'tot_assim_steps': 3, 'inflation_param': [2, 4, 4]}, + [['tot_assim_steps', 3], ['inflation_param', [2, 4, 4]]], +]) +def test_restart_preserves_full_schedule(options): + ensemble = enif_mda.__new__(enif_mda) + ensemble.keys_da = {'mda': options} + ensemble.restart = True + ensemble.iteration = 2 + ensemble.loop_ind = 1 + assert ensemble._ext_assim_steps() == [0, 1, 2] + assert ensemble._ext_inflation_param() == [2, 4, 4] From ace17fca13c362beb86321c0a246e243dd9c22dd Mon Sep 17 00:00:00 2001 From: "Rolf J. Lorentzen" Date: Fri, 25 Sep 2026 12:58:05 +0200 Subject: [PATCH 2/2] Update pyproject.toml --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index c32f072e..803a86d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,6 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "PET" +version = "0.2" description = "Python Ensemble Toolbox" authors = [ { name = "Data assimilation and optimization group" }