iterative_ensemble_smoother.AdaptiveESMDA#

class iterative_ensemble_smoother.AdaptiveESMDA(covariance: ndarray[tuple[int, ...], dtype[floating]], observations: ndarray[tuple[int, ...], dtype[floating]], alpha: int | ndarray[tuple[int, ...], dtype[floating]] = 5, seed: Generator | int | None = None)[source]#

Bases: BaseESMDA

Hard Thresholding Adaptive ESMDA.

References

  • Adaptive Correlation- and Distance-Based Localization for Iterative Ensemble Smoothers in a Coupled Nonlinear Multiscale Model. Vossepoel, Femke & Evensen, Geir & Van Leeuwen, Peter Jan. (2025) http://doi.org/10.1175/MWR-D-24-0269.1

Examples

A full example where the forward model maps 10 parameters to 3 outputs. We will use 100 realizations. First we define the forward model:

>>> rng = np.random.default_rng(42)
>>> A = rng.normal(size=(3, 10))
>>> def forward_model(x):
...     return A @ x

Then we set up the AdaptiveESMDA instance and the prior realizations X:

>>> covariance = np.ones(3, dtype=float)  # Covariance of the observations / outputs
>>> observations = np.array([1, 2, 3], dtype=float)  # The observed data
>>> smoother = AdaptiveESMDA(covariance=covariance,
...                          observations=observations, alpha=3, seed=42)
>>> X = rng.normal(size=(10, 100))

To assimilate data, we iterate over the assimilation steps in an outer loop, then over parameter batches:

>>> def yield_param_indices():
...     yield [1, 2, 3, 4]
...     yield [5, 6, 7, 8, 9]
>>> for iteration in range(smoother.num_assimilations()):
...
...     Y = np.array([forward_model(x) for x in X.T]).T
...
...     # Prepare for assimilation
...     smoother.prepare_assimilation(Y=Y, truncation=0.99)
...
...     def func(corr_XY, ensemble_members_per_parameter):
...         # Takes an array of shape (params_batch, obs)
...         # and an array representing number of non-missing ensemble
...         # members per parameter. Returns which pairs (param, obs) to keep.
...         return np.ones_like(corr_XY, dtype=np.bool_)
...
...     for param_idx in yield_param_indices():
...         X[param_idx, :] = smoother.assimilate_batch(X=X[param_idx, :])
...

Methods

assimilate_batch(*, X[, missing, ...])

Assimilate a batch of parameters against all observations.

perturb_observations(*, ensemble_size, alpha)

Create a matrix D with perturbed observations.

prepare_assimilation(*, Y[, truncation, ...])

Prepare assimilation of one or several batches of parameters.

three_over_sqrt_ensemble_members(corr_XY, ...)

Use the correlation threshold 3 / sqrt(n).

num_assimilations
__init__(covariance: ndarray[tuple[int, ...], dtype[floating]], observations: ndarray[tuple[int, ...], dtype[floating]], alpha: int | ndarray[tuple[int, ...], dtype[floating]] = 5, seed: Generator | int | None = None) → None#
Parameters:

  • covariance (np.ndarray) – Either a 1D array of diagonal covariances, or a 2D covariance matrix. The shape is (num_observations,) or (num_observations, num_observations). This is C_D in Emerick (2013), and represents observation or measurement errors. We observe d from the real world, y from the model g(x), and assume that d = y + e, where the error e is multivariate normal with covariance given by covariance.

  • observations (np.ndarray) – 1D array of shape (num_observations,) representing real-world observations. This is d_obs in Emerick (2013).

  • alpha (int or 1D np.ndarray, optional) – Multiplicative factor for the covariance. If an integer alpha is given, an array with length alpha and elements alpha is constructed. If an 1D array is given, it is normalized so sum_i 1/alpha_i = 1 and used. The default is 5, which corresponds to np.array([5, 5, 5, 5, 5]).

  • seed (integer or numpy.random.Generator, optional) – A seed or numpy.random.Generator used for random number generation. The argument is passed to numpy.random.default_rng(). The default is None.

Methods definition

_compute_delta_M(*, X: ndarray[tuple[int, ...], dtype[floating]], missing: ndarray[tuple[int, ...], dtype[bool]] | None = None) → ndarray[tuple[int, ...], dtype[floating]]#

Prepare delta_M := X - center(X), dealing with missing values as needed.

assimilate_batch(*, X: ndarray[tuple[int, ...], dtype[floating]], missing: ndarray[tuple[int, ...], dtype[bool]] | None = None, correlation_callback: Callable[[ndarray[tuple[int, ...], dtype[floating]], ndarray[tuple[int, ...], dtype[int64]] | int], ndarray[tuple[int, ...], dtype[bool]]] | str = 'three_over_sqrt_ensemble_members', overwrite: bool = False) → ndarray[tuple[int, ...], dtype[floating]][source]#

Assimilate a batch of parameters against all observations.

The internal storage used by the class is 2 * ensemble_size * num_observations, so a good batch size that is of the same order of magnitude as the internal storage is 2 * num_observations. This is only a rough guideline.

Parameters:

  • X (np.ndarray) – A 2D array of shape (num_parameters_batch, ensemble_size). Each row corresponds to a parameter in the model, and each column corresponds to an ensemble member (realization).

  • missing (np.ndarray or None) – A boolean 2D array of shape (num_parameters_batch, ensemble_size). If an entry is True, then that value is assumed missing. This can happen if the ensemble members use different grids, where each ensemble member has a slightly different grid layout. If None, then all entries are assumed to be valid.

  • correlation_callback (callable or string, optional) – A callable with signature (corr_XY, ensemble_members_per_parameter) -> mask. corr_XY is a cross-correlation 2D array of shape (num_parameters_batch, num_observations). ensemble_members_per_parameter is either an int (when there are no missing values, equal to the ensemble size) or a 1D int array of shape (num_parameters_batch,) giving the number of non-missing ensemble members for each parameter. The returned array must have the same shape as corr_XY and represents any kind of correlation thresholding or softening.

  • overwrite (bool) – If False (the default), the input arrays will not be overwritten (mutated). If True, the method may overwrite the input arrays.

Returns:

X_posterior – 2D array of shape (num_parameters_batch, ensemble_size).

Return type:

np.ndarray

perturb_observations(*, ensemble_size: int, alpha: float) → ndarray[tuple[int, ...], dtype[floating]]#

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:

  • ensemble_size (int) – The ensemble size, i.e., the number of columns in the returned array, which is of shape (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:

D – Each column consists of perturbed observations, observation std is scaled by sqrt(alpha).

Return type:

np.ndarray

prepare_assimilation(*, Y: ndarray[tuple[int, ...], dtype[floating]], truncation: float = 0.99, overwrite: bool = False) → None#

Prepare assimilation of one or several batches of parameters.

This method call pre-computes everything that is needed to assimilate a set of batches once. This saves time since we do not have to repeat the same computations for every single batch.

Parameters:

  • Y (np.ndarray) – 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.

  • truncation (float) – How large a fraction of the singular values to keep in the inversion routine (if the inversion routine supports it). Must be a float in the range (0, 1]. A lower number means a more approximate answer and a slightly faster computation. The default is 0.99, which is recommended by Emerick in the reference paper.

  • overwrite (bool) – If False (the default), the input array will not be overwritten (mutated). If True, the method may overwrite the input array.

Returns:

The instance with mutated state.

Return type:

self

Notes

In the equation:

M + (delta M)

(delta D)^T [(delta D) (delta D)^T + alpha (N_e - 1) C_D]^(-1) (D_obs - D)

This method call corresponds to computing:

  • The next-to-last factors: (delta D)^T [(delta D) (delta D)^T +

    alpha (N_e - 1) C_D]^(-1)

  • The last factor: (D_obs - D)

The total internal storage is: 2 * ensemble_size * num_observations, since shapes are:

  • The next-to-last factors: (ensemble_size, num_observations)

  • The last factor: (num_observations, ensemble_size)

static three_over_sqrt_ensemble_members(corr_XY: ndarray[tuple[int, ...], dtype[floating]], ensemble_members_per_parameter: ndarray[tuple[int, ...], dtype[int64]] | int) → ndarray[tuple[int, ...], dtype[bool]][source]#

Use the correlation threshold 3 / sqrt(n). Note that unless the number of ensemble members is > 9, all responses are removed and no update happens at all.

This simple thresholding rule is equation (6) in the adaptive localization paper: http://doi.org/10.1175/MWR-D-24-0269.1

The rule is also mentioned in the paper:

Continuous Hyper-parameter OPtimization (CHOP) in an ensemble Kalman filter https://arxiv.org/abs/2206.03050