From 8d6a6339c1528b23e8c377d36ba558c8d19fd44e Mon Sep 17 00:00:00 2001 From: Tommy Odland Date: Fri, 27 Oct 2023 07:15:48 +0200 Subject: [PATCH] added back basic row scaling --- .../experimental.py | 147 +++++++++++------- 1 file changed, 91 insertions(+), 56 deletions(-) diff --git a/src/iterative_ensemble_smoother/experimental.py b/src/iterative_ensemble_smoother/experimental.py index d9e46514..cff213c9 100644 --- a/src/iterative_ensemble_smoother/experimental.py +++ b/src/iterative_ensemble_smoother/experimental.py @@ -2,68 +2,103 @@ Contains (publicly available, but not officially supported) experimental features of iterative_ensemble_smoother """ +from typing import Union + import numpy as np +import numpy.typing as npt + +from iterative_ensemble_smoother import ESMDA -from iterative_ensemble_smoother.ies import create_coefficient_matrix -from iterative_ensemble_smoother.utils import covariance_to_correlation -rng = np.random.default_rng() +class RowScaling: + def multiply(self, X, K): + """Takes a matrix X and a matrix K and performs X @ K.""" + return X @ K def ensemble_smoother_update_step_row_scaling( - response_ensemble, - A_with_row_scaling, - observation_errors, - observation_values, - noise=None, - truncation=0.98, - inversion="exact", + *, + covariance: npt.NDArray[np.double], + observations: npt.NDArray[np.double], + X_with_row_scaling: list[tuple[npt.NDArray[np.double], RowScaling]], + Y: npt.NDArray[np.double], + seed: Union[np.random._generator.Generator, int, None] = None, + inversion: str = "exact", + truncation: float = 1.0, ): - """This is an experimental feature.""" - - # ============================================================================= - # # The old API used standard devisions for 1D arrays, but covariance for 2D - # covariance = observation_values**2 if - # observation_values.ndim == 1 else observation_values - # - # smoother = SIES(parameters=, - # covariance=covariance, observations=observation_values) - # ============================================================================= - - ensemble_size = response_ensemble.shape[1] - if noise is None: - num_obs = len(observation_values) - noise = rng.standard_normal(size=(num_obs, ensemble_size)) - - # Columns of E should be sampled from N(0,Cdd) and centered, Evensen 2019 - if len(observation_errors.shape) == 2: - E = np.linalg.cholesky(observation_errors) @ noise - else: - E = np.linalg.cholesky(np.diag(observation_errors**2)) @ noise - E = E @ ( - np.identity(ensemble_size) - - np.ones((ensemble_size, ensemble_size)) / ensemble_size + """Perform a single ESMDA update (ES) with row scaling. + The matrices in X_with_row_scaling WILL BE MUTATED. + + + Explanation of row scaling + -------------------------- + + The ESMDA update can be written as: + + X_post = X_prior + X_prior @ K + + where K is a transition matrix. The core of the row scaling approach is that + for each row i in the matrix X, we apply an update with strength alpha: + + X_post = X_prior + alpha * X_prior @ K + + Clearly 0 <= alpha <= 1 denotes the 'strength' of the update; alpha == 1 + corresponds to a normal smoother update and alpha == 0 corresponds to no + update. With the per row transformation of X the operation is no longer matrix + multiplication but the pseudo code looks like: + + for i in rows: + X_i_post = X_i_prior + alpha * X_i_prior @ K + + See also original code: + https://github.com/equinor/ert/blob/963f9bc08ebc87374b7ed3403c8ba78c20909ae9/src/clib/lib/enkf/row_scaling.cpp#L51 + + """ + + # https://github.com/equinor/ert/blob/963f9bc08ebc87374b7ed3403c8ba78c20909ae9/src/clib/lib/enkf/row_scaling.cpp#L51 + + # Create ESMDA instance + # Setting alpha=1 means we run a single data assimilation + smoother = ESMDA( + covariance=covariance, + observations=observations, + seed=seed, + inversion=inversion, + alpha=1, ) - R, observation_errors = covariance_to_correlation(observation_errors) - - D = (E + observation_values.reshape(num_obs, 1)) - response_ensemble - D = (D.T / observation_errors).T - E = (E.T / observation_errors).T - response_ensemble = (response_ensemble.T / observation_errors).T - for A, row_scale in A_with_row_scaling: - W = create_coefficient_matrix( - (response_ensemble - response_ensemble.mean(axis=1, keepdims=True)) - / np.sqrt(ensemble_size - 1), - R, - E, - D, - inversion, - truncation, - np.zeros((ensemble_size, ensemble_size)), - 1.0, - ) - I = np.identity(ensemble_size) # noqa: E741 - transition_matrix = I + W / np.sqrt(ensemble_size - 1) - row_scale.multiply(A, transition_matrix) - return A_with_row_scaling + # Create transition matrix - common to all parameters in X + transition_matrix = smoother.compute_transition_matrix(Y=Y, alpha=1, truncation=1.0) + + # Loop over groups of rows (parameters) + for X, row_scale in X_with_row_scaling: + X += row_scale.multiply(X, transition_matrix) + + return X_with_row_scaling + + +if __name__ == "__main__": + + # Example showing how to use row scaling + num_parameters = 100 + num_observations = 20 + num_ensemble = 10 + + rng = np.random.default_rng(42) + + X = rng.normal(size=(num_parameters, num_ensemble)) + Y = rng.normal(size=(num_observations, num_ensemble)) + covariance = np.exp(rng.normal(size=num_observations)) + observations = rng.normal(size=num_observations, loc=1) + + row_groups = [(0,), (1, 2), (4, 5, 6), tuple(range(7, 100))] + X_with_row_scaling = [(X[idx, :], RowScaling()) for idx in row_groups] + X_with_row_scaling = X_with_row_scaling.copy() + + X_with_row_scaling_updated = ensemble_smoother_update_step_row_scaling( + covariance=covariance, + observations=observations, + X_with_row_scaling=X_with_row_scaling, + Y=Y, + seed=rng, + )