"""
Implement the ES-MDA-RS algorithms.
@author: acollet
"""
from typing import Any, Callable, Dict, List, Optional, Sequence
import numpy as np
import numpy.typing as npt
from pyesmda.esmda import ESMDA
from pyesmda.utils import compute_ensemble_average_normalized_objective_function
# pylint: disable=C0103 # Does not conform to snake_case naming style
[docs]class ESMDA_RS(ESMDA):
r"""
Restricted Step Ensemble Smoother with Multiple Data Assimilation.
Implement an adaptative version of the original ES-MDA algorithm proposed by
Emerick, A. A. and A. C. Reynolds
:cite:p:`emerickEnsembleSmootherMultiple2013,
emerickHistoryMatchingProductionSeismic2013`. This adaptative version introduced by
:cite:p:`leAdaptiveEnsembleSmoother2016` provides an automatic procedure for
choosing the inflation factor for the next data-assimilation step adaptively
as the history match proceeds. The procedure also decides when to stop,
i.e. the number of assimilation, which is no longer a user input.
Attributes
----------
d_dim : int
Number of observation values :math:`N_{obs}`, and consequently of
predicted values.
obs : npt.NDArray[np.float64]
Obsevrations vector with dimensions (:math:`N_{obs}`).
cov_obs: npt.NDArray[np.float64]
Covariance matrix of observed data measurement errors with dimensions
(:math:`N_{obs}`, :math:`N_{obs}`). Also denoted :math:`R`.
std_m_prior: npt.NDArray[np.float64]
Vector of a priori standard deviation :math:`sigma` of the estimated
parameter. The expected dimension is (:math:`N_{m}`).
It is the diagonal of :math:`C_{M}`.
d_obs_uc: npt.NDArray[np.float64]
Vectors of pertubed observations with dimension
(:math:`N_{e}`, :math:`N_{obs}`).
d_pred: npt.NDArray[np.float64]
Vectors of predicted values (one for each ensemble member)
with dimensions (:math:`N_{e}`, :math:`N_{obs}`).
d_history: List[npt.NDArray[np.float64]]
List of vectors of predicted values obtained at each assimilation step.
m_prior:
Vectors of parameter values (one vector for each ensemble member) used in the
last assimilation step. Dimensions are (:math:`N_{e}`, :math:`N_{m}`).
m_bounds : npt.NDArray[np.float64]
Lower and upper bounds for the :math:`N_{m}` parameter values.
Expected dimensions are (:math:`N_{m}`, 2) with lower bounds on the first
column and upper on the second one.
m_history: List[npt.NDArray[np.float64]]
List of successive `m_prior`.
cov_md: npt.NDArray[np.float64]
Cross-covariance matrix between the forecast state vector and predicted data.
Dimensions are (:math:`N_{m}, N_{obs}`).
cov_dd: npt.NDArray[np.float64]
Autocovariance matrix of predicted data.
Dimensions are (:math:`N_{obs}, N_{obs}`).
cov_mm: npt.NDArray[np.float64]
Autocovariance matrix of estimated parameters.
Dimensions are (:math:`N_{m}, N_{m}`).
forward_model: callable
Function calling the non-linear observation model (forward model)
for all ensemble members and returning the predicted data for
each ensemble member.
forward_model_args: Tuple[Any]
Additional args for the callable forward_model.
forward_model_kwargs: Dict[str, Any]
Additional kwargs for the callable forward_model.
n_assimilations : int
Number of data assimilations (:math:`N_{a}`) performed.
Automatically determined. Initially at 0.
cov_obs_inflation_factors : List[float]
List of multiplication factor used to inflate the covariance matrix of the
measurement errors. It is determined automatically.
cov_mm_inflation_factors: List[float]
List of factors used to inflate the adjusted parameters covariance among
iterations:
Each realization of the ensemble at the end of each update step i,
is linearly inflated around its mean.
See :cite:p:`andersonExploringNeedLocalization2007`.
dd_correlation_matrix : Optional[npt.NDArray[np.float64]]
Correlation matrix based on spatial and temporal distances between
observations and observations :math:`\rho_{DD}`. It is used to localize the
autocovariance matrix of predicted data by applying an elementwise
multiplication by this matrix.
Expected dimensions are (:math:`N_{obs}`, :math:`N_{obs}`).
md_correlation_matrix : Optional[npt.NDArray[np.float64]]
Correlation matrix based on spatial and temporal distances between
parameters and observations :math:`\rho_{MD}`. It is used to localize the
cross-covariance matrix between the forecast state vector (parameters)
and predicted data by applying an elementwise
multiplication by this matrix.
Expected dimensions are (:math:`N_{m}`, :math:`N_{obs}`).
save_ensembles_history: bool
Whether to save the history predictions and parameters over the assimilations.
rng: np.random.Generator
The random number generator used in the predictions perturbation step.
is_forecast_for_last_assimilation: bool
Whether to compute the predictions for the ensemble obtained at the
last assimilation step.
"""
# pylint: disable=R0902 # Too many instance attributes
__slots__: List[str] = ["std_m_prior"]
[docs] def __init__(
self,
obs: npt.NDArray[np.float64],
m_init: npt.NDArray[np.float64],
cov_obs: npt.NDArray[np.float64],
std_m_prior: npt.NDArray[np.float64],
forward_model: Callable[..., npt.NDArray[np.float64]],
forward_model_args: Sequence[Any] = (),
forward_model_kwargs: Optional[Dict[str, Any]] = None,
cov_mm_inflation_factors: Optional[Sequence[float]] = None,
dd_correlation_matrix: Optional[npt.NDArray[np.float64]] = None,
md_correlation_matrix: Optional[npt.NDArray[np.float64]] = None,
m_bounds: Optional[npt.NDArray[np.float64]] = None,
save_ensembles_history: bool = False,
seed: Optional[int] = None,
is_forecast_for_last_assimilation: bool = True,
) -> None:
# pylint: disable=R0913 # Too many arguments
# pylint: disable=R0914 # Too many local variables
r"""Construct the instance.
Parameters
----------
obs : npt.NDArray[np.float64]
Obsevrations vector with dimension :math:`N_{obs}`.
m_init : npt.NDArray[np.float64]
Initial ensemble of parameters vector with dimensions
(:math:`N_{e}`, :math:`N_{m}`).
cov_obs: npt.NDArray[np.float64]
Covariance matrix of observed data measurement errors with dimensions
(:math:`N_{obs}`, :math:`N_{obs}`). Also denoted :math:`R`.
std_m_prior: npt.NDArray[np.float64]
Vector of a priori standard deviation :math:`sigma` of the estimated
parameter. The expected dimension is (:math:`N_{m}`).
It is the diagonal of :math:`C_{M}`.
forward_model: callable
Function calling the non-linear observation model (forward model)
for all ensemble members and returning the predicted data for
each ensemble member.
forward_model_args: Optional[Tuple[Any]]
Additional args for the callable forward_model. The default is None.
forward_model_kwargs: Optional[Dict[str, Any]]
Additional kwargs for the callable forward_model. The default is None.
cov_mm_inflation_factors: Optional[Sequence[float]]
List of factors used to inflate the adjusted parameters covariance
among iterations:
Each realization of the ensemble at the end of each update step i,
is linearly inflated around its mean.
Must match the number of data assimilations (:math:`N_{a}`).
See :cite:p:`andersonExploringNeedLocalization2007`.
If None, the default is 1.0. at each iteration (no inflation).
The default is None.
dd_correlation_matrix : Optional[npt.NDArray[np.float64]]
Correlation matrix based on spatial and temporal distances between
observations and observations :math:`\rho_{DD}`. It is used to localize the
autocovariance matrix of predicted data by applying an elementwise
multiplication by this matrix.
Expected dimensions are (:math:`N_{obs}`, :math:`N_{obs}`).
The default is None.
md_correlation_matrix : Optional[npt.NDArray[np.float64]]
Correlation matrix based on spatial and temporal distances between
parameters and observations :math:`\rho_{MD}`. It is used to localize the
cross-covariance matrix between the forecast state vector (parameters)
and predicted data by applying an elementwise
multiplication by this matrix.
Expected dimensions are (:math:`N_{m}`, :math:`N_{obs}`).
The default is None.
m_bounds : Optional[npt.NDArray[np.float64]], optional
Lower and upper bounds for the :math:`N_{m}` parameter values.
Expected dimensions are (:math:`N_{m}`, 2) with lower bounds on the first
column and upper on the second one. The default is None.
save_ensembles_history: bool, optional
Whether to save the history predictions and parameters over
the assimilations. The default is False.
seed: Optional[int]
Seed for the white noise generator used in the perturbation step.
If None, the default :func:`numpy.random.default_rng()` is used.
The default is None.
is_forecast_for_last_assimilation: bool, optional
Whether to compute the predictions for the ensemble obtained at the
last assimilation step. The default is True.
"""
self.obs: npt.NDArray[np.float64] = obs
self.m_prior: npt.NDArray[np.float64] = m_init
self.save_ensembles_history: bool = save_ensembles_history
self.m_history: list[npt.NDArray[np.float64]] = []
self.d_history: list[npt.NDArray[np.float64]] = []
self.d_pred: npt.NDArray[np.float64] = np.zeros([self.n_ensemble, self.d_dim])
self.cov_obs = cov_obs
self.std_m_prior: npt.NDArray[np.float64] = std_m_prior
self.d_obs_uc: npt.NDArray[np.float64] = np.array([])
self.cov_md: npt.NDArray[np.float64] = np.array([])
self.cov_dd: npt.NDArray[np.float64] = np.array([])
self.forward_model: Callable = forward_model
self.forward_model_args: Sequence[Any] = forward_model_args
if forward_model_kwargs is None:
forward_model_kwargs = {}
self.forward_model_kwargs: Dict[str, Any] = forward_model_kwargs
self._assimilation_step: int = 0
self.cov_obs_inflation_factors = []
self.cov_mm_inflation_factors = cov_mm_inflation_factors
self.dd_correlation_matrix = dd_correlation_matrix
self.md_correlation_matrix = md_correlation_matrix
self.m_bounds = m_bounds
self.rng: np.random.Generator = np.random.default_rng(seed)
self.is_forecast_for_last_assimilation = is_forecast_for_last_assimilation
@property
def n_assimilations(self) -> int:
"""Get the number of assimilations performed. Read-only."""
return self._assimilation_step
@property
def cov_obs_inflation_factors(self) -> List[float]:
r"""
Get the inlfation factors for the covariance matrix of the measurement errors.
Single and multiple data assimilation are equivalent for the
linear-Gaussian case as long as the factor :math:`\alpha_{l}` used to
inflate the covariance matrix of the measurement errors satisfy the
following condition:
.. math::
\sum_{l=1}^{N_{a}} \frac{1}{\alpha_{l}} = 1
In practise, :math:`\alpha_{l} = N_{a}` is a good choice
:cite:p:`emerickEnsembleSmootherMultiple2013`.
"""
return self._cov_obs_inflation_factors
@cov_obs_inflation_factors.setter
def cov_obs_inflation_factors(self, a: List[float]) -> None:
"""Set the inflation factors the covariance matrix of the measurement errors."""
self._cov_obs_inflation_factors = a
[docs] def solve(self) -> None:
"""Solve the optimization problem with ES-MDA-RS algorithm."""
if self.save_ensembles_history:
self.m_history.append(self.m_prior) # save m_init
current_inflation_factor: float = 10.0 # to initiate the while
while not self._is_unity_reached(current_inflation_factor):
self._assimilation_step += 1
print(f"Assimilation # {self._assimilation_step}")
self._forecast()
# Divide per 2, because it is multiplied by 2 as the beginning
# of the second while loop
current_inflation_factor: float = (
self._compute_initial_inflation_factor() / 2
)
is_valid_parameter_change: bool = False
while not is_valid_parameter_change:
current_inflation_factor *= 2 # double the inflation (dumping) factor
self._pertrub(current_inflation_factor)
self._approximate_covariance_matrices()
m_pred = self._apply_bounds(self._analyse(current_inflation_factor))
is_valid_parameter_change: bool = self._is_valid_parameter_change(
m_pred
)
# If the criteria is reached -> Get exactly one for the sum
if self._is_unity_reached(current_inflation_factor):
current_inflation_factor = 1 / (
1 - np.sum([1 / a for a in self.cov_obs_inflation_factors])
)
self._pertrub(current_inflation_factor)
self._approximate_covariance_matrices()
m_pred = self._analyse(current_inflation_factor)
is_valid_parameter_change: bool = self._is_valid_parameter_change(
m_pred
)
self.cov_obs_inflation_factors.append(current_inflation_factor)
print(f"- Inflation factor = {current_inflation_factor:.3f}")
# Update the prior parameter for next iteration
self.m_prior = m_pred
# Saving the parameters history
if self.save_ensembles_history:
self.m_history.append(m_pred)
# Last assimilation
if self.is_forecast_for_last_assimilation:
self._forecast()
[docs] def _compute_initial_inflation_factor(self) -> float:
r"""Compute the :math:`\alpha_{l}` inflation (dumping) factor."""
return 0.25 * compute_ensemble_average_normalized_objective_function(
self.d_pred, self.obs, self.cov_obs
)
[docs] def _is_unity_reached(self, current_inflation_factor: float) -> bool:
"""
Whether the sum of the inverse inflation factors is above one.
It includes all factors up to the current iteration.
Parameters
----------
current_inflation_factor: float
Multiplication factor used to inflate the covariance matrix of the
measurement errors for the current (last) iteration.
"""
return (
np.sum([1 / a for a in self.cov_obs_inflation_factors])
+ 1 / current_inflation_factor
>= 1
)
[docs] def _is_valid_parameter_change(self, m_pred: npt.NDArray[np.float64]) -> bool:
"""Check if all change residuals are below 2 sigma.
Parameters
----------
m_pred : npt.NDArray[np.float64]
_description_
Returns
-------
bool
_description_
"""
def is_lower(residuals) -> bool:
return np.all(residuals < 2 * self.std_m_prior)
return np.all(list(map(is_lower, np.abs(m_pred - self.m_prior))))