From 083476f850f2bd5863cbfc0ec52c8dcebf404081 Mon Sep 17 00:00:00 2001 From: tommyod Date: Thu, 19 Oct 2023 07:13:46 +0200 Subject: [PATCH 1/7] removed 'ensemble_mask' argument --- src/iterative_ensemble_smoother/esmda.py | 28 +++----------- tests/test_esmda.py | 48 ------------------------ 2 files changed, 6 insertions(+), 70 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index 8255d5f3..e66f5cb3 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -22,7 +22,7 @@ """ import numbers -from typing import Optional, Union +from typing import Union import numpy as np import numpy.typing as npt @@ -165,7 +165,6 @@ def assimilate( self, X: npt.NDArray[np.double], Y: npt.NDArray[np.double], - ensemble_mask: Optional[npt.NDArray[np.bool_]] = None, overwrite: bool = False, truncation: float = 1.0, ) -> npt.NDArray[np.double]: @@ -179,10 +178,6 @@ def assimilate( 2D array of shape (num_parameters, ensemble_size). Y : np.ndarray 2D array of shape (num_parameters, ensemble_size). - ensemble_mask : np.ndarray - 1D boolean array of length `ensemble_size`, describing which - ensemble members are active. Inactive realizations are ignored. - Defaults to all active. overwrite : bool If True, then arguments X and Y may be overwritten. If False, then the method will not permute inputs in any way. @@ -206,9 +201,6 @@ def assimilate( assert ( num_ensemble == num_emsemble2 ), "Number of ensemble members in X and Y must match" - assert (ensemble_mask is None) or ( - ensemble_mask.ndim == 1 and len(ensemble_mask) == num_ensemble - ) if not np.issubdtype(X.dtype, np.floating): raise TypeError("Argument `X` must be contain floats") if not np.issubdtype(Y.dtype, np.floating): @@ -220,14 +212,6 @@ def assimilate( if not overwrite: X, Y = np.copy(X), np.copy(Y) - # If no ensemble mask was given, we use the entire ensemble - if ensemble_mask is None: - ensemble_mask = np.ones(num_ensemble, dtype=bool) - - # No ensemble members means no update - if ensemble_mask.sum() == 0: - return X - # Line 2 (b) in the description of ES-MDA in the 2013 Emerick paper # Draw samples from zero-centered multivariate normal with cov=alpha * C_D, @@ -235,7 +219,7 @@ def assimilate( # if C_D = L L.T by the cholesky factorization, then drawing y from # 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, ensemble_mask.sum()) + size = (num_outputs, num_ensemble) if self.C_D.ndim == 2: D = self.observations.reshape(-1, 1) + np.sqrt( self.alpha[self.iteration] @@ -244,7 +228,7 @@ def assimilate( D = self.observations.reshape(-1, 1) + np.sqrt( self.alpha[self.iteration] ) * self.rng.normal(size=size) * self.C_D_L.reshape(-1, 1) - assert D.shape == (num_outputs, ensemble_mask.sum()) + assert D.shape == (num_outputs, num_ensemble) # Line 2 (c) in the description of ES-MDA in the 2013 Emerick paper # Choose inversion method, e.g. 'exact'. The inversion method computes @@ -252,12 +236,12 @@ def assimilate( inversion_func = self._inversion_methods[self.inversion] # Update and return - X[:, ensemble_mask] += inversion_func( + X += inversion_func( alpha=self.alpha[self.iteration], C_D=self.C_D, D=D, - Y=Y[:, ensemble_mask], - X=X[:, ensemble_mask], + Y=Y, + X=X, truncation=truncation, ) diff --git a/tests/test_esmda.py b/tests/test_esmda.py index b60c4eff..1c218c22 100644 --- a/tests/test_esmda.py +++ b/tests/test_esmda.py @@ -68,54 +68,6 @@ def test_that_diagonal_covariance_gives_same_answer_as_dense(self, seed, inversi assert np.allclose(X_posterior1, X_posterior2) - @pytest.mark.parametrize("num_ensemble", [2**i for i in range(2, 10)]) - def test_that_using_example_mask_only_updates_those_parameters(self, num_ensemble): - seed = num_ensemble - rng = np.random.default_rng(seed) - alpha = rng.choice(np.array([5, 10, 25, 50])) # Random alpha - - num_outputs = 2 - num_inputs = 1 - - def g(x): - """Transform a single ensemble member.""" - return np.array([np.sin(x / 2), x]) - - def G(X): - """Transform all ensemble members.""" - return np.array([g(x_i) for x_i in X.T]).squeeze().T - - # Create an ensemble mask and set half the entries randomly to True - ensemble_mask = np.zeros(num_ensemble, dtype=bool) - idx = rng.choice(num_ensemble, size=num_ensemble // 2, replace=False) - ensemble_mask[idx] = True - - # Prior is N(0, 1) - X_prior = rng.normal(size=(num_inputs, num_ensemble)) - - # Measurement errors - covariance = np.eye(num_outputs) - - # The true inputs and observations, a result of running with X_true = 1 - X_true = np.ones(shape=(num_inputs,)) - observations = G(X_true) - - # Prepare ESMDA instance running with lower number of ensemble members - esmda_subset = ESMDA(covariance, observations, alpha=alpha, seed=seed) - X_i_subset = np.copy(X_prior[:, ensemble_mask]) - - # Prepare ESMDA instance running with all ensemble members - esmda_masked = ESMDA(covariance, observations, alpha=alpha, seed=seed) - X_i = np.copy(X_prior) - - # Run both - for _ in range(esmda_subset.num_assimilations()): - X_i_subset = esmda_subset.assimilate(X_i_subset, G(X_i_subset)) - X_i = esmda_masked.assimilate(X_i, G(X_i), ensemble_mask=ensemble_mask) - - # Exactly the same result (seed value for both must be equal ) - assert np.allclose(X_i_subset, X_i[:, ensemble_mask]) - @pytest.mark.parametrize( "num_ensemble", [10, 100, 1000], From a0edab18da8beddd249a1b2896ef7f06127775d1 Mon Sep 17 00:00:00 2001 From: tommyod Date: Thu, 19 Oct 2023 07:37:49 +0200 Subject: [PATCH 2/7] create method get_D --- src/iterative_ensemble_smoother/esmda.py | 54 +++++++++++++++++++----- 1 file changed, 44 insertions(+), 10 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index 8255d5f3..5a831eb4 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -22,7 +22,7 @@ """ import numbers -from typing import Optional, Union +from typing import Optional, Tuple, Union import numpy as np import numpy.typing as npt @@ -42,7 +42,7 @@ class ESMDA: Parameters ---------- - covariance : npt.NDArray[np.double] + covariance : np.ndarray Either a 1D array of diagonal covariances, or a 2D covariance matrix. The shape is either (num_observations,) or (num_observations, num_observations). This is C_D in Emerick (2013), and represents observation or measurement @@ -236,14 +236,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, ensemble_mask.sum()) - if self.C_D.ndim == 2: - D = self.observations.reshape(-1, 1) + np.sqrt( - self.alpha[self.iteration] - ) * self.C_D_L @ self.rng.normal(size=size) - else: - D = self.observations.reshape(-1, 1) + np.sqrt( - self.alpha[self.iteration] - ) * self.rng.normal(size=size) * self.C_D_L.reshape(-1, 1) + D = self.get_D(size=size, alpha=self.alpha[self.iteration]) assert D.shape == (num_outputs, ensemble_mask.sum()) # Line 2 (c) in the description of ES-MDA in the 2013 Emerick paper @@ -264,6 +257,47 @@ def assimilate( self.iteration += 1 return X + def get_D(self, *, size: Tuple[int, int], alpha: float) -> npt.NDArray[np.double]: + """Create a matrix D with perturbed observations. + + In the Emerick (2013) paper, the matrix D is defined in section 6. + See section 2(b) of the ES-MDA algorithm in the paper. + + + Parameters + ---------- + size : Tuple[int, int] + The size, a tuple with (num_observations, ensemble_size). + alpha : float + The inflation factor for the covariance. + + Returns + ------- + D : np.ndarray + Each column consists of perturbed observations, scaled by alpha. + + """ + # Draw samples from zero-centered multivariate normal with cov=alpha * C_D, + # and add them to the observations. Notice that + # if C_D = L @ L.T by the cholesky factorization, then drawing y from + # 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). + + D: npt.NDArray[np.double] + + # Two cases, depending on whether C_D was given as 1D or 2D array + if self.C_D.ndim == 2: + D = self.observations.reshape(-1, 1) + np.sqrt( + self.alpha[self.iteration] + ) * self.C_D_L @ self.rng.normal(size=size) + else: + D = self.observations.reshape(-1, 1) + np.sqrt( + self.alpha[self.iteration] + ) * self.rng.normal(size=size) * self.C_D_L.reshape(-1, 1) + assert D.shape == size + + return D + if __name__ == "__main__": import pytest From 64e04e02ff3b3ee86376b797adaaf9cb1c23bf9d Mon Sep 17 00:00:00 2001 From: tommyod Date: Thu, 19 Oct 2023 09:59:57 +0200 Subject: [PATCH 3/7] test failing - unsure why --- src/iterative_ensemble_smoother/esmda.py | 45 +++++++++++++ .../esmda_inversion.py | 26 ++++++-- tests/test_esmda.py | 66 ++++++++++++++++++- tests/test_esmda_inversion.py | 36 +++++++++- 4 files changed, 166 insertions(+), 7 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index 5a831eb4..99c0c5f2 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -165,6 +165,7 @@ def assimilate( self, X: npt.NDArray[np.double], Y: npt.NDArray[np.double], + *, ensemble_mask: Optional[npt.NDArray[np.bool_]] = None, overwrite: bool = False, truncation: float = 1.0, @@ -257,6 +258,50 @@ def assimilate( self.iteration += 1 return X + def get_K( + self, + Y: npt.NDArray[np.double], + *, + alpha: float, + truncation: float = 1.0, + ) -> npt.NDArray[np.double]: + """Return a matrix K such that X_posterior = X_prior + center(X_prior) @ K. + + 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. + + + Parameters + ---------- + Y : npt.NDArray[np.double] + DESCRIPTION. + * : TYPE + DESCRIPTION. + alpha : float + DESCRIPTION. + truncation : float, optional + DESCRIPTION. The default is 1.0. + : TYPE + DESCRIPTION. + + Returns + ------- + None. + + """ + + D = self.get_D(size=Y.shape, alpha=alpha) + inversion_func = self._inversion_methods[self.inversion] + return inversion_func( + alpha=alpha, + C_D=self.C_D, + D=D, + Y=Y, + X=None, + truncation=truncation, + return_K=True, # Ensures that we don't need X + ) + def get_D(self, *, size: Tuple[int, int], alpha: float) -> npt.NDArray[np.double]: """Create a matrix D with perturbed observations. diff --git a/src/iterative_ensemble_smoother/esmda_inversion.py b/src/iterative_ensemble_smoother/esmda_inversion.py index 5d1586d9..a4e4a678 100644 --- a/src/iterative_ensemble_smoother/esmda_inversion.py +++ b/src/iterative_ensemble_smoother/esmda_inversion.py @@ -1,3 +1,5 @@ +from typing import Optional + import numpy as np import numpy.typing as npt import scipy as sp # type: ignore @@ -178,8 +180,9 @@ def inversion_exact_cholesky( C_D: npt.NDArray[np.double], D: npt.NDArray[np.double], Y: npt.NDArray[np.double], - X: npt.NDArray[np.double], + X: Optional[npt.NDArray[np.double]], truncation: float = 1.0, + return_K: 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. @@ -211,12 +214,16 @@ def inversion_exact_cholesky( C_DD.flat[:: C_DD.shape[1] + 1] += alpha * C_D K = sp.linalg.solve(C_DD, D - Y, **solver_kwargs) - # Form C_MD = center(X) @ center(Y).T / (num_ensemble - 1) - X = X - np.mean(X, axis=1, keepdims=True) + # Center matrix Y = Y - np.mean(Y, axis=1, keepdims=True) - _, num_ensemble = X.shape + _, num_ensemble = Y.shape - return np.linalg.multi_dot([X / (num_ensemble - 1), Y.T, K]) # type: ignore + # Don't left-multiply the X + if return_K: + return (Y.T @ K) / (num_ensemble - 1) # type: ignore + + X = X - np.mean(X, axis=1, keepdims=True) + return np.linalg.multi_dot([X, Y.T / (num_ensemble - 1), K]) # type: ignore def inversion_exact_lstsq( @@ -396,6 +403,7 @@ def inversion_subspace( Y: npt.NDArray[np.double], X: npt.NDArray[np.double], truncation: float = 1.0, + return_K: bool = False, ) -> npt.NDArray[np.double]: """See Appendix A.2 in Emerick et al (2012) @@ -427,6 +435,9 @@ def inversion_subspace( [0., 0., 0.]]) """ + print( + {"alpha": alpha, "C_D": C_D, "D": D, "Y": Y, "X": X, "truncation": truncation} + ) # N_n is the number of observations # N_e is the number of members in the ensemble @@ -468,6 +479,11 @@ def inversion_subspace( # and finally we multiiply by (D - Y) term = U_r_w_inv @ Z + if return_K: + return np.linalg.multi_dot( # type: ignore + [D_delta.T, (term / (1 + T)), term.T, (D - Y)] + ) + # Compute C_MD = center(X) @ center(Y).T / (num_ensemble - 1) X_shift = X - np.mean(X, axis=1, keepdims=True) return np.linalg.multi_dot( # type: ignore diff --git a/tests/test_esmda.py b/tests/test_esmda.py index b60c4eff..7cca5612 100644 --- a/tests/test_esmda.py +++ b/tests/test_esmda.py @@ -422,6 +422,70 @@ def test_that_float_dtypes_are_preserved(self, inversion, dtype, diagonal): assert X_posterior.dtype == dtype +@pytest.mark.parametrize("inversion", ESMDA._inversion_methods.keys()) +def test_row_by_row_assimilation(inversion): + # Create problem instance + rng = np.random.default_rng(42) + + num_outputs = 4 + num_inputs = 5 + num_ensemble = 3 + + A = rng.normal(size=(num_outputs, num_inputs)) + + def g(X): + return A @ X + + # Prior is N(0, 1) + X_prior = rng.normal(size=(num_inputs, num_ensemble)) + + covariance = np.exp(rng.normal(size=num_outputs)) + observations = A @ np.linspace(0, 1, num=num_inputs) + rng.normal( + size=num_outputs, scale=0.01 + ) + + # =========== Use the high level API =========== + smoother = ESMDA( + covariance=covariance, + observations=observations, + alpha=2, + inversion=inversion, + seed=1, + ) + X = np.copy(X_prior) + print(X_prior) + for iteration in range(smoother.num_assimilations()): + X = smoother.assimilate(X, g(X)) + print(X) + + X_posterior_highlevel_API = np.copy(X) + + # =========== Use the low-level level API =========== + print("# =========== Use the low-level level API ===========") + smoother = ESMDA( + covariance=covariance, + observations=observations, + alpha=2, + inversion=inversion, + seed=1, + ) + X = np.copy(X_prior) + print(X_prior) + for alpha_i in smoother.alpha: + K = smoother.get_K(Y=g(X), alpha=alpha_i) + + # TODO: Why is this equivalent? + X_centered = X - np.mean(X, axis=1, keepdims=True) + assert np.allclose(X_centered @ K, X @ K) + + X += X_centered @ K + print(X) + + X_posterior_lowlevel_API = np.copy(X) + + assert np.allclose(X_posterior_highlevel_API, X_posterior_lowlevel_API) + + if __name__ == "__main__": import pytest @@ -429,6 +493,6 @@ def test_that_float_dtypes_are_preserved(self, inversion, dtype, diagonal): args=[ __file__, "-v", - "-k test_that_float_dtypes_are_preserved", + "-k test_row_by_row_assimilation", ] ) diff --git a/tests/test_esmda_inversion.py b/tests/test_esmda_inversion.py index 465087a1..b1f49e05 100644 --- a/tests/test_esmda_inversion.py +++ b/tests/test_esmda_inversion.py @@ -37,6 +37,40 @@ class TestEsmdaInversion: + # Functions that take the parameter return_K + @pytest.mark.parametrize( + "function", + [inversion_exact_cholesky, inversion_subspace], + ) + @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( + self, function, ensemble_members, num_outputs, num_inputs + ): + # ensemble_members, num_outputs, num_inputs = 5, 10, 15 + + np.random.seed(ensemble_members * 100 + num_outputs * 10 + num_inputs) + + # Create positive symmetric definite covariance C_D + E = np.random.randn(num_outputs, num_outputs) + C_D = E.T @ E + + # Set alpha to something other than 1 to test that it works + alpha = 3 + + # Create observations + D = np.random.randn(num_outputs, ensemble_members) + Y = np.random.randn(num_outputs, ensemble_members) + X = np.random.randn(num_inputs, ensemble_members) + + # Test both with and without X / return_K + 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) + + X_centered = X - np.mean(X, axis=1, keepdims=True) + assert np.allclose(X + ans, X + X_centered @ K) + @pytest.mark.parametrize("length", list(range(1, 101, 5))) def test_that_the_sum_of_normalize_alpha_is_one(self, length): # Generate random array @@ -350,6 +384,6 @@ def test_timing(num_outputs=100, num_inputs=50, ensemble_members=25): args=[ __file__, "-v", - "-k test_that_exact_inverions_are_equal_with_few_ensemble_members", + "-k test_that_returning_K_is_equivalent_to_full_computation", ] ) From 5a77c23ec1bc1f41c075a54efdac29a8cfc7e51b Mon Sep 17 00:00:00 2001 From: tommyod Date: Thu, 19 Oct 2023 10:25:51 +0200 Subject: [PATCH 4/7] updated docstrings --- src/iterative_ensemble_smoother/esmda.py | 46 +++++++++++++----------- tests/test_esmda.py | 8 ++--- tests/test_esmda_inversion.py | 4 +-- 3 files changed, 30 insertions(+), 28 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index a8d1f489..07408ebb 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -171,17 +171,20 @@ def assimilate( ) -> npt.NDArray[np.double]: """Assimilate data and return an updated ensemble X_posterior. - num_parameters, ensemble_size - Parameters ---------- X : np.ndarray - 2D array of shape (num_parameters, ensemble_size). + A 2D array of shape (num_parameters, ensemble_size). Each row + 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). + 2D array of shape (num_parameters, ensemble_size), containing + responses when evaluating the model at X. In other words, Y = g(X), + where g is the forward model. overwrite : bool - If True, then arguments X and Y may be overwritten. - If False, then the method will not permute inputs in any way. + If True, then arguments X and Y may be overwritten and mutated. + If False, then the method will not mutate inputs in any way. + Setting this to True might save memory. truncation : float How large a fraction of the singular values to keep in the inversion routine. Must be a float in the range (0, 1]. A lower number means @@ -254,24 +257,26 @@ def get_K( 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. - Parameters ---------- - Y : npt.NDArray[np.double] - DESCRIPTION. - * : TYPE - DESCRIPTION. + Y : np.ndarray + 2D array of shape (num_parameters, ensemble_size), containing + responses when evaluating the model at X. In other words, Y = g(X), + where g is the forward model. alpha : float - DESCRIPTION. - truncation : float, optional - DESCRIPTION. The default is 1.0. - : TYPE - DESCRIPTION. + 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. + truncation : float + How large a fraction of the singular values to keep in the inversion + routine. Must be a float in the range (0, 1]. A lower number means + a more approximate answer and a slightly faster computation. Returns ------- - None. - + K : np.ndarray + A matrix K such that X_posterior = X_prior + X_prior @ K. + It has shape (num_ensemble_members, num_ensemble_members). """ D = self.get_D(size=Y.shape, alpha=alpha) @@ -292,13 +297,14 @@ def get_D(self, *, size: Tuple[int, int], alpha: float) -> npt.NDArray[np.double In the Emerick (2013) paper, the matrix D is defined in section 6. See section 2(b) of the ES-MDA algorithm in the paper. - Parameters ---------- size : Tuple[int, int] The size, a tuple with (num_observations, ensemble_size). alpha : float - The inflation factor for the covariance. + 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 ------- diff --git a/tests/test_esmda.py b/tests/test_esmda.py index ecea768e..1536bef9 100644 --- a/tests/test_esmda.py +++ b/tests/test_esmda.py @@ -405,15 +405,12 @@ def g(X): seed=1, ) X = np.copy(X_prior) - print(X_prior) for iteration in range(smoother.num_assimilations()): X = smoother.assimilate(X, g(X)) - print(X) X_posterior_highlevel_API = np.copy(X) # =========== Use the low-level level API =========== - print("# =========== Use the low-level level API ===========") smoother = ESMDA( covariance=covariance, observations=observations, @@ -422,16 +419,15 @@ def g(X): seed=1, ) X = np.copy(X_prior) - print(X_prior) for alpha_i in smoother.alpha: K = smoother.get_K(Y=g(X), alpha=alpha_i) - # TODO: Why is this equivalent? + # TODO: Why is this equivalent? ... X_centered = X - np.mean(X, axis=1, keepdims=True) assert np.allclose(X_centered @ K, X @ K) + # Here we could loop over each row in X and multiply by K X += X_centered @ K - print(X) X_posterior_lowlevel_API = np.copy(X) diff --git a/tests/test_esmda_inversion.py b/tests/test_esmda_inversion.py index b1f49e05..dc14ec5b 100644 --- a/tests/test_esmda_inversion.py +++ b/tests/test_esmda_inversion.py @@ -68,8 +68,8 @@ def test_that_returning_K_is_equivalent_to_full_computation( 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) - X_centered = X - np.mean(X, axis=1, keepdims=True) - assert np.allclose(X + ans, X + X_centered @ K) + X - np.mean(X, axis=1, keepdims=True) + assert np.allclose(X + ans, X + X @ K) @pytest.mark.parametrize("length", list(range(1, 101, 5))) def test_that_the_sum_of_normalize_alpha_is_one(self, length): From a1c26fccb2bb611e1e0fcbe5f74a67c258be9304 Mon Sep 17 00:00:00 2001 From: tommyod Date: Thu, 19 Oct 2023 10:39:26 +0200 Subject: [PATCH 5/7] more descriptive method names --- src/iterative_ensemble_smoother/esmda.py | 13 +++++++++---- src/iterative_ensemble_smoother/esmda_inversion.py | 3 --- tests/test_esmda.py | 2 +- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index 07408ebb..2ecd99d1 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -224,7 +224,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.get_D(size=size, alpha=self.alpha[self.iteration]) + D = self.perturb_observations(size=size, alpha=self.alpha[self.iteration]) assert D.shape == (num_outputs, num_ensemble) # Line 2 (c) in the description of ES-MDA in the 2013 Emerick paper @@ -245,7 +245,7 @@ def assimilate( self.iteration += 1 return X - def get_K( + def compute_transition_matrix( self, Y: npt.NDArray[np.double], *, @@ -279,7 +279,10 @@ def get_K( It has shape (num_ensemble_members, num_ensemble_members). """ - D = self.get_D(size=Y.shape, alpha=alpha) + # Recall the update equation + # X += C_MD @ (C_DD + alpha * C_D)^(-1) @ (D - Y) + + D = self.perturb_observations(size=Y.shape, alpha=alpha) inversion_func = self._inversion_methods[self.inversion] return inversion_func( alpha=alpha, @@ -291,7 +294,9 @@ def get_K( return_K=True, # Ensures that we don't need X ) - def get_D(self, *, size: Tuple[int, int], alpha: float) -> npt.NDArray[np.double]: + def perturb_observations( + self, *, size: Tuple[int, int], alpha: float + ) -> npt.NDArray[np.double]: """Create a matrix D with perturbed observations. In the Emerick (2013) paper, the matrix D is defined in section 6. diff --git a/src/iterative_ensemble_smoother/esmda_inversion.py b/src/iterative_ensemble_smoother/esmda_inversion.py index a4e4a678..f9001cea 100644 --- a/src/iterative_ensemble_smoother/esmda_inversion.py +++ b/src/iterative_ensemble_smoother/esmda_inversion.py @@ -435,9 +435,6 @@ def inversion_subspace( [0., 0., 0.]]) """ - print( - {"alpha": alpha, "C_D": C_D, "D": D, "Y": Y, "X": X, "truncation": truncation} - ) # N_n is the number of observations # N_e is the number of members in the ensemble diff --git a/tests/test_esmda.py b/tests/test_esmda.py index 1536bef9..143764aa 100644 --- a/tests/test_esmda.py +++ b/tests/test_esmda.py @@ -420,7 +420,7 @@ def g(X): ) X = np.copy(X_prior) for alpha_i in smoother.alpha: - K = smoother.get_K(Y=g(X), alpha=alpha_i) + K = smoother.compute_transition_matrix(Y=g(X), alpha=alpha_i) # TODO: Why is this equivalent? ... X_centered = X - np.mean(X, axis=1, keepdims=True) From 59923069a480bbda842402e411c919505971d23c Mon Sep 17 00:00:00 2001 From: tommyod Date: Thu, 19 Oct 2023 15:10:48 +0200 Subject: [PATCH 6/7] realized that we don't need to center X --- src/iterative_ensemble_smoother/esmda.py | 12 ++++-- .../esmda_inversion.py | 37 ++++++++++--------- tests/test_esmda.py | 6 +-- 3 files changed, 29 insertions(+), 26 deletions(-) diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py index 2ecd99d1..9de8a13e 100644 --- a/src/iterative_ensemble_smoother/esmda.py +++ b/src/iterative_ensemble_smoother/esmda.py @@ -252,7 +252,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 + center(X_prior) @ K. + """Return a matrix K such that X_posterior = X_prior + X_prior @ K. 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. @@ -279,8 +279,14 @@ def compute_transition_matrix( It has shape (num_ensemble_members, num_ensemble_members). """ - # Recall the update equation + # 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), + # so that + # X_new = X_old + X_old @ K + # or + # X += X @ K D = self.perturb_observations(size=Y.shape, alpha=alpha) inversion_func = self._inversion_methods[self.inversion] @@ -289,7 +295,7 @@ def compute_transition_matrix( C_D=self.C_D, D=D, Y=Y, - X=None, + X=None, # We don't need X to compute the factor K truncation=truncation, return_K=True, # Ensures that we don't need X ) diff --git a/src/iterative_ensemble_smoother/esmda_inversion.py b/src/iterative_ensemble_smoother/esmda_inversion.py index f9001cea..ac42551a 100644 --- a/src/iterative_ensemble_smoother/esmda_inversion.py +++ b/src/iterative_ensemble_smoother/esmda_inversion.py @@ -66,9 +66,15 @@ def empirical_cross_covariance( assert X.shape[1] == Y.shape[1], "Ensemble size must be equal" # https://en.wikipedia.org/wiki/Estimation_of_covariance_matrices - # Subtract means - X = X - np.mean(X, axis=1, keepdims=True) - Y = Y - np.mean(Y, axis=1, keepdims=True) + # Subtract mean. Even though the equation says E[(X - E[X])(Y - E[Y])^T], + # we actually only need to subtract the mean value from one matrix, since + # (X - E[X])(Y - E[Y])^T = E[(X - E[X])Y] - E[(X - E[X])E[Y]^T] + # = E[(X - E[X])Y] - E[(0)E[Y]^T] = E[(X - E[X])Y] + # We choose to subtract from the matrix with the smaller number of rows + if X.shape[0] > Y.shape[0]: + Y = Y - np.mean(Y, axis=1, keepdims=True) + else: + X = X - np.mean(X, axis=1, keepdims=True) # Compute outer product and divide cov = X @ Y.T / (X.shape[1] - 1) @@ -222,7 +228,6 @@ def inversion_exact_cholesky( if return_K: return (Y.T @ K) / (num_ensemble - 1) # type: ignore - X = X - np.mean(X, axis=1, keepdims=True) return np.linalg.multi_dot([X, Y.T / (num_ensemble - 1), K]) # type: ignore @@ -254,10 +259,9 @@ def inversion_exact_lstsq( lhs, D - Y, overwrite_a=True, overwrite_b=True, lapack_driver="gelsy" ) - # Compute C_MD := center(X) @ center(Y).T / (X.shape[1] - 1) - X_shift = (X - np.mean(X, axis=1, keepdims=True)) / (X.shape[1] - 1) - Y_shift = Y - np.mean(Y, axis=1, keepdims=True) - return np.linalg.multi_dot([X_shift, Y_shift.T, ans]) # type: ignore + # Compute C_MD := X @ center(Y).T / (Y.shape[1] - 1) + Y_shift = (Y - np.mean(Y, axis=1, keepdims=True)) / (Y.shape[1] - 1) + return np.linalg.multi_dot([X, Y_shift.T, ans]) # type: ignore def inversion_exact_rescaled( @@ -318,11 +322,10 @@ def inversion_exact_rescaled( term = C_D_L_inv.T @ U_r if C_D.ndim == 2 else (C_D_L_inv * U_r.T).T # Compute the first factors, which make up C_MD - X_shift = (X - np.mean(X, axis=1, keepdims=True)) / (N_e - 1) - Y_shift = Y - np.mean(Y, axis=1, keepdims=True) + Y_shift = (Y - np.mean(Y, axis=1, keepdims=True)) / (N_e - 1) return np.linalg.multi_dot( # type: ignore - [X_shift, Y_shift.T, term / s_r, term.T, (D - Y)] + [X, Y_shift.T, term / s_r, term.T, (D - Y)] ) @@ -359,7 +362,7 @@ def inversion_exact_subspace_woodbury( D_delta /= np.sqrt(N_e - 1) # Compute the first factors, which make up C_MD - X_shift = (X - np.mean(X, axis=1, keepdims=True)) / np.sqrt(N_e - 1) + # X_shift = (X - np.mean(X, axis=1, keepdims=True)) / np.sqrt(N_e - 1) # A full covariance matrix was given if C_D.ndim == 2: @@ -377,7 +380,7 @@ def inversion_exact_subspace_woodbury( # Compute the woodbury inversion, then return inverted = C_D_inv - np.linalg.multi_dot([term, sp.linalg.inv(center), term.T]) return np.linalg.multi_dot( # type: ignore - [X_shift, D_delta.T, inverted, (D - Y)] + [X, D_delta.T / np.sqrt(N_e - 1), inverted, (D - Y)] ) # A diagonal covariance matrix was given as a 1D array. @@ -391,7 +394,7 @@ def inversion_exact_subspace_woodbury( [UT_D.T, sp.linalg.inv(center), UT_D] ) return np.linalg.multi_dot( # type: ignore - [X_shift, D_delta.T, inverted, (D - Y)] + [X, D_delta.T / np.sqrt(N_e - 1), inverted, (D - Y)] ) @@ -482,9 +485,8 @@ def inversion_subspace( ) # Compute C_MD = center(X) @ center(Y).T / (num_ensemble - 1) - X_shift = X - np.mean(X, axis=1, keepdims=True) return np.linalg.multi_dot( # type: ignore - [X_shift, D_delta.T, (term / (1 + T)), term.T, (D - Y)] + [X, D_delta.T, (term / (1 + T)), term.T, (D - Y)] ) @@ -543,9 +545,8 @@ def inversion_rescaled_subspace( diag = 1 / (1 + T_r) # Compute C_MD - X_shift = X - np.mean(X, axis=1, keepdims=True) return np.linalg.multi_dot( # type: ignore - [X_shift, D_delta.T, (term * diag), term.T, (D - Y)] + [X, D_delta.T, (term * diag), term.T, (D - Y)] ) diff --git a/tests/test_esmda.py b/tests/test_esmda.py index 143764aa..3b7d3358 100644 --- a/tests/test_esmda.py +++ b/tests/test_esmda.py @@ -422,12 +422,8 @@ def g(X): for alpha_i in smoother.alpha: K = smoother.compute_transition_matrix(Y=g(X), alpha=alpha_i) - # TODO: Why is this equivalent? ... - X_centered = X - np.mean(X, axis=1, keepdims=True) - assert np.allclose(X_centered @ K, X @ K) - # Here we could loop over each row in X and multiply by K - X += X_centered @ K + X += X @ K X_posterior_lowlevel_API = np.copy(X) From 4dfa3912327e571cebab865f83c1d4bf00e1a3f2 Mon Sep 17 00:00:00 2001 From: Tommy Odland Date: Mon, 23 Oct 2023 11:18:13 +0200 Subject: [PATCH 7/7] remove unused centering in test --- tests/test_esmda_inversion.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_esmda_inversion.py b/tests/test_esmda_inversion.py index dc14ec5b..df534741 100644 --- a/tests/test_esmda_inversion.py +++ b/tests/test_esmda_inversion.py @@ -68,7 +68,6 @@ def test_that_returning_K_is_equivalent_to_full_computation( 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) - X - np.mean(X, axis=1, keepdims=True) assert np.allclose(X + ans, X + X @ K) @pytest.mark.parametrize("length", list(range(1, 101, 5)))