Source code for pyesmda.esmda

"""
Implement the ES-MDA algorithms.

@author: acollet
"""
from concurrent.futures import ProcessPoolExecutor
from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Union

import numpy as np
from scipy.sparse import csr_matrix

from pyesmda.utils import (
    NDArrayFloat,
    approximate_cov_mm,
    approximate_covariance_matrix_from_ensembles,
    check_nans_in_predictions,
    inflate_ensemble_around_its_mean,
)

# pylint: disable=C0103 # Does not conform to snake_case naming style


[docs]class ESMDA: r""" Ensemble Smoother with Multiple Data Assimilation. Implement the ES-MDA as proposed by Emerick, A. A. and A. C. Reynolds :cite:p:`emerickEnsembleSmootherMultiple2013, emerickHistoryMatchingProductionSeismic2013`. Attributes ---------- d_dim : int Number of observation values :math:`N_{obs}`, and consequently of predicted values. obs : NDArrayFloat Obsevrations vector with dimensions (:math:`N_{obs}`). cov_obs: NDArrayFloat Covariance matrix of observed data measurement errors with dimensions (:math:`N_{obs}`, :math:`N_{obs}`). Also denoted :math:`R`. d_obs_uc: NDArrayFloat Vectors of pertubed observations with dimension (:math:`N_{e}`, :math:`N_{obs}`). d_pred: NDArrayFloat Vectors of predicted values (one for each ensemble member) with dimensions (:math:`N_{e}`, :math:`N_{obs}`). d_history: List[NDArrayFloat] 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 : NDArrayFloat 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[NDArrayFloat] List of successive `m_prior`. cov_md: NDArrayFloat Cross-covariance matrix between the forecast state vector and predicted data. Dimensions are (:math:`N_{m}, N_{obs}`). cov_obs: csr_matrix Autocovariance matrix of predicted data. Dimensions are (:math:`N_{obs}, N_{obs}`). cov_mm: NDArrayFloat 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}`). cov_obs_inflation_factors : List[float] List of multiplication factor used to inflate the covariance matrix of the measurement errors. 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[Union[NDArrayFloat, csr_matrix]] 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}`). A sparse matrix format can be provided to save some memory. md_correlation_matrix : Optional[Union[NDArrayFloat, csr_matrix]] 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}`). . A sparse matrix format can be provided to save some memory. 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. batch_size: int Number of parameters that are assimilated at once. This option is available to overcome memory limitations when the number of parameters is large. In that case, the size of the covariance matrices tends to explode and the update step must be performed by chunks of parameters. is_parallel_analyse_step: bool, optional Whether to use parallel computing for the analyse step if the number of batch is above one. The default is True. n_batches: int Number of batches required during the update step. """ # pylint: disable=R0902 # Too many instance attributes __slots__: List[str] = [ "obs", "_cov_obs", "d_obs_uc", "d_pred", "d_history", "m_prior", "_m_bounds", "m_history", "cov_md", "cov_dd", "forward_model", "forward_model_args", "forward_model_kwargs", "_n_assimilations", "_assimilation_step", "_cov_obs_inflation_factors", "_cov_mm_inflation_factors", "_dd_correlation_matrix", "_md_correlation_matrix", "save_ensembles_history", "rng", "is_forecast_for_last_assimilation", "batch_size", "is_parallel_analyse_step", ]
[docs] def __init__( self, obs: NDArrayFloat, m_init: NDArrayFloat, cov_obs: Union[NDArrayFloat, csr_matrix], forward_model: Callable[..., NDArrayFloat], forward_model_args: Sequence[Any] = (), forward_model_kwargs: Optional[Dict[str, Any]] = None, n_assimilations: int = 4, cov_obs_inflation_factors: Optional[Sequence[float]] = None, cov_mm_inflation_factors: Optional[Sequence[float]] = None, dd_correlation_matrix: Optional[Union[NDArrayFloat, csr_matrix]] = None, md_correlation_matrix: Optional[Union[NDArrayFloat, csr_matrix]] = None, m_bounds: Optional[NDArrayFloat] = None, save_ensembles_history: bool = False, seed: Optional[int] = None, is_forecast_for_last_assimilation: bool = True, batch_size: int = 5000, is_parallel_analyse_step: bool = True, ) -> None: # pylint: disable=R0913 # Too many arguments # pylint: disable=R0914 # Too many local variables r"""Construct the instance. Parameters ---------- obs : NDArrayFloat Obsevrations vector with dimension :math:`N_{obs}`. m_init : NDArrayFloat Initial ensemble of parameters vector with dimensions (:math:`N_{e}`, :math:`N_{m}`). cov_obs: NDArrayFloat Covariance matrix of observed data measurement errors with dimensions (:math:`N_{obs}`, :math:`N_{obs}`). Also denoted :math:`R`. It can be a numpy array or a sparse matrix (scipy.linalg). 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. n_assimilations : int, optional Number of data assimilations (:math:`N_{a}`). The default is 4. cov_obs_inflation_factors : Optional[Sequence[float]] Multiplication factor used to inflate the covariance matrix of the measurement errors. Must match the number of data assimilations (:math:`N_{a}`). 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[NDArrayFloat] 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[NDArrayFloat] 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[NDArrayFloat], 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. batch_size: int Number of parameters that are assimilated at once. This option is available to overcome memory limitations when the number of parameters is large. In that case, the size of the covariance matrices tends to explode and the update step must be performed by chunks of parameters. The default is 5000. is_parallel_analyse_step: bool, optional Whether to use parallel computing for the analyse step if the number of batch is above one. It relies on `concurrent.futures` multiprocessing. The default is True. """ self.obs: NDArrayFloat = obs self.m_prior: NDArrayFloat = m_init self.save_ensembles_history: bool = save_ensembles_history self.m_history: list[NDArrayFloat] = [] self.d_history: list[NDArrayFloat] = [] self.d_pred: NDArrayFloat = np.zeros([self.n_ensemble, self.d_dim]) self.cov_obs = cov_obs self.d_obs_uc: NDArrayFloat = np.array([]) self.cov_md: NDArrayFloat = np.array([]) self.cov_dd: NDArrayFloat = 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._set_n_assimilations(n_assimilations) self._assimilation_step: int = 0 self.set_cov_obs_inflation_factors(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 self.batch_size = batch_size self.is_parallel_analyse_step = is_parallel_analyse_step
@property def n_assimilations(self) -> int: """Return the number of assimilations to perform. Read-only attribute.""" return self._n_assimilations
[docs] def _set_n_assimilations(self, n: int) -> None: """Set the number of assimilations to perform.""" try: if int(n) < 1: raise ValueError("The number of assimilations must be 1 or more.") if int(n) != float(n): raise TypeError() except TypeError as e: raise TypeError( "The number of assimilations must be a positive integer." ) from e self._n_assimilations = int(n)
@property def n_ensemble(self) -> int: """Return the number of ensemble members.""" return self.m_prior.shape[0] @property def m_dim(self) -> int: """Return the length of the parameters vector.""" return self.m_prior.shape[1] @property def d_dim(self) -> int: """Return the number of forecast data.""" return len(self.obs) @property def cov_obs(self) -> csr_matrix: """Get the observation errors covariance matrix.""" return self._cov_obs @cov_obs.setter def cov_obs(self, s: Union[NDArrayFloat, csr_matrix]) -> None: """Set the observation errors covariance matrix.""" if len(s.shape) != 2 or s.shape[0] != s.shape[1] or s.shape[0] != self.d_dim: raise ValueError( "cov_obs must be a 2D square matrix with " f"dimensions ({self.d_dim}, {self.d_dim})." ) self._cov_obs: csr_matrix = csr_matrix(s) @property def cov_mm(self) -> NDArrayFloat: r""" Get the estimated parameters autocovariance matrix. It is a read-only attribute. The covariance matrice :math:`C^{l}_{MM}` is approximated from the ensemble in the standard way of EnKF :cite:p:`evensenDataAssimilationEnsemble2007,aanonsenEnsembleKalmanFilter2009`: .. math:: C^{l}_{MM} = \frac{1}{N_{e} - 1} \sum_{j=1}^{N_{e}}\left(m^{l}_{j} - \overline{m^{l}}\right)\left(m^{l}_{j} - \overline{m^{l}} \right)^{T} with :math:`\overline{m^{l}}`, the parameters ensemble means, at iteration :math:`l`. """ return approximate_cov_mm(self.m_prior) @property def m_bounds(self) -> NDArrayFloat: """Get the parameter errors covariance matrix.""" return self._m_bounds @m_bounds.setter def m_bounds(self, mb: Optional[NDArrayFloat]) -> None: """Set the parameter errors covariance matrix.""" if mb is None: # In that case, create an array of nan. self._m_bounds: NDArrayFloat = np.empty([self.m_dim, 2], dtype=np.float64) self._m_bounds[:] = np.nan elif mb.shape[0] != self.m_dim: raise ValueError( f"m_bounds is of shape {mb.shape} while it " f"should be of shape ({self.m_dim}, 2)" ) else: self._m_bounds = mb @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
[docs] def set_cov_obs_inflation_factors(self, a: Optional[Sequence[float]]) -> None: """Set the inflation factors the covariance matrix of the measurement errors.""" if a is None: self._cov_obs_inflation_factors: List[float] = [ 1 / self.n_assimilations ] * self.n_assimilations elif len(a) != self.n_assimilations: raise ValueError( "The length of cov_obs_inflation_factors should match n_assimilations" ) else: self._cov_obs_inflation_factors = list(a)
@property def cov_mm_inflation_factors(self) -> List[float]: r""" Get the inlfation factors for the adjusted parameters covariance matrix. Covariance inflation is a method used to counteract the tendency of ensemble Kalman methods to underestimate the uncertainty because of either undersampling, inbreeding, or spurious correlations :cite:p:`todaroAdvancedTechniquesSolving2021`. The spread of the ensemble is artificially increased before the assimilation of the observations, according to scheme introduced by :cite:`andersonExploringNeedLocalization2007`: .. math:: m^{l}_{j} \leftarrow r^{l}\left(m^{l}_{j} - \frac{1}{N_{e}} \sum_{j}^{N_{e}}m^{l}_{j}\right) + \frac{1}{N_{e}}\sum_{j}^{N_{e}}m^{l}_{j} where :math:`r` is the inflation factor for the assimilation step :math:`l`. """ return list(self._cov_mm_inflation_factors) @cov_mm_inflation_factors.setter def cov_mm_inflation_factors(self, a: Optional[Sequence[float]]) -> None: """ Set the inflation factors the adjusted parameters covariance matrix. If no values have been provided by the user, the default is 1.0. at each iteration (no inflation). """ if a is None: self._cov_mm_inflation_factors: List[float] = [1.0] * self.n_assimilations elif len(a) != self.n_assimilations: raise ValueError( "The length of cov_mm_inflation_factors should match n_assimilations" ) else: self._cov_mm_inflation_factors = list(a) @property def dd_correlation_matrix(self) -> Optional[csr_matrix]: """Get the observations-observations localization matrix.""" return self._dd_correlation_matrix @dd_correlation_matrix.setter def dd_correlation_matrix( self, s: Optional[Union[NDArrayFloat, csr_matrix]] ) -> None: """Set the observations-observations localization matrix.""" if s is None: self._dd_correlation_matrix = None return if len(s.shape) != 2 or s.shape[0] != s.shape[1] or s.shape[0] != self.d_dim: raise ValueError( "dd_correlation_matrix must be a 2D square matrix with " f"dimensions ({self.d_dim}, {self.d_dim})." ) self._dd_correlation_matrix: Optional[csr_matrix] = csr_matrix(s) @property def md_correlation_matrix(self) -> Optional[csr_matrix]: """Get the parameters-observations localization matrix.""" return self._md_correlation_matrix @md_correlation_matrix.setter def md_correlation_matrix( self, s: Optional[Union[NDArrayFloat, csr_matrix]] ) -> None: """Set the parameters-observations localization matrix.""" if s is None: self._md_correlation_matrix = None return if len(s.shape) != 2: raise ValueError( "md_correlation_matrix must be a 2D matrix with " f"dimensions ({self.m_dim}, {self.d_dim})." ) if s.shape[0] != self.m_dim or s.shape[1] != self.d_dim: raise ValueError( "md_correlation_matrix must be a 2D matrix with " f"dimensions ({self.m_dim}, {self.d_dim})." ) self._md_correlation_matrix: Optional[csr_matrix] = csr_matrix(s) @property def n_batches(self) -> int: """Number of batch used in the optimization.""" return int(self.m_dim / self.batch_size) + 1
[docs] def solve(self) -> None: """Solve the optimization problem with ES-MDA algorithm.""" if self.n_batches == 1: self._solve() else: # assimilate chunk of parameters rather than everything all at once. self._solve_locally()
[docs] def _solve(self) -> None: """Solve the optimization problem with ES-MDA algorithm.""" if self.save_ensembles_history: self.m_history.append(self.m_prior) # save m_init for self._assimilation_step in range(self.n_assimilations): print(f"Assimilation # {self._assimilation_step + 1}") # inflating the covariance self.m_prior = self._apply_bounds( inflate_ensemble_around_its_mean( self.m_prior, self.cov_mm_inflation_factors[self._assimilation_step], ) ) self._forecast() self._pertrub(self.cov_obs_inflation_factors[self._assimilation_step]) self._approximate_covariance_matrices() # Update the prior parameter for next iteration self.m_prior = self._apply_bounds( self._analyse(self.cov_obs_inflation_factors[self._assimilation_step]) ) # Saving the parameters history if self.save_ensembles_history: self.m_history.append(self.m_prior) if self.is_forecast_for_last_assimilation: self._forecast()
[docs] def _solve_locally(self) -> None: """Solve the optimization problem with ES-MDA algorithm.""" if self.save_ensembles_history: self.m_history.append(self.m_prior) # save m_init for self._assimilation_step in range(self.n_assimilations): print(f"Assimilation # {self._assimilation_step + 1}") self._forecast() self._pertrub(self.cov_obs_inflation_factors[self._assimilation_step]) # covariance approximation dd self.cov_dd = approximate_covariance_matrix_from_ensembles( self.d_pred, self.d_pred ) # Spatial and temporal localization: obs - obs if self.dd_correlation_matrix is not None: self.cov_dd = self.dd_correlation_matrix.multiply(self.cov_dd).toarray() # Update the prior parameter for next iteration self.m_prior = self._apply_bounds( inflate_ensemble_around_its_mean( self._local_analyse( self.cov_obs_inflation_factors[self._assimilation_step] ), self.cov_mm_inflation_factors[self._assimilation_step], ) ) # Saving the parameters history if self.save_ensembles_history: self.m_history.append(self.m_prior) if self.is_forecast_for_last_assimilation: self._forecast()
[docs] def _forecast(self) -> None: r""" Forecast step of ES-MDA. Run the forward model from time zero until the end of the historical period from time zero until the end of the historical period to compute the vector of predicted data .. math:: d^{l}_{j}=g\left(m^{l}_{j}\right),\textrm{for }j=1,2,...,N_{e}, where :math:`g(ยท)` denotes the nonlinear observation model, i.e., :math:`d^{l}_{j}` is the :math:`N_{d}`-dimensional vector of predicted data obtained by running the forward model reservoir simulation with the model parameters given by the vector :math:`m^{l}_{j}` from time zero. Note that we use :math:`N_{d}` to denote the total number of measurements in the entire history. """ self.d_pred = self.forward_model( self.m_prior, *self.forward_model_args, **self.forward_model_kwargs ) if self.save_ensembles_history: self.d_history.append(self.d_pred) # Check if no nan values are found in the predictions. # If so, stop the assimilation check_nans_in_predictions(self.d_pred, self._assimilation_step)
[docs] def _pertrub(self, inflation_factor: float) -> None: r""" Perturbation of the observation vector step of ES-MDA. Perturb the vector of observations .. math:: d^{l}_{uc,j} = d_{obs} + \sqrt{\alpha_{l+1}}C_{D}^{1/2}Z_{d}, \textrm{for } j=1,2,...,N_{e}, where :math:`Z_{d} \sim \mathcal{N}(O, I_{N_{d}})`. Notes ----- To get reproducible behavior, use a seed when creating the ESMDA instance. """ self.d_obs_uc = np.zeros([self.n_ensemble, self.d_dim]) for i in range(self.d_dim): self.d_obs_uc[:, i] = self.obs[i] + np.sqrt( inflation_factor ) * self.rng.normal(0, np.abs(self.cov_obs.diagonal()[i]), self.n_ensemble)
[docs] def _approximate_covariance_matrices(self) -> None: r""" Approximate the covariance matrices. The covariance matrices :math:`C^{l}_{MD}` and :math:`C^{l}_{DD}` are approximated from the ensemble in the standard way of EnKF :cite:p:`evensenDataAssimilationEnsemble2007,aanonsenEnsembleKalmanFilter2009`: .. math:: C^{l}_{MD} = \frac{1}{N_{e} - 1} \sum_{j=1}^{N_{e}}\left(m^{l}_{j} - \overline{m^{l}}\right)\left(d^{l}_{j} - \overline{d^{l}} \right)^{T} .. math:: C^{l}_{DD} = \frac{1}{N_{e} - 1} \sum_{j=1}^{N_{e}}\left(d^{l}_{j} -\overline{d^{l}} \right)\left(d^{l}_{j} - \overline{d^{l}} \right)^{T} with :math:`\overline{m^{l}}` and :math:`\overline{d^{l}}`, the the ensemble means, at iteration :math:`l`, of parameters and predictions, respectively. If defined by the user, covariances localization is applied by element-wise multiplication (Schur product or Hadamard product) of the original covariance matrix and a distance dependent correlation function that smoothly reduces the correlations between points for increasing distances and cuts off long-range correlations above a specific distance: .. math:: \tilde{C}^{l}_{MD} = \rho_{MD} \odot C^{l}_{MD} .. math:: \tilde{C}^{l}_{DD} = \rho_{DD} \odot C^{l}_{DD} with :math:`\odot` the element wise multiplication. """ self.cov_md = approximate_covariance_matrix_from_ensembles( self.m_prior, self.d_pred ) self.cov_dd = approximate_covariance_matrix_from_ensembles( self.d_pred, self.d_pred ) # Spatial and temporal localization: obs - obs if self.dd_correlation_matrix is not None: self.cov_dd = self.dd_correlation_matrix.multiply(self.cov_dd).toarray() # Spatial and temporal localization: parameters - obs if self.md_correlation_matrix is not None: self.cov_md = self.md_correlation_matrix.multiply(self.cov_md).toarray()
[docs] def _analyse(self, inflation_factor: float) -> NDArrayFloat: r""" Analysis step of the ES-MDA. Update the vector of model parameters using .. math:: m^{l+1}_{j} = m^{l}_{j} + C^{l}_{MD}\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right)^{-1} \left(d^{l}_{uc,j} - d^{l}_{j} \right), \textrm{for } j=1,2,...,N_{e}. Notes ----- To avoid the inversion of :math:`\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right)`, the product :math:`\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right) ^{-1} \left(d^{l}_{uc,j} - d^{l}_{j} \right)` is solved linearly as :math:`A^{-1}b = x` which is equivalent to solve :math:`Ax = b`. """ # predicted parameters return ( self.m_prior + ( self.cov_md
[docs] @ np.linalg.solve( self.cov_dd + inflation_factor * self.cov_obs, (self.d_obs_uc - self.d_pred).T, ) ).T )
def _local_analyse(self, inflation_factor: float) -> NDArrayFloat: r""" Analysis step of the ES-MDA. Update the vector of model parameters using .. math:: m^{l+1}_{j} = m^{l}_{j} + C^{l}_{MD}\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right)^{-1} \left(d^{l}_{uc,j} - d^{l}_{j} \right), \textrm{for } j=1,2,...,N_{e}. Notes ----- To avoid the inversion of :math:`\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right)`, the product :math:`\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right) ^{-1} \left(d^{l}_{uc,j} - d^{l}_{j} \right)` is solved linearly as :math:`A^{-1}b = x` which is equivalent to solve :math:`Ax = b`. """ m_pred: NDArrayFloat = np.zeros(self.m_prior.shape) if self.is_parallel_analyse_step: with ProcessPoolExecutor() as executor: results: Iterator[NDArrayFloat] = executor.map( self._get_batch_m_update, range(self.n_batches), [inflation_factor] * self.n_batches, ) for index, res in enumerate(results): _slice = slice( index * self.batch_size, (index + 1) * self.batch_size ) m_pred[:, _slice] = res # self.simu_n += n_ensemble else: for index in range(self.n_batches): _slice = slice(index * self.batch_size, (index + 1) * self.batch_size) m_pred[:, _slice] = self._get_batch_m_update(index, inflation_factor) return m_pred
def _get_batch_m_update(self, index: int, inflation_factor: float) -> NDArrayFloat: _slice = slice(index * self.batch_size, (index + 1) * self.batch_size) batch_cov_md = approximate_covariance_matrix_from_ensembles( self.m_prior[:, _slice], self.d_pred ) # apply localization if self.md_correlation_matrix is not None: batch_cov_md = self.md_correlation_matrix[_slice, :].multiply(batch_cov_md) return ( self.m_prior[:, _slice] + ( batch_cov_md
[docs] @ np.linalg.solve( self.cov_dd + inflation_factor * self.cov_obs, (self.d_obs_uc - self.d_pred).T, ) ).T ) def _apply_bounds(self, m_pred: NDArrayFloat) -> NDArrayFloat: """Apply bounds constraints to the adjusted parameters.""" m_pred = np.where( m_pred < self.m_bounds[:, 0], self.m_bounds[:, 0], m_pred ) # lower bounds m_pred = np.where( m_pred > self.m_bounds[:, 1], self.m_bounds[:, 1], m_pred ) # upper bounds return m_pred