From 6599b27d4c8bcf570242c4f514064503b3f202cf Mon Sep 17 00:00:00 2001 From: Tommy <10076072+tommyod@users.noreply.github.com> Date: Thu, 16 Nov 2023 11:03:04 +0100 Subject: [PATCH 1/7] Notebook with Gauss-Linear model (#171) * added linear regression (Gauss-Linear model) * added notebook .py file --------- Co-authored-by: Tommy Odland --- docs/source/LinearRegression.py | 172 ++++++++++++++++++++++++++++++++ docs/source/conf.py | 5 +- docs/source/examples.rst | 1 + 3 files changed, 175 insertions(+), 3 deletions(-) create mode 100644 docs/source/LinearRegression.py diff --git a/docs/source/LinearRegression.py b/docs/source/LinearRegression.py new file mode 100644 index 00000000..32a4b07d --- /dev/null +++ b/docs/source/LinearRegression.py @@ -0,0 +1,172 @@ +# --- +# jupyter: +# jupytext: +# formats: ipynb,py:percent +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.15.2 +# kernelspec: +# display_name: Python 3 (ipykernel) +# language: python +# name: python3 +# --- + +# %% +# ruff: noqa: E402 + +# %% [markdown] +# # Linear regression with ESMDA +# +# We solve a linear regression problem using ESMDA. +# First we define the forward model as $g(x) = Ax$, +# then we set up a prior ensemble on the linear +# regression coefficients, so $x \sim \mathcal{N}(0, 1)$. +# +# As shown in the 2013 paper by Emerick et al, when a set of +# inflation weights $\alpha_i$ is chosen so that $\sum_i \alpha_i^{-1} = 1$, +# ESMDA yields the correct posterior mean for the linear-Gaussian case. + +# %% [markdown] +# ## Import packages + +# %% +import numpy as np +from matplotlib import pyplot as plt + +from iterative_ensemble_smoother import ESMDA + +# %% [markdown] +# ## Create problem data +# +# Some settings worth experimenting with: +# +# - Decreasing `prior_std=1` will pull the posterior solution toward zero. +# - Increasing `num_ensemble` will increase the quality of the solution. +# - Increasing `num_observations / num_parameters` +# will increase the quality of the solution. + +# %% +num_parameters = 25 +num_observations = 100 +num_ensemble = 30 +prior_std = 1 + +# %% +rng = np.random.default_rng(42) + +# Create a problem with g(x) = A @ x +A = rng.standard_normal(size=(num_observations, num_parameters)) + + +def g(X): + """Forward model.""" + return A @ X + + +# Create observations: obs = g(x) + N(0, 1) +x_true = np.linspace(-1, 1, num=num_parameters) +observation_noise = rng.standard_normal(size=num_observations) +observations = g(x_true) + observation_noise + +# Initial ensemble X ~ N(0, prior_std) and diagonal covariance with ones +X = rng.normal(size=(num_parameters, num_ensemble)) * prior_std + +# Covariance matches the noise added to observations above +covariance = np.ones(num_observations) + +# %% [markdown] +# ## Solve the maximum likelihood problem +# +# We can solve $Ax = b$, where $b$ is the observations, +# for the maximum likelihood estimate. +# Notice that unlike using a Ridge model, +# solving $Ax = b$ directly does not use any prior information. + +# %% +x_ml, *_ = np.linalg.lstsq(A, observations, rcond=None) + +plt.figure(figsize=(8, 3)) +plt.scatter(np.arange(len(x_true)), x_true, label="True parameter values") +plt.scatter(np.arange(len(x_true)), x_ml, label="ML estimate (no prior)") +plt.xlabel("Parameter index") +plt.ylabel("Parameter value") +plt.grid(True, ls="--", zorder=0, alpha=0.33) +plt.legend() +plt.show() + +# %% [markdown] +# ## Solve using ESMDA +# +# We crease an `ESMDA` instance and solve the Guass-linear problem. + +# %% +smoother = ESMDA( + covariance=covariance, + observations=observations, + alpha=5, + seed=1, +) + +X_i = np.copy(X) +for i, alpha_i in enumerate(smoother.alpha, 1): + print( + f"ESMDA iteration {i}/{smoother.num_assimilations()}" + + f" with inflation factor alpha_i={alpha_i}" + ) + X_i = smoother.assimilate(X_i, Y=g(X_i)) + + +X_posterior = np.copy(X_i) + +# %% [markdown] +# ## Plot and compare solutions +# +# Compare the true parameters with both the ML estimate +# from linear regression and the posterior means obtained using `ESMDA`. + +# %% +plt.figure(figsize=(8, 3)) +plt.scatter(np.arange(len(x_true)), x_true, label="True parameter values") +plt.scatter(np.arange(len(x_true)), x_ml, label="ML estimate (no prior)") +plt.scatter( + np.arange(len(x_true)), np.mean(X_posterior, axis=1), label="Posterior mean" +) +plt.xlabel("Parameter index") +plt.ylabel("Parameter value") +plt.grid(True, ls="--", zorder=0, alpha=0.33) +plt.legend() +plt.show() + +# %% [markdown] +# We now include the posterior samples as well. + +# %% +plt.figure(figsize=(8, 3)) +plt.scatter(np.arange(len(x_true)), x_true, label="True parameter values") +plt.scatter(np.arange(len(x_true)), x_ml, label="ML estimate (no prior)") +plt.scatter( + np.arange(len(x_true)), np.mean(X_posterior, axis=1), label="Posterior mean" +) + +# Loop over every ensemble member and plot it +for j in range(num_ensemble): + # Jitter along the x-axis a little bit + x_jitter = np.arange(len(x_true)) + rng.normal(loc=0, scale=0.1, size=len(x_true)) + + # Plot this ensemble member + plt.scatter( + x_jitter, + X_posterior[:, j], + label=("Posterior values" if j == 0 else None), + color="black", + alpha=0.2, + s=5, + zorder=0, + ) +plt.xlabel("Parameter index") +plt.ylabel("Parameter value") +plt.grid(True, ls="--", zorder=0, alpha=0.33) +plt.legend() +plt.show() diff --git a/docs/source/conf.py b/docs/source/conf.py index eb42c062..30177164 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -20,11 +20,10 @@ version = re.sub(r"\.dev.*$", r".dev", ies.__version__) release = version -# convert the python file to a notebook +# Convert Python files to notebooks check_output(["jupytext", "Polynomial.py", "-o", "Polynomial.ipynb"]) - -# Do the same for this file check_output(["jupytext", "Oscillator.py", "-o", "Oscillator.ipynb"]) +check_output(["jupytext", "LinearRegression.py", "-o", "LinearRegression.ipynb"]) # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones. diff --git a/docs/source/examples.rst b/docs/source/examples.rst index d6506b3e..c043aee6 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -7,5 +7,6 @@ Example Usage Polynomial Oscillator + LinearRegression * :ref:`genindex` From 56ab349341ceb71dd35fad8fd3658d5784f1763a Mon Sep 17 00:00:00 2001 From: Tommy <10076072+tommyod@users.noreply.github.com> Date: Mon, 20 Nov 2023 11:51:12 +0100 Subject: [PATCH 2/7] Update README installation instructions (#172) * Update README.md with development install --------- Co-authored-by: Tommy Odland --- README.md | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 554b13c2..1bcd4f27 100644 --- a/README.md +++ b/README.md @@ -26,22 +26,20 @@ This algorithm is particularly effective for problems with a large number of par pip install iterative_ensemble_smoother ``` -## Usage - -**iterative_ensemble_smoother** mainly implements `SIES` and `ESMDA` classes. Check out -the examples section to see how to use them. - - -## Building from source - -To build **iterative_ensemble_smoother** from source: +If you want to do development, then run: -```bash +```text git clone https://github.com/equinor/iterative_ensemble_smoother.git cd iterative_ensemble_smoother -pip install . + +pip install --editable '.[doc,dev]' ``` +## Usage + +**iterative_ensemble_smoother** mainly implements the two classes `SIES` and `ESMDA`. +Check out the examples section to see how to use them. + ## Building the documentation ```bash From 1cc21ee6e3979fc1c85e4c0a287a9102a16d428a Mon Sep 17 00:00:00 2001 From: Tommy <10076072+tommyod@users.noreply.github.com> Date: Mon, 20 Nov 2023 13:16:33 +0100 Subject: [PATCH 3/7] added LITERATURE.md (#173) Co-authored-by: Tommy Odland --- LITERATURE.md | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 LITERATURE.md diff --git a/LITERATURE.md b/LITERATURE.md new file mode 100644 index 00000000..0c12bde9 --- /dev/null +++ b/LITERATURE.md @@ -0,0 +1,30 @@ +# Literature + +A list of literature that we find useful for Data Assimilation. +Could be moved to `docs/source` eventually. + +## Papers + +### SIES + +- (2019) [Efficient Implementation of an Iterative Ensemble Smoother for Data Assimilation and Reservoir History Matching](https://www.frontiersin.org/articles/10.3389/fams.2019.00047/full) by Evensen et al + +### ESMDA + +- (2012) [History matching time-lapse seismic data using the ensemble Kalman filter with multiple data assimilations](https://link.springer.com/article/10.1007/s10596-012-9275-5) by Emerick et al +- (2013) [Ensemble smoother with multiple data assimilation](https://www.sciencedirect.com/science/article/abs/pii/S0098300412000994) by Emerick et al + +### Covariance regularization + +- (2022) [GraphSPME: Markov Precision Matrix Estimation and Asymptotic Stein-Type Shrinkage](https://arxiv.org/abs/2205.07584) by Lunde et al + +## Books + +### Preliminaries + +- (2006) [Pattern Recognition and Machine Learning](https://www.amazon.com/Pattern-Recognition-Learning-Information-Statistics/dp/0387310738) by Bishop + +### Data Assimilation + +- (2014) [Data Assimilation: The Ensemble Kalman Filter](https://www.amazon.com/Data-Assimilation-Ensemble-Kalman-Filter/dp/3642424767/) by Evensen +- (2022) [Data Assimilation Fundamentals: A Unified Formulation of the State and Parameter Estimation Problem](https://www.amazon.com/Data-Assimilation-Fundamentals-Formulation-Environment/dp/3030967085/) by Evensen From 8a221fad9409a70bcae402db41a03630f8e9d935 Mon Sep 17 00:00:00 2001 From: Feda Curic Date: Thu, 23 Nov 2023 11:34:18 +0100 Subject: [PATCH 4/7] Fix readthedocs (#182) --- .readthedocs.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.readthedocs.yaml b/.readthedocs.yaml index d026bfd7..9e46c891 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -10,7 +10,7 @@ python: - method: pip path: . extra_requirements: - - docs + - doc sphinx: builder: html From fa74f6b8b33a56a6617d2cd261f84b1c527d0728 Mon Sep 17 00:00:00 2001 From: Feda Curic Date: Thu, 23 Nov 2023 11:38:42 +0100 Subject: [PATCH 5/7] Remove unused alpha from perturb_observations (#178) --- src/iterative_ensemble_smoother/esmda.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index 3f83cde5..de8ef9eb 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -227,7 +227,7 @@ def assimilate( # a zero cented normal means that y := L @ z, where z ~ norm(0, 1) # Therefore, scaling C_D by alpha is equivalent to scaling L with sqrt(alpha) size = (num_outputs, num_ensemble) - D = self.perturb_observations(size=size, alpha=self.alpha[self.iteration]) + D = self.perturb_observations(size=size) assert D.shape == (num_outputs, num_ensemble) # Line 2 (c) in the description of ES-MDA in the 2013 Emerick paper @@ -291,7 +291,7 @@ def compute_transition_matrix( # or # X += X @ K - D = self.perturb_observations(size=Y.shape, alpha=alpha) + D = self.perturb_observations(size=Y.shape) inversion_func = self._inversion_methods[self.inversion] return inversion_func( alpha=alpha, @@ -303,9 +303,7 @@ def compute_transition_matrix( return_K=True, # Ensures that we don't need X ) - def perturb_observations( - self, *, size: Tuple[int, int], alpha: float - ) -> npt.NDArray[np.double]: + def perturb_observations(self, *, size: Tuple[int, int]) -> npt.NDArray[np.double]: """Create a matrix D with perturbed observations. In the Emerick (2013) paper, the matrix D is defined in section 6. @@ -315,10 +313,6 @@ def perturb_observations( ---------- size : Tuple[int, int] The size, a tuple with (num_observations, ensemble_size). - alpha : float - The covariance inflation factor. The sequence of alphas should - obey the equation sum_i (1/alpha_i) = 1. However, this is NOT enforced - in this method call. The user/caller is responsible for this. Returns ------- From 6a48047e714de84c23442a78e5bc76fb323f7f80 Mon Sep 17 00:00:00 2001 From: Feda Curic Date: Thu, 23 Nov 2023 11:38:59 +0100 Subject: [PATCH 6/7] Remove redundant assert (#176) --- src/iterative_ensemble_smoother/esmda.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index de8ef9eb..5357f0a8 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -159,7 +159,6 @@ def __init__( raise TypeError("Argument `covariance` must be 1D or 2D array") self.C_D = covariance - assert isinstance(self.C_D, np.ndarray) and self.C_D.ndim in (1, 2) def num_assimilations(self) -> int: return len(self.alpha) From 31e9b5479f1790e0f67a85b86672576b7cadf634 Mon Sep 17 00:00:00 2001 From: Feda Curic Date: Thu, 23 Nov 2023 11:39:25 +0100 Subject: [PATCH 7/7] Use T for transition matrix instead of K (#181) K usually refers to Kalman gain --- src/iterative_ensemble_smoother/esmda.py | 26 ++++++++--------- .../esmda_inversion.py | 28 +++++++++---------- tests/test_esmda_inversion.py | 8 +++--- 3 files changed, 31 insertions(+), 31 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index 5357f0a8..70ca4f39 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -180,7 +180,7 @@ def assimilate( corresponds to a parameter in the model, and each column corresponds to an ensemble member (realization). Y : np.ndarray - 2D array of shape (num_parameters, ensemble_size), containing + 2D array of shape (num_observations, ensemble_size), containing responses when evaluating the model at X. In other words, Y = g(X), where g is the forward model. overwrite : bool @@ -195,7 +195,7 @@ def assimilate( Returns ------- X_posterior : np.ndarray - 2D array of shape (num_inputs, num_ensemble_members). + 2D array of shape (num_parameters, num_ensemble_members). """ if self.iteration >= self.num_assimilations(): @@ -208,9 +208,9 @@ def assimilate( num_ensemble == num_emsemble2 ), "Number of ensemble members in X and Y must match" if not np.issubdtype(X.dtype, np.floating): - raise TypeError("Argument `X` must be contain floats") + raise TypeError("Argument `X` must contain floats") if not np.issubdtype(Y.dtype, np.floating): - raise TypeError("Argument `Y` must be contain floats") + raise TypeError("Argument `Y` must contain floats") assert 0 < truncation <= 1.0 @@ -254,7 +254,7 @@ def compute_transition_matrix( alpha: float, truncation: float = 1.0, ) -> npt.NDArray[np.double]: - """Return a matrix K such that X_posterior = X_prior + X_prior @ K. + """Return a matrix T such that X_posterior = X_prior + X_prior @ T. The purpose of this method is to facilitate row-by-row, or batch-by-batch, updates of X. This is useful if X is too large to fit in memory. @@ -262,7 +262,7 @@ def compute_transition_matrix( Parameters ---------- Y : np.ndarray - 2D array of shape (num_parameters, ensemble_size), containing + 2D array of shape (num_observations, ensemble_size), containing responses when evaluating the model at X. In other words, Y = g(X), where g is the forward model. alpha : float @@ -276,19 +276,19 @@ def compute_transition_matrix( Returns ------- - K : np.ndarray - A matrix K such that X_posterior = X_prior + X_prior @ K. + T : np.ndarray + A matrix T such that X_posterior = X_prior + X_prior @ T. It has shape (num_ensemble_members, num_ensemble_members). """ # Recall the update equation: # X += C_MD @ (C_DD + alpha * C_D)^(-1) @ (D - Y) # X += X @ center(Y).T / (N-1) @ (C_DD + alpha * C_D)^(-1) @ (D - Y) - # We form K := center(Y).T / (N-1) @ (C_DD + alpha * C_D)^(-1) @ (D - Y), + # We form T := center(Y).T / (N-1) @ (C_DD + alpha * C_D)^(-1) @ (D - Y), # so that - # X_new = X_old + X_old @ K + # X_new = X_old + X_old @ T # or - # X += X @ K + # X += X @ T D = self.perturb_observations(size=Y.shape) inversion_func = self._inversion_methods[self.inversion] @@ -297,9 +297,9 @@ def compute_transition_matrix( C_D=self.C_D, D=D, Y=Y, - X=None, # We don't need X to compute the factor K + X=None, # We don't need X to compute the factor T truncation=truncation, - return_K=True, # Ensures that we don't need X + return_T=True, # Ensures that we don't need X ) def perturb_observations(self, *, size: Tuple[int, int]) -> npt.NDArray[np.double]: diff --git a/src/iterative_ensemble_smoother/esmda_inversion.py b/src/iterative_ensemble_smoother/esmda_inversion.py index 700efe9e..ca96e2e1 100644 --- a/src/iterative_ensemble_smoother/esmda_inversion.py +++ b/src/iterative_ensemble_smoother/esmda_inversion.py @@ -188,16 +188,16 @@ def inversion_exact_cholesky( Y: npt.NDArray[np.double], X: Optional[npt.NDArray[np.double]], truncation: float = 1.0, - return_K: bool = False, + return_T: bool = False, ) -> npt.NDArray[np.double]: """Computes an exact inversion using `sp.linalg.solve`, which uses a Cholesky factorization in the case of symmetric, positive definite matrices. The goal is to compute: C_MD @ inv(C_DD + alpha * C_D) @ (D - Y) - First we solve (C_DD + alpha * C_D) @ K = (D - Y) for K, so that - K = inv(C_DD + alpha * C_D) @ (D - Y), then we compute - C_MD @ K, but we don't explicitly form C_MD, since it might be more + First we solve (C_DD + alpha * C_D) @ T = (D - Y) for T, so that + T = inv(C_DD + alpha * C_D) @ (D - Y), then we compute + C_MD @ T, but we don't explicitly form C_MD, since it might be more efficient to perform the matrix products in another order. """ C_DD = empirical_covariance_upper(Y) # Only compute upper part @@ -210,25 +210,25 @@ def inversion_exact_cholesky( "lower": False, # Only use the upper part while solving } - # Compute K := sp.linalg.inv(C_DD + alpha * C_D) @ (D - Y) + # Compute T := sp.linalg.inv(C_DD + alpha * C_D) @ (D - Y) if C_D.ndim == 2: # C_D is a covariance matrix C_DD += alpha * C_D # Save memory by mutating - K = sp.linalg.solve(C_DD, D - Y, **solver_kwargs) + T = sp.linalg.solve(C_DD, D - Y, **solver_kwargs) elif C_D.ndim == 1: # C_D is an array, so add it to the diagonal without forming diag(C_D) C_DD.flat[:: C_DD.shape[1] + 1] += alpha * C_D - K = sp.linalg.solve(C_DD, D - Y, **solver_kwargs) + T = sp.linalg.solve(C_DD, D - Y, **solver_kwargs) # Center matrix Y = Y - np.mean(Y, axis=1, keepdims=True) _, num_ensemble = Y.shape # Don't left-multiply the X - if return_K: - return (Y.T @ K) / (num_ensemble - 1) # type: ignore + if return_T: + return (Y.T @ T) / (num_ensemble - 1) # type: ignore - return np.linalg.multi_dot([X, Y.T / (num_ensemble - 1), K]) # type: ignore + return np.linalg.multi_dot([X, Y.T / (num_ensemble - 1), T]) # type: ignore def inversion_exact_lstsq( @@ -253,8 +253,8 @@ def inversion_exact_lstsq( lhs = C_DD lhs.flat[:: lhs.shape[0] + 1] += alpha * C_D - # K = lhs^-1 @ (D - Y) - # lhs @ K = (D - Y) + # T = lhs^-1 @ (D - Y) + # lhs @ T = (D - Y) ans, *_ = sp.linalg.lstsq( lhs, D - Y, overwrite_a=True, overwrite_b=True, lapack_driver="gelsy" ) @@ -405,7 +405,7 @@ def inversion_subspace( Y: npt.NDArray[np.double], X: npt.NDArray[np.double], truncation: float = 1.0, - return_K: bool = False, + return_T: bool = False, ) -> npt.NDArray[np.double]: """See Appendix A.2 in :cite:t:`emerickHistoryMatchingTimelapse2012`. @@ -478,7 +478,7 @@ def inversion_subspace( # and finally we multiiply by (D - Y) term = U_r_w_inv @ Z - if return_K: + if return_T: return np.linalg.multi_dot( # type: ignore [D_delta.T, (term / (1 + T)), term.T, (D - Y)] ) diff --git a/tests/test_esmda_inversion.py b/tests/test_esmda_inversion.py index df534741..9335f2c5 100644 --- a/tests/test_esmda_inversion.py +++ b/tests/test_esmda_inversion.py @@ -45,7 +45,7 @@ class TestEsmdaInversion: @pytest.mark.parametrize("ensemble_members", [5, 10, 15]) @pytest.mark.parametrize("num_outputs", [5, 10, 15]) @pytest.mark.parametrize("num_inputs", [5, 10, 15]) - def test_that_returning_K_is_equivalent_to_full_computation( + def test_that_returning_T_is_equivalent_to_full_computation( self, function, ensemble_members, num_outputs, num_inputs ): # ensemble_members, num_outputs, num_inputs = 5, 10, 15 @@ -64,11 +64,11 @@ def test_that_returning_K_is_equivalent_to_full_computation( Y = np.random.randn(num_outputs, ensemble_members) X = np.random.randn(num_inputs, ensemble_members) - # Test both with and without X / return_K + # Test both with and without X / return_T ans = function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X) - K = function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=None, return_K=True) + T = function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=None, return_T=True) - assert np.allclose(X + ans, X + X @ K) + assert np.allclose(X + ans, X + X @ T) @pytest.mark.parametrize("length", list(range(1, 101, 5))) def test_that_the_sum_of_normalize_alpha_is_one(self, length):