Repository: equinor/iterative_ensemble_smoother (language: python)

Issue:
Consider adding `ESMDA`
Could I think have a similar api to `SIES` but where we pass in a list of weights.
API design ideas
Now that I have read through the code a bit more thoroughly, here are some things I don't like about the current API:

- The `fit` and `update` methods might not need to be separate. `update` picks up where `fit` left off, and `update` does not even return anything. Is there any reason to separate them? 
- Some parameters, such as `observation_errors` (C_dd in paper) and `observation_values` (d in Paper) are sent in to `fit` in every iteration, but they never change, right? I think we should send them into the initialization instead.


To be conrete, I propose an API like this:

```python

# Every parameter that does not change in each Newton iteration (update() call) goes into the init()
model = SIES(X, C_dd, d, steplength_schedule, inversion="exact", seed=1)
# After model is initialized, E is drawn (once), and dimensions (N=realization, m=output size, n=input size)
# are inferred from X and d.

X_i = X.copy()
for iteration in range(10):
    Y_i = g(X_i)  # Evaluate the model g()
    X_i = model.update(Y_i)
    # Here X_i does not need to be explicitly passed, even if  section 2.4.3 is triggered,
    # since we can keep it internally in the model object from the previous iteration
```

Some comments:
- I want step-length to either be defined at initialization time, or at iteration time. Currently both is possible. Maybe there are some advantages, but it seems a bit redundant. You get more control if you have to input it at iteration time, because you can make a decision based on how the iterations are behaving. If you define it ahead of time you lose some control, but defining it at initialization time is more user-friendly and might be good enough.
- I want every argument that is one-time-only to go into init. Data too. Even if it's not what sklearn would've done. I think it's better that way because it separates what it updated and what is constant.
- I want to merge `fit` and `update`. I particularly don't like the name `fit`, because it has heavy connotations to sklearn and ML. So I would call the new method `update` or `iterate` or `apply_newton_step` or something.
- I want the `update` step to take in as few args as possible. Objects that can be held in state from previous iteration should be, objects that are constant go into initialization.

This breaks the typical sklearn api, where `init()` is only allowed to take in simple parameters (ints, strings, floats, etc) and not data. But the SIES is not an estimator, there's no fitting and predicting in the same ways as in an ML model. So mimicing sklearn is not as relevant as if we were implementing an ML algorithm.


Let me know yours thoughts. Maybe best to discuss over a call. I suspect there are some reasons/usecases why it's not as simple as I want it to be here. Let me know.


Discussion / hints:
Based on one of our previous discussions it sounds like the weights, or the inflation of the covariance, could be calculated inside of the fit/update algorithm. Do we have a reference paper to `ESMDA` or perhaps a png-picture of a pdf-slide that may guide our understanding?
The reference paper is this one: https://www.sciencedirect.com/science/article/pii/S0098300412000994

I think we need to pass in the weights as they can be set by users in ERT.
I can attempt this issue @dafeda 
Thanks, interesting suggestions. Let's discuss this over a call.
Leave this alone for now. We should think even more generally. What kind of API would facilitate many different types of data assimilation algorithms?

- Reporting, localization, etc

Fix patch:
diff --git a/src/iterative_ensemble_smoother/__init__.py b/src/iterative_ensemble_smoother/__init__.py
index 1f83e31f..1af2fdda 100644
--- a/src/iterative_ensemble_smoother/__init__.py
+++ b/src/iterative_ensemble_smoother/__init__.py
@@ -8,8 +8,10 @@
     ES,
     SIES,
 )
+from iterative_ensemble_smoother.esmda import ESMDA
 
 __all__ = [
     "ES",
     "SIES",
+    "ESMDA",
 ]
diff --git a/src/iterative_ensemble_smoother/esmda.py b/src/iterative_ensemble_smoother/esmda.py
new file mode 100644
index 00000000..311161f0
--- /dev/null
+++ b/src/iterative_ensemble_smoother/esmda.py
@@ -0,0 +1,258 @@
+"""
+Ensemble Smoother with Multiple Data Assimilation (ES-MDA)
+----------------------------------------------------------
+
+Implementation of the 2013 paper "Ensemble smoother with multiple data assimilation"
+
+References
+----------
+
+Emerick, A.A., Reynolds, A.C. History matching time-lapse seismic data using
+the ensemble Kalman filter with multiple data assimilations.
+Comput Geosci 16, 639–659 (2012). https://doi.org/10.1007/s10596-012-9275-5
+
+Alexandre A. Emerick, Albert C. Reynolds.
+Ensemble smoother with multiple data assimilation.
+Computers & Geosciences, Volume 55, 2013, Pages 3-15, ISSN 0098-3004,
+https://doi.org/10.1016/j.cageo.2012.03.011.
+
+https://gitlab.com/antoinecollet5/pyesmda
+
+https://helper.ipam.ucla.edu/publications/oilws3/oilws3_14147.pdf
+
+"""
+
+import numbers
+from typing import Optional, Union
+
+import numpy as np
+import numpy.typing as npt
+import scipy as sp  # type: ignore
+
+from iterative_ensemble_smoother.esmda_inversion import (
+    inversion_exact_cholesky,
+    inversion_subspace,
+    normalize_alpha,
+)
+
+
+class ESMDA:
+    # Available inversion methods. The inversion methods all compute
+    # C_MD @ (C_DD + alpha * C_D)^(-1)  @ (D - Y)
+    _inversion_methods = {
+        "exact": inversion_exact_cholesky,
+        "subspace": inversion_subspace,
+    }
+
+    def __init__(
+        self,
+        C_D: npt.NDArray[np.double],
+        observations: npt.NDArray[np.double],
+        alpha: Union[int, npt.NDArray[np.double]] = 5,
+        seed: Union[np.random._generator.Generator, int, None] = None,
+        inversion: str = "exact",
+    ) -> None:
+        """Initialize Ensemble Smoother with Multiple Data Assimilation (ES-MDA).
+
+        The implementation follows the 2013 paper by Emerick et al.
+
+        Parameters
+        ----------
+        C_D : np.ndarray
+            Covariance matrix of outputs of shape (num_outputs, num_outputs).
+            If a 1D array is passed, it represents a diagonal covariance matrix.
+        observations : np.ndarray
+            1D array of shape (num_inputs,) representing real-world observations.
+        alpha : int or 1D np.ndarray, optional
+            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.Generator, optional
+            A seed or numpy.random._generator.Generator used for random number
+            generation. The argument is passed to numpy.random.default_rng().
+            The default is None.
+        inversion : str, optional
+            Which inversion method to use. The default is "exact".
+
+        """
+        # Validate inputs
+        if not (isinstance(C_D, np.ndarray) and C_D.ndim in (1, 2)):
+            raise TypeError("Argument `C_D` must be a NumPy array of dimension 1 or 2.")
+
+        if C_D.ndim == 2 and C_D.shape[0] != C_D.shape[1]:
+            raise ValueError("Argument `C_D` must be square if it's 2D.")
+
+        if not (isinstance(observations, np.ndarray) and observations.ndim == 1):
+            raise TypeError("Argument `observations` must be a 1D NumPy array.")
+
+        if not observations.shape[0] == C_D.shape[0]:
+            raise ValueError("Shapes of `observations` and `C_D` must match.")
+
+        if not (
+            (isinstance(alpha, np.ndarray) and alpha.ndim == 1)
+            or isinstance(alpha, numbers.Integral)
+        ):
+            raise TypeError("Argument `alpha` must be an integer or a 1D NumPy array.")
+
+        if not (
+            isinstance(seed, (int, np.random._generator.Generator)) or seed is None
+        ):
+            raise TypeError(
+                "Argument `seed` must be an integer or numpy.random._generator.Generator."
+            )
+
+        if not isinstance(inversion, str):
+            raise TypeError(
+                f"Argument `inversion` must be a string in {tuple(self._inversion_methods.keys())}"
+            )
+        if inversion not in self._inversion_methods.keys():
+            raise ValueError(
+                f"Argument `inversion` must be a string in {tuple(self._inversion_methods.keys())}"
+            )
+
+        # Store data
+        self.observations = observations
+        self.iteration = 0
+        self.rng = np.random.default_rng(seed)
+        self.inversion = inversion
+
+        # Alpha can either be an integer (num iterations) or a list of weights
+        if isinstance(alpha, np.ndarray) and alpha.ndim == 1:
+            self.alpha = normalize_alpha(alpha)
+        elif isinstance(alpha, numbers.Integral):
+            self.alpha = normalize_alpha(np.ones(alpha))
+            assert np.allclose(self.alpha, normalize_alpha(self.alpha))
+        else:
+            raise TypeError("Alpha must be integer or 1D array.")
+
+        # Only compute the covariance factorization once
+        # If it's a full matrix, we gain speedup by only computing cholesky once
+        # If it's a diagonal, we gain speedup by never having to compute cholesky
+        num_outputs = C_D.shape[0]
+
+        if isinstance(C_D, np.ndarray) and C_D.ndim == 2:
+            self.C_D_L = sp.linalg.cholesky(C_D, lower=False)
+        elif isinstance(C_D, np.ndarray) and C_D.ndim == 1:
+            self.C_D_L = np.sqrt(C_D)
+        else:
+            raise TypeError("Argument `C_D` must be 1D or 2D array")
+
+        self.C_D = C_D
+        assert isinstance(self.C_D, np.ndarray) and self.C_D.ndim in (1, 2)
+
+    def num_assimilations(self) -> int:
+        return len(self.alpha)
+
+    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]:
+        """Assimilate data and return an updated ensemble X_posterior.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            2D array of shape (num_inputs, num_ensemble_members).
+        Y : np.ndarray
+            2D array of shape (num_ouputs, num_ensemble_members).
+        ensemble_mask : np.ndarray
+            1D boolean array of length `num_ensemble_members`, 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.
+        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
+        -------
+        X_posterior : np.ndarray
+            2D array of shape (num_inputs, num_ensemble_members).
+
+        """
+        if self.iteration >= self.num_assimilations():
+            raise Exception("No more assimilation steps to run.")
+
+        # Verify shapes
+        _, num_ensemble = X.shape
+        num_outputs, num_emsemble2 = Y.shape
+        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):
+            raise TypeError("Argument `Y` must be contain floats")
+
+        assert 0 < truncation <= 1.0
+
+        # Do not overwrite input arguments
+        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,
+        # 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)
+        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)
+        assert D.shape == (num_outputs, ensemble_mask.sum())
+
+        # Line 2 (c) in the description of ES-MDA in the 2013 Emerick paper
+        # Choose inversion method, e.g. 'exact'. The inversion method computes
+        # C_MD @ sp.linalg.inv(C_DD + C_D_alpha) @ (D - Y)
+        inversion_func = self._inversion_methods[self.inversion]
+
+        # Update and return
+        X[:, ensemble_mask] += inversion_func(
+            alpha=self.alpha[self.iteration],
+            C_D=self.C_D,
+            D=D,
+            Y=Y[:, ensemble_mask],
+            X=X[:, ensemble_mask],
+            truncation=truncation,
+        )
+
+        self.iteration += 1
+        return X
+
+
+if __name__ == "__main__":
+    import pytest
+
+    pytest.main(
+        args=[
+            __file__,
+            "--doctest-modules",
+            "-v",
+        ]
+    )
diff --git a/src/iterative_ensemble_smoother/esmda_inversion.py b/src/iterative_ensemble_smoother/esmda_inversion.py
new file mode 100644
index 00000000..78de6fac
--- /dev/null
+++ b/src/iterative_ensemble_smoother/esmda_inversion.py
@@ -0,0 +1,538 @@
+import numpy as np
+import numpy.typing as npt
+import scipy as sp  # type: ignore
+
+
+def empirical_covariance_upper(X: npt.NDArray[np.double]) -> npt.NDArray[np.double]:
+    """Compute the upper triangular part of the empirical covariance matrix X
+    with shape (num_variables, num_observations).
+
+    Examples
+    --------
+    >>> X = np.array([[-2.4, -0.3,  0.7,  0.2,  1.1],
+    ...               [-1.5,  0.4, -0.4, -0.9,  1. ],
+    ...               [-0.1, -0.4, -0. , -0.5,  1.1]])
+    >>> empirical_covariance_upper(X)
+    array([[1.873, 0.981, 0.371],
+           [0.   , 0.997, 0.392],
+           [0.   , 0.   , 0.407]])
+
+    Naive computation:
+
+    >>> empirical_cross_covariance(X, X)
+    array([[1.873, 0.981, 0.371],
+           [0.981, 0.997, 0.392],
+           [0.371, 0.392, 0.407]])
+    """
+    _, num_observations = X.shape
+    X = (X - np.mean(X, axis=1, keepdims=True)) / np.sqrt(num_observations - 1)
+    # https://www.math.utah.edu/software/lapack/lapack-blas/dsyrk.html
+    XXT: npt.NDArray[np.double] = sp.linalg.blas.dsyrk(alpha=1.0, a=X)
+    return XXT
+
+
+def empirical_cross_covariance(
+    X: npt.NDArray[np.double], Y: npt.NDArray[np.double]
+) -> npt.NDArray[np.double]:
+    """Both X and Y have shape (parameters, ensemble_size).
+
+    We use this function instead of np.cov to handle cross-correlation,
+    where X and Y have a different number of parameters.
+
+    Examples
+    --------
+    >>> X = np.array([[-2.4, -0.3,  0.7],
+    ...               [ 0.2,  1.1, -1.5]])
+    >>> Y = np.array([[ 0.4, -0.4, -0.9],
+    ...               [ 1. , -0.1, -0.4],
+    ...               [-0. , -0.5,  1.1],
+    ...               [-1.8, -1.1,  0.3]])
+    >>> empirical_cross_covariance(X, Y)
+    array([[-1.035     , -1.15833333,  0.66      ,  1.56333333],
+           [ 0.465     ,  0.36166667, -1.08      , -1.09666667]])
+
+    Verify against numpy.cov
+
+    >>> np.cov(X, rowvar=True, ddof=1)
+    array([[ 2.50333333, -0.99666667],
+           [-0.99666667,  1.74333333]])
+    >>> empirical_cross_covariance(X, X)
+    array([[ 2.50333333, -0.99666667],
+           [-0.99666667,  1.74333333]])
+
+    """
+    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)
+
+    # Compute outer product and divide
+    cov = X @ Y.T / (X.shape[1] - 1)
+    assert cov.shape == (X.shape[0], Y.shape[0])
+    return cov
+
+
+def normalize_alpha(alpha: npt.NDArray[np.double]) -> npt.NDArray[np.double]:
+    """Assure that sum_i (1/alpha_i) = 1.
+
+    This is Eqn (22) in the 2013 Emerick paper.
+
+    Examples
+    --------
+    >>> alpha = np.arange(10) + 1
+    >>> np.sum(1/normalize_alpha(alpha))
+    1.0
+    """
+    factor = np.sum(1 / alpha)
+    rescaled: npt.NDArray[np.double] = alpha * factor
+    return rescaled
+
+
+def singular_values_to_keep(
+    singular_values: npt.NDArray[np.double], truncation: float = 1.0
+) -> int:
+    """Find the index of the singular values to keep when truncating.
+
+    Examples
+    --------
+    >>> singular_values = np.array([3, 2, 1, 0, 0, 0])
+    >>> i = singular_values_to_keep(singular_values, truncation=1.0)
+    >>> singular_values[:i]
+    array([3, 2, 1])
+
+    >>> singular_values = np.array([4, 3, 2, 1])
+    >>> i = singular_values_to_keep(singular_values, truncation=1.0)
+    >>> singular_values[:i]
+    array([4, 3, 2, 1])
+
+    >>> singular_values = np.array([4, 3, 2, 1])
+    >>> singular_values_to_keep(singular_values, truncation=0.95)
+    4
+    >>> singular_values_to_keep(singular_values, truncation=0.9)
+    3
+    >>> singular_values_to_keep(singular_values, truncation=0.7)
+    2
+
+    """
+    assert np.all(
+        np.diff(singular_values) <= 0
+    ), "Singular values must be sorted decreasing"
+    assert 0 < truncation <= 1, "Threshold must be in range (0, 1]"
+    singular_values = np.array(singular_values, dtype=float)
+
+    # Take cumulative sum and normalize
+    cumsum = np.cumsum(singular_values)
+    cumsum /= cumsum[-1]
+    return int(np.searchsorted(cumsum, v=truncation, side="left") + 1)
+
+
+# =============================================================================
+# INVERSION FUNCTIONS
+# =============================================================================
+# All of these functions compute (exactly, or approximately), the product
+#
+#  C_MD @ inv(C_DD + alpha * C_D) @ (D - Y)
+#
+# where C_MD = empirical_cross_covariance(X, Y) = center(X) @ center(Y).T / (X.shape[1] - 1)
+#       C_DD = empirical_cross_covariance(Y, Y) = center(Y) @ center(Y).T / (Y.shape[1] - 1)
+#
+# The methods can be classified as
+#   - exact : with truncation=1.0, these methods compute the exact solution
+#   - exact : with truncation<1.0, these methods may approximate the solution
+#   - approximate: if ensemble_members <= num_outputs, then the solution is
+#                  always approximated, regardless of the truncation
+#   - approximate: if ensemble_members > num_outputs, then the solution is
+#                  exact when truncation is 1.0
+
+# Every inversion function has the form
+# inversion_<exact/approximate>_<name>
+
+
+def inversion_exact_naive(
+    *,
+    alpha: float,
+    C_D: npt.NDArray[np.double],
+    D: npt.NDArray[np.double],
+    Y: npt.NDArray[np.double],
+    X: npt.NDArray[np.double],
+    truncation: float = 1.0,
+) -> npt.NDArray[np.double]:
+    """Naive inversion, used for testing only.
+
+    Computes C_MD @ inv(C_DD + alpha * C_D) @ (D - Y) naively.
+    """
+    # Naive implementation of Equation (3) in Emerick (2013)
+    C_MD = empirical_cross_covariance(X, Y)
+    C_DD = empirical_cross_covariance(Y, Y)
+    C_D = np.diag(C_D) if C_D.ndim == 1 else C_D
+    return C_MD @ sp.linalg.inv(C_DD + alpha * C_D) @ (D - Y)  # type: ignore
+
+
+def inversion_exact_cholesky(
+    *,
+    alpha: float,
+    C_D: npt.NDArray[np.double],
+    D: npt.NDArray[np.double],
+    Y: npt.NDArray[np.double],
+    X: npt.NDArray[np.double],
+    truncation: float = 1.0,
+) -> 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.
+
+    The goal is to compute: C_MD @ inv(C_DD + alpha * C_D) @ (D - Y)
+
+    First we solve (C_DD + alpha * C_D) @ K = (D - Y) for K, so that
+    K = inv(C_DD + alpha * C_D) @ (D - Y), then we compute
+    C_MD @ K, but we don't explicitly form C_MD, since it might be more
+    efficient to perform the matrix products in another order.
+    """
+    C_DD = empirical_covariance_upper(Y)  # Only compute upper part
+
+    # Arguments for sp.linalg.solve
+    solver_kwargs = {
+        "overwrite_a": True,
+        "overwrite_b": True,
+        "assume_a": "pos",  # Assume positive definite matrix (use cholesky)
+        "lower": False,  # Only use the upper part while solving
+    }
+
+    # Compute K := sp.linalg.inv(C_DD + alpha * C_D) @ (D - Y)
+    if C_D.ndim == 2:
+        # C_D is a covariance matrix
+        C_DD += alpha * C_D  # Save memory by mutating
+        K = sp.linalg.solve(C_DD, D - Y, **solver_kwargs)
+    elif C_D.ndim == 1:
+        # C_D is an array, so add it to the diagonal without forming diag(C_D)
+        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)
+    Y = Y - np.mean(Y, axis=1, keepdims=True)
+    _, num_ensemble = X.shape
+
+    return np.linalg.multi_dot([X / (num_ensemble - 1), Y.T, K])  # type: ignore
+
+
+def inversion_exact_lstsq(
+    *,
+    alpha: float,
+    C_D: npt.NDArray[np.double],
+    D: npt.NDArray[np.double],
+    Y: npt.NDArray[np.double],
+    X: npt.NDArray[np.double],
+    truncation: float = 1.0,
+) -> npt.NDArray[np.double]:
+    """Computes inversion using least squares. While this method can deal with
+    rank-deficient C_D, it should not be used since it's very slow.
+    """
+    C_DD = empirical_cross_covariance(Y, Y)
+
+    # A covariance matrix was given
+    if C_D.ndim == 2:
+        lhs = C_DD + alpha * C_D
+    # A diagonal covariance matrix was given as a vector
+    else:
+        lhs = C_DD
+        lhs.flat[:: lhs.shape[0] + 1] += alpha * C_D
+
+    # K = lhs^-1 @ (D - Y)
+    # lhs @ K = (D - Y)
+    ans, *_ = sp.linalg.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
+
+
+def inversion_exact_rescaled(
+    *,
+    alpha: float,
+    C_D: npt.NDArray[np.double],
+    D: npt.NDArray[np.double],
+    Y: npt.NDArray[np.double],
+    X: npt.NDArray[np.double],
+    truncation: float = 1.0,
+) -> npt.NDArray[np.double]:
+    """Compute a rescaled inversion.
+
+    See Appendix A.1 in Emerick et al (2012) for details regarding this approach."""
+    C_DD = empirical_cross_covariance(Y, Y)
+
+    if C_D.ndim == 2:
+        # Eqn (57). Cholesky factorize the covariance matrix C_D
+        # TODO: here we compute the cholesky factor in every call, but C_D
+        # never changes. it would be better to compute it once
+        C_D_L = sp.linalg.cholesky(C_D, lower=True)  # Lower triangular cholesky
+        C_D_L_inv, _ = sp.linalg.lapack.dtrtri(
+            C_D_L, lower=1, overwrite_c=1
+        )  # Invert lower triangular using BLAS routine
+        C_D_L_inv /= np.sqrt(alpha)
+
+        # Eqn (59). Form C_tilde
+        # TODO: Use BLAS routine for triangular times dense matrix
+        # sp.linalg.blas.strmm(alpha=1, a=C_D_L_inv, b=C_DD, lower=1)
+        C_tilde = C_D_L_inv @ C_DD @ C_D_L_inv.T
+        C_tilde.flat[:: C_tilde.shape[0] + 1] += 1.0  # Add to diagonal
+
+    # When C_D is a diagonal covariance matrix, there is no need to perform
+    # the cholesky factorization
+    elif C_D.ndim == 1:
+        C_D_L_inv = 1 / np.sqrt(C_D * alpha)
+        C_tilde = (C_D_L_inv * (C_DD * C_D_L_inv).T).T
+        C_tilde.flat[:: C_tilde.shape[0] + 1] += 1.0  # Add to diagonal
+
+    # Eqn (60). Compute SVD, which is equivalent to taking eigendecomposition
+    # since C_tilde is PSD. Using eigh() is faster than svd().
+    # Note that svd() returns eigenvalues in decreasing order, while eigh()
+    # returns eigenvalues in increasing order.
+    # driver="evr" => fastest option
+    s, U = sp.linalg.eigh(C_tilde, driver="evr", overwrite_a=True)
+    # Truncate the SVD ( U_r @ np.diag(s_r) @ U_r.T == C_tilde )
+    # N_n is the number of observations
+    # N_e is the number of members in the ensemble
+    N_n, N_e = Y.shape
+
+    idx = singular_values_to_keep(s[::-1], truncation=truncation)
+    N_r = min(N_n, N_e - 1, idx)  # Number of values in SVD to keep
+    U_r, s_r = U[:, -N_r:], s[-N_r:]
+    # U_r @ np.diag(s_r) @ U_r.T == C_tilde
+
+    # Eqn (61). Compute symmetric term once first, then multiply together and
+    # finally multiply with (D - Y)
+    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)
+
+    return np.linalg.multi_dot([X_shift, Y_shift.T, term / s_r, term.T, (D - Y)])  # type: ignore
+
+
+def inversion_exact_subspace_woodbury(
+    *,
+    alpha: float,
+    C_D: npt.NDArray[np.double],
+    D: npt.NDArray[np.double],
+    Y: npt.NDArray[np.double],
+    X: npt.NDArray[np.double],
+    truncation: float = 1.0,
+) -> npt.NDArray[np.double]:
+    """Use the Woodbury lemma to compute the inversion.
+
+    This approach uses the Woodbury lemma to compute:
+        C_MD @ inv(C_DD + alpha * C_D) @ (D - Y)
+
+    Since C_DD = U @ U.T, where U := center(Y) / sqrt(N_e - 1), we can use:
+
+    (A + U @ U.T)^-1 = A^-1 - A^-1 @ U @ (1 + U.T @ A^-1 @ U )^-1 @ U.T @ A^-1
+
+    to compute inv(C_DD + alpha * C_D).
+
+
+
+    """
+
+    # Woodbury: (A + U @ U.T)^-1 = A^-1 - A^-1 @ U @ (1 + U.T @ A^-1 @ U )^-1 @ U.T @ A^-1
+
+    # Compute D_delta. N_n = number of outputs, N_e = number of ensemble members
+    N_n, N_e = Y.shape
+    D_delta = Y - np.mean(Y, axis=1, keepdims=True)  # Subtract average
+    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)
+
+    # A full covariance matrix was given
+    if C_D.ndim == 2:
+        # Invert C_D
+        # TODO: This inverse could be cached
+        C_D_inv = np.linalg.inv(C_D) / alpha
+
+        # Compute the center part of the rhs in woodburry
+        center = np.linalg.multi_dot([D_delta.T, C_D_inv, D_delta])
+        center.flat[:: center.shape[0] + 1] += 1.0  # Add to diagonal
+
+        # Compute the symmetric term of the rhs in woodbury
+        term = C_D_inv @ D_delta
+
+        # 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([X_shift, D_delta.T, inverted, (D - Y)])  # type: ignore
+
+    # A diagonal covariance matrix was given as a 1D array.
+    # Same computation as above, but exploit the diagonal structure
+    else:
+        C_D_inv = 1 / (C_D * alpha)  # Invert diagonal
+        center = np.linalg.multi_dot([D_delta.T * C_D_inv, D_delta])
+        center.flat[:: center.shape[0] + 1] += 1.0
+        UT_D = D_delta.T * C_D_inv
+        inverted = np.diag(C_D_inv) - np.linalg.multi_dot(
+            [UT_D.T, sp.linalg.inv(center), UT_D]
+        )
+        return np.linalg.multi_dot([X_shift, D_delta.T, inverted, (D - Y)])  # type: ignore
+
+
+def inversion_subspace(
+    *,
+    alpha: float,
+    C_D: npt.NDArray[np.double],
+    D: npt.NDArray[np.double],
+    Y: npt.NDArray[np.double],
+    X: npt.NDArray[np.double],
+    truncation: float = 1.0,
+) -> npt.NDArray[np.double]:
+    """See Appendix A.2 in Emerick et al (2012)
+
+    This is an approximate solution. The approximation is that when
+    U, w, V.T = svd(D_delta)
+    then we assume that U @ U.T = I.
+    This is not true in general, for instance:
+
+    >>> Y = np.array([[2, 0],
+    ...               [0, 0],
+    ...               [0, 0]])
+    >>> D_delta = Y - np.mean(Y, axis=1, keepdims=True)  # Subtract average
+    >>> D_delta
+    array([[ 1., -1.],
+           [ 0.,  0.],
+           [ 0.,  0.]])
+    >>> U, w, VT = sp.linalg.svd(D_delta)
+    >>> U, w
+    (array([[1., 0., 0.],
+           [0., 1., 0.],
+           [0., 0., 1.]]), array([1.41421356, 0.        ]))
+    >>> U[:, :1] @ np.diag(w[:1]) @ VT[:1, :] # Reconstruct D_Delta
+    array([[ 1., -1.],
+           [ 0.,  0.],
+           [ 0.,  0.]])
+    >>> U[:, :1] @ U[:, :1].T # But U_r @ U_r.T != I
+    array([[1., 0., 0.],
+           [0., 0., 0.],
+           [0., 0., 0.]])
+
+    """
+
+    # N_n is the number of observations
+    # N_e is the number of members in the ensemble
+    N_n, N_e = Y.shape
+
+    # Subtract the mean of every observation, see Eqn (67)
+    D_delta = Y - np.mean(Y, axis=1, keepdims=True)  # Subtract average
+
+    # Eqn (68)
+    # TODO: Approximately 50% of the time in the function is spent here
+    # consider using randomized svd for further speed gains
+    U, w, _ = sp.linalg.svd(D_delta, overwrite_a=True, full_matrices=False)
+
+    # Clip the singular value decomposition
+    idx = singular_values_to_keep(w, truncation=truncation)
+    N_r = min(N_n, N_e - 1, idx)  # Number of values in SVD to keep
+    U_r, w_r = U[:, :N_r], w[:N_r]
+
+    # Eqn (70). First compute the symmetric term, then form X
+    U_r_w_inv = U_r / w_r
+    if C_D.ndim == 1:
+        X1 = (N_e - 1) * np.linalg.multi_dot([U_r_w_inv.T * C_D * alpha, U_r_w_inv])
+    else:
+        X1 = (N_e - 1) * np.linalg.multi_dot([U_r_w_inv.T, alpha * C_D, U_r_w_inv])
+
+    # Eqn (72)
+    # Z, T, _ = sp.linalg.svd(X, overwrite_a=True, full_matrices=False)
+    # Compute SVD, which is equivalent to taking eigendecomposition
+    # since X is PSD. Using eigh() is faster than svd().
+    # Note that svd() returns eigenvalues in decreasing order, while eigh()
+    # returns eigenvalues in increasing order.
+    # driver="evr" => fastest option
+    T, Z = sp.linalg.eigh(X1, driver="evr", overwrite_a=True)
+
+    # Eqn (74).
+    # C^+ = (N_e - 1) hat{C}^+
+    #     = (N_e - 1) (U / w @ Z) * (1 / (1 + T)) (U / w @ Z)^T
+    #     = (N_e - 1) (term) * (1 / (1 + T)) (term)^T
+    # and finally we multiiply by (D - Y)
+    term = U_r_w_inv @ Z
+
+    # 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([X_shift, D_delta.T, (term / (1 + T)), term.T, (D - Y)])  # type: ignore
+
+
+def inversion_rescaled_subspace(
+    *,
+    alpha: float,
+    C_D: npt.NDArray[np.double],
+    D: npt.NDArray[np.double],
+    Y: npt.NDArray[np.double],
+    X: npt.NDArray[np.double],
+    truncation: float = 1.0,
+) -> npt.NDArray[np.double]:
+    """See Appendix A.2 in Emerick et al (2012)
+
+    Subspace inversion with rescaling.
+
+    """
+    # TODO: I don't understand why this approach is not approximate, when
+    # `inversion_subspace` is approximate
+
+    N_n, N_e = Y.shape
+    D_delta = Y - np.mean(Y, axis=1, keepdims=True)  # Subtract average
+
+    if C_D.ndim == 2:
+        # Eqn (76). Cholesky factorize the covariance matrix C_D
+        # TODO: here we compute the cholesky factor in every call, but C_D
+        # never changes. it would be better to compute it once
+        C_D_L = sp.linalg.cholesky(C_D * alpha, lower=True)  # Lower triangular cholesky
+        # Here C_D_L is C^{1/2} in equation (57)
+        # assert np.allclose(C_D_L @ C_D_L.T, C_D * alpha)
+        C_D_L_inv, _ = sp.linalg.lapack.dtrtri(
+            C_D_L, lower=1
+        )  # Invert lower triangular
+
+        # Use BLAS to compute product of lower triangular matrix C_D_L_inv and D_Delta
+        # This line is equal to C_D_L_inv @ D_delta
+        C_D_L_times_D_delta = sp.linalg.blas.dtrmm(
+            alpha=1.0, a=C_D_L_inv, b=D_delta, lower=1
+        )
+
+    else:
+        # Same as above, but C_D is a vector
+        C_D_L_inv = 1 / np.sqrt(alpha * C_D)  # Invert the Cholesky factor a diagonal
+        C_D_L_times_D_delta = (D_delta.T * C_D_L_inv).T
+
+    U, w, _ = sp.linalg.svd(C_D_L_times_D_delta, overwrite_a=True, full_matrices=False)
+    idx = singular_values_to_keep(w, truncation=truncation)
+
+    # assert np.allclose(VT @ VT.T, np.eye(VT.shape[0]))
+    N_r = min(N_n, N_e - 1, idx)  # Number of values in SVD to keep
+    U_r, w_r = U[:, :N_r], w[:N_r]
+
+    # Eqn (78) - taking into account that C_D_L_inv could be an array
+    term = C_D_L_inv.T @ (U_r / w_r) if C_D.ndim == 2 else ((U_r / w_r).T * C_D_L_inv).T
+    T_r = (N_e - 1) / w_r**2  # Equation (79)
+    diag = 1 / (1 + T_r)
+
+    # Compute C_MD
+    X_shift = X - np.mean(X, axis=1, keepdims=True)
+    return np.linalg.multi_dot([X_shift, D_delta.T, (term * diag), term.T, (D - Y)])  # type: ignore
+
+
+if __name__ == "__main__":
+    import pytest
+
+    # --durations=10  <- May be used to show potentially slow tests
+    pytest.main(
+        args=[
+            __file__,
+            "--doctest-modules",
+            "-v",
+            "-v",
+            # "-k simple",
+        ]
+    )


Test patch:
diff --git a/tests/test_esmda.py b/tests/test_esmda.py
new file mode 100644
index 00000000..52a59347
--- /dev/null
+++ b/tests/test_esmda.py
@@ -0,0 +1,384 @@
+"""
+Ensemble Smoother with Multiple Data Assimilation (ES-MDA)
+----------------------------------------------------------
+
+
+
+References
+----------
+
+Emerick, A.A., Reynolds, A.C. History matching time-lapse seismic data using
+the ensemble Kalman filter with multiple data assimilations.
+Comput Geosci 16, 639–659 (2012). https://doi.org/10.1007/s10596-012-9275-5
+
+Alexandre A. Emerick, Albert C. Reynolds.
+Ensemble smoother with multiple data assimilation.
+Computers & Geosciences, Volume 55, 2013, Pages 3-15, ISSN 0098-3004,
+https://doi.org/10.1016/j.cageo.2012.03.011.
+
+https://gitlab.com/antoinecollet5/pyesmda
+
+"""
+
+import numpy as np
+import pytest
+
+from iterative_ensemble_smoother.esmda import ESMDA
+from iterative_ensemble_smoother.esmda_inversion import empirical_cross_covariance
+
+
+class TestESMDA:
+    @pytest.mark.parametrize("seed", list(range(10)))
+    @pytest.mark.parametrize("inversion", ["exact", "subspace"])
+    def test_that_diagonal_covariance_gives_same_answer_as_dense(self, seed, inversion):
+        rng = np.random.default_rng(seed)
+
+        num_outputs = rng.choice([5, 10, 15, 25])
+        num_inputs = rng.choice([5, 10, 15, 25])
+        num_ensemble = rng.choice([5, 10, 15, 25])
+        alpha = rng.choice([1, 3, 5])
+
+        # Prior is N(0, 1)
+        X_prior = rng.normal(size=(num_inputs, num_ensemble))
+        Y_prior = rng.normal(size=(num_outputs, num_ensemble))
+
+        # Measurement errors
+        C_D = np.exp(rng.normal(size=num_outputs))
+
+        # Observations
+        observations = rng.normal(size=num_outputs, loc=1)
+
+        esmda = ESMDA(C_D, observations, alpha=alpha, seed=seed, inversion=inversion)
+        X_posterior1 = np.copy(X_prior)
+        for _ in range(esmda.num_assimilations()):
+            X_posterior1 = esmda.assimilate(X_posterior1, Y_prior)
+
+        esmda = ESMDA(
+            np.diag(C_D), observations, alpha=alpha, seed=seed, inversion=inversion
+        )
+        X_posterior2 = np.copy(X_prior)
+        for _ in range(esmda.num_assimilations()):
+            X_posterior2 = esmda.assimilate(X_posterior2, Y_prior)
+
+        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
+        C_D = 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(C_D, 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(C_D, 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],
+    )
+    def test_that_alpha_as_integer_and_array_returns_same_result(self, num_ensemble):
+        seed = num_ensemble
+        rng = np.random.default_rng(seed)
+
+        num_outputs = 2
+        num_inputs = 1
+
+        def g(x):
+            """Transform a single ensemble member."""
+            return np.array([np.sin(x / 2), x]) + 5
+
+        def G(X):
+            """Transform all ensemble members."""
+            return np.array([g(x_i) for x_i in X.T]).squeeze().T
+
+        # Prior is N(0, 1)
+        X_prior = rng.normal(size=(num_inputs, num_ensemble))
+
+        # Measurement errors
+        C_D = np.eye(num_outputs)
+
+        # The true inputs and observations, a result of running with N(1, 1)
+        X_true = rng.normal(size=(num_inputs,)) + 1
+        observations = G(X_true)
+
+        # Create ESMDA instance from an integer `alpha` and run it
+        esmda_integer = ESMDA(C_D, observations, alpha=5, seed=seed)
+        X_i_int = np.copy(X_prior)
+        for _ in range(esmda_integer.num_assimilations()):
+            X_i_int = esmda_integer.assimilate(X_i_int, G(X_i_int))
+
+        # Create another ESMDA instance from a vector `alpha` and run it
+        esmda_array = ESMDA(C_D, observations, alpha=np.ones(5), seed=seed)
+        X_i_array = np.copy(X_prior)
+        for _ in range(esmda_array.num_assimilations()):
+            X_i_array = esmda_array.assimilate(X_i_array, G(X_i_array))
+
+        # Exactly the same result with equal seeds
+        assert np.allclose(X_i_int, X_i_array)
+
+    # Linear transformation a x + b
+    @pytest.mark.parametrize("a", [-2, 1, 3, 9])
+    @pytest.mark.parametrize("b", [-2, 0, 1])
+    # Likelihood information
+    @pytest.mark.parametrize("X_true", [-2, 5, 10, 15])
+    @pytest.mark.parametrize("C_D", [1, 3, 9])
+    # Prior information
+    @pytest.mark.parametrize("mu", [-1, 0, 6, 12])
+    @pytest.mark.parametrize("S", [1, 3, 9])
+    def test_that_result_corresponds_with_theory_for_1D_gauss_linear_case(
+        self, mu, S, X_true, C_D, a, b
+    ):
+        # Here we test on a Guass-linear case.
+        # The section "2.3.3 Bayes’ theorem for Gaussian variables" in the book
+        # Pattern Recognition and Machine Learning by Bishop (2006) states that if
+        # p(x)   = N(x | mu, S)
+        # p(y|x) = N(y | a * x + b, C_D)
+        # then
+        # p(x | y) has
+        # covariance COV = inv(inv(S) + a * inv(C_D) * a)
+        # mean      MEAN = COV (a * inv(C_D) * (y - b) + inv(S) * mu)
+        def G(X):
+            """Transform all ensemble members."""
+            return a * X + b
+
+        # Analytical solution given by Bishop
+        inv = lambda x: 1 / x  # Matrix inversion for 1D matrix
+        COV = inv(inv(S) + a * inv(C_D) * a)
+        MEAN = COV * ((a * inv(C_D)) * (G(X_true) - b) + inv(S) * mu)
+
+        num_ensemble = 1_000
+        num_inputs = 1
+
+        # Create a random number generator
+        parameters = [mu, S, X_true, C_D, a, b]
+        seed = abs(sum(p_i * 10**i for i, p_i in enumerate(parameters)))
+        rng = np.random.default_rng(seed)
+
+        # Prior is p(x) ~ N(mu, S)
+        X_prior = rng.normal(size=(num_inputs, num_ensemble), loc=mu, scale=np.sqrt(S))
+
+        # Create ESMDA instance
+        esmda = ESMDA(np.atleast_2d([C_D]), G(np.atleast_1d(X_true)), alpha=1, seed=rng)
+        X_i = X_prior
+        for _ in range(esmda.num_assimilations()):
+            X_i = esmda.assimilate(X_i, G(X_i))
+
+        # Check that analytical solution is close to ESMDA posterior
+        # np.isclose(a, b) := abs(`a` - `b`) <= (`atol` + `rtol` * abs(`b`))
+        assert np.isclose(X_i.mean(), MEAN, rtol=0.1, atol=0.2)
+        assert np.isclose(X_i.var(ddof=1), COV, rtol=0.1, atol=0.1)
+
+    @pytest.mark.parametrize("seed", list(range(100)))
+    def test_that_result_corresponds_with_theory_for_gauss_linear_case(self, seed):
+        # Here we test on a Guass-linear case.
+        # The section "2.3.3 Bayes’ theorem for Gaussian variables" in the book
+        # Pattern Recognition and Machine Learning by Bishop (2006) states that if
+        # p(x)   = N(x | mu, C_M)
+        # p(y|x) = N(y | A x + b, C_D)
+        # then
+        # p(x | y) has
+        # covariance COV = inv(inv(C_M) + A^T @ inv(C_D) @ A)
+        # mean      MEAN = COV (A^T * inv(C_D) * (y - b) + inv(C_M) * mu)
+        #                = mu + C_M A^T (A C_M A^T + C_D)^-1 A (x - mu)
+
+        # Generate data
+        rng = np.random.default_rng(seed)
+        num_ensemble = 10_000
+        num_inputs, num_outputs = 5, 4
+
+        # Inputs
+        mu = rng.normal(size=num_inputs)
+        C_M_factor = rng.normal(size=(num_inputs, num_inputs))
+        C_M = C_M_factor.T @ C_M_factor + np.eye(num_inputs)
+
+        # Transformation
+        A = rng.normal(size=(num_outputs, num_inputs))
+        b = rng.normal(size=num_outputs)
+
+        def G(x):
+            if x.ndim == 1:
+                return A @ x + b
+            else:
+                return (b.reshape(-1, 1) + (A @ x)).squeeze()
+
+        assert np.allclose(
+            G(np.arange(num_inputs)), G(np.arange(num_inputs).reshape(-1, 1)).squeeze()
+        )
+
+        # Output covariance
+        C_D_factor = rng.normal(size=(num_outputs, num_outputs))
+        C_D = C_D_factor.T @ C_D_factor + np.eye(num_outputs)
+
+        # Analytical solution given by Bishop
+        X_true = mu + 10
+        inv = np.linalg.inv
+        COV = inv(inv(C_M) + A.T @ inv(C_D) @ A)
+        MEAN = COV @ ((A.T @ inv(C_D)) @ (G(X_true) - b) + inv(C_M) @ mu)
+        MEAN2 = mu + C_M @ A.T @ inv(A @ C_M @ A.T + C_D) @ A @ (X_true - mu)
+        assert np.allclose(MEAN, MEAN2)  # Both ways to compute mean is equivalent
+
+        # Prior is p(x) ~ N(mu, C_M)
+        X_prior = rng.multivariate_normal(mean=mu, cov=C_M, size=num_ensemble).T
+        assert G(X_prior).shape == (num_outputs, num_ensemble)
+
+        # Create ESMDA instance
+        esmda = ESMDA(C_D, G(X_true), alpha=1, seed=rng)
+        X_i = np.copy(X_prior)
+        for _ in range(esmda.num_assimilations()):
+            X_i = esmda.assimilate(X_i, G(X_i))
+
+        # Check that analytical solution is close to ESMDA posterior
+        relative_error_mean = np.linalg.norm(X_i.mean(axis=1) - MEAN) / np.linalg.norm(
+            MEAN
+        )
+        print("relative_error_mean", relative_error_mean)
+
+        assert relative_error_mean < 0.05
+
+        # TODO: The errors in the covariance matrix estimation is always pretty high
+        # why is this?
+        covariance = empirical_cross_covariance(X_i, X_i)
+        relative_error_covariance = np.linalg.norm((covariance - COV)) / np.linalg.norm(
+            COV
+        )
+        print("relative_error_covariance", relative_error_covariance)
+
+        assert relative_error_covariance < 1.2
+
+    # Likelihood information
+    @pytest.mark.parametrize("X_true", [5, 10])
+    @pytest.mark.parametrize("C_D", [1, 3])
+    # Prior information
+    @pytest.mark.parametrize("mu", [0, 2, 4])
+    @pytest.mark.parametrize("S", [1, 3, 9])
+    # Number of iterations
+    @pytest.mark.parametrize("alpha", [5, 10, 15])
+    def test_that_single_and_multiple_assimilations_achieve_same_result_for_1D_gauss_linear_case(
+        self, alpha, mu, S, X_true, C_D
+    ):
+        # Here we test on a Guass-linear case.
+        # The section "2.3.3 Bayes' theorem for Gaussian variables" in the book
+        # Pattern Recognition and Machine Learning by Bishop (2006) states that if
+        # p(x)   = N(x | mu, S)
+        # p(y|x) = N(y | a * x + b, C_D)
+        # then
+        # p(x | y) has
+        # covariance COV = inv(inv(S) + a * inv(C_D) * a)
+        # mean      MEAN = COV (a * inv(C_D) * (y - b) + inv(S) * mu)
+
+        def G(X):
+            """Transform all ensemble members."""
+            a, b = 1, 1
+            return a * X + b
+
+        num_ensemble = 10_000
+        num_inputs = 1
+
+        # Create a random number generator
+        parameters = [alpha, mu, S, X_true, C_D]
+        seed = abs(sum(p_i * 10**i for i, p_i in enumerate(parameters)))
+        rng = np.random.default_rng(seed)
+
+        # Prior is p(x) ~ N(mu, S)
+        X_prior = rng.normal(size=(num_inputs, num_ensemble), loc=mu, scale=np.sqrt(S))
+
+        # Create ESMDA instance with single iteration
+        esmda = ESMDA(np.diag([C_D]), G(np.atleast_1d(X_true)), alpha=1, seed=rng)
+        X_i_single = np.copy(X_prior)
+        for _ in range(esmda.num_assimilations()):
+            X_i_single = esmda.assimilate(X_i_single, G(X_i_single))
+
+        # Create ESMDA instance with multiple iterations
+        esmda = ESMDA(np.diag([C_D]), G(np.atleast_1d(X_true)), alpha=alpha, seed=rng)
+        X_i_multiple = np.copy(X_prior)
+        for _ in range(esmda.num_assimilations()):
+            X_i_multiple = esmda.assimilate(X_i_multiple, G(X_i_multiple))
+
+        # Check that summary statistics of solutions are close to each other
+        # np.isclose(a, b) := abs(`a` - `b`) <= (`atol` + `rtol` * abs(`b`))
+        assert np.isclose(X_i_single.mean(), X_i_multiple.mean(), rtol=0.05)
+        assert np.isclose(X_i_single.var(), X_i_multiple.var(), rtol=0.1)
+
+
+class TestESMDAMemory:
+    @pytest.fixture
+    def setup(self):
+        rng = np.random.default_rng(42)
+
+        num_outputs = 10_000
+        num_inputs = 1_000
+        num_ensemble = 100
+
+        # Prior is N(0, 1)
+        X_prior = rng.normal(size=(num_inputs, num_ensemble))
+        Y_prior = rng.normal(size=(num_outputs, num_ensemble))
+
+        # Measurement errors
+        C_D = np.exp(rng.normal(size=num_outputs))
+
+        # Observations
+        observations = rng.normal(size=num_outputs, loc=1)
+
+        return X_prior, Y_prior, C_D, observations
+
+    @pytest.mark.limit_memory("138 MB")
+    def test_ESMDA_memory_usage_subspace_inversion(self, setup):
+        # TODO: Currently this is a regression test. Work to improve memory usage.
+
+        X_prior, Y_prior, C_D, observations = setup
+
+        # Create ESMDA instance from an integer `alpha` and run it
+        esmda = ESMDA(C_D, observations, alpha=1, seed=1, inversion="subspace")
+
+        for _ in range(esmda.num_assimilations()):
+            esmda.assimilate(X_prior, Y_prior)
+
+
+if __name__ == "__main__":
+    import pytest
+
+    pytest.main(
+        args=[
+            __file__,
+            "-v",
+            "-k test_that_diagonal_covariance_gives_same_answer_as_dense",
+        ]
+    )
diff --git a/tests/test_esmda_inversion.py b/tests/test_esmda_inversion.py
new file mode 100644
index 00000000..f0a235f3
--- /dev/null
+++ b/tests/test_esmda_inversion.py
@@ -0,0 +1,356 @@
+import numpy as np
+import pytest
+
+from iterative_ensemble_smoother.esmda_inversion import (
+    inversion_exact_cholesky,
+    inversion_exact_lstsq,
+    inversion_exact_naive,
+    inversion_exact_rescaled,
+    inversion_rescaled_subspace,
+    inversion_subspace,
+    inversion_exact_subspace_woodbury,
+    normalize_alpha,
+)
+
+
+# If `truncation` is 1.0, then all of these produce the same result
+EXACT_INVERSIONS = [
+    inversion_exact_naive,
+    inversion_exact_cholesky,
+    inversion_exact_rescaled,
+    inversion_exact_lstsq,
+    inversion_exact_subspace_woodbury,
+    inversion_rescaled_subspace,
+]
+
+# Even when `truncation` is 1.0, these functions approximate the solution
+# when ensemble_members <= num_outputs
+APPROX_INVERSIONS = [
+    inversion_subspace,
+]
+
+# The functions that make use of the truncation parameter
+TRUNCATION_INVERSIONS = [
+    inversion_exact_rescaled,
+    inversion_subspace,
+    inversion_rescaled_subspace,
+]
+
+
+class TestEsmdaInversion:
+    @pytest.mark.parametrize("length", list(range(1, 101, 5)))
+    def test_that_the_sum_of_normalize_alpha_is_one(self, length):
+        # Generate random array
+        rng = np.random.default_rng(length)
+        alpha = np.exp(rng.normal(size=length))
+
+        # Test the defining property of the function
+        assert np.isclose(np.sum(1 / normalize_alpha(alpha)), 1)
+
+    @pytest.mark.parametrize(
+        "function",
+        EXACT_INVERSIONS,
+    )
+    def test_exact_inversion_on_a_simple_example(self, function):
+        """Test on one of the simplest cases imaginable.
+
+        If C_D = diag([1, 1, 1]) and Y = array([[2, 0],
+                                                [0, 0],
+                                                [0, 0]])
+
+        then C_DD = diag([2, 0, 0])
+
+        and inv(C_DD + C_D) = diag([1/3, 1, 1])
+        """
+
+        # Create positive symmetric definite covariance C_D
+        C_D = np.identity(3)
+
+        # Create observations
+        D = np.ones((3, 2))
+        Y = np.array([[2, 0], [0, 0], [0, 0]], dtype=float)
+        X = np.zeros((3, 2), dtype=float)
+        np.fill_diagonal(X, 1.0)
+
+        # inv(diag([3, 1, 1])) @ (D - Y)
+        K0 = function(alpha=1, C_D=C_D, D=D, Y=Y, X=X)
+        assert np.allclose(K0, np.array([[-1 / 3, 1 / 3], [1 / 3, -1 / 3], [0, 0]]))
+
+        # Same thing, but with a diagonal covariance represented as an array
+        K0 = function(alpha=1, C_D=np.diag(C_D), D=D, Y=Y, X=X)
+        assert np.allclose(K0, np.array([[-1 / 3, 1 / 3], [1 / 3, -1 / 3], [0, 0]]))
+
+    @pytest.mark.parametrize(
+        "function", EXACT_INVERSIONS + [inversion_rescaled_subspace]
+    )
+    @pytest.mark.parametrize("ensemble_members", [5, 10, 25, 50])
+    @pytest.mark.parametrize("num_outputs", [5, 10, 25, 50])
+    @pytest.mark.parametrize("num_inputs", [5, 10, 25, 50])
+    def test_that_exact_inverions_are_equal_with_few_ensemble_members(
+        self, function, ensemble_members, num_outputs, num_inputs
+    ):
+        """With few ensemble members, every exact inversion method should be equal."""
+        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)
+
+        # All non-subspace methods
+        K1 = inversion_exact_naive(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X)
+        K2 = function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X, truncation=1.0)
+        assert np.allclose(K1, K2)
+
+    @pytest.mark.parametrize("function", APPROX_INVERSIONS)
+    @pytest.mark.parametrize("num_outputs", [10])
+    def test_that_approximate_inversions_do_not_compute_exact_answer_with_few_ensemble_members(
+        self, function, num_outputs
+    ):
+        """With few ensemble members, the approximate methods should not return
+        the exact answer, even when truncation=1.0."""
+        np.random.seed(num_outputs)
+
+        num_inputs = num_outputs
+        ensemble_members = num_outputs // 2
+
+        # 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)
+
+        # All non-subspace methods
+        K1 = inversion_exact_naive(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X)
+        K2 = function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X, truncation=1.0)
+        assert not np.allclose(K1, K2)
+
+    @pytest.mark.parametrize(
+        "function",
+        EXACT_INVERSIONS + APPROX_INVERSIONS,
+    )
+    @pytest.mark.parametrize(
+        "num_outputs,num_emsemble",
+        [(10, 11), (10, 20), (10, 100), (100, 101), (100, 200), (100, 500)],
+    )
+    def test_that_all_inversions_all_equal_with_many_ensemble_members(
+        self, function, num_outputs, num_emsemble
+    ):
+        """As long as ensemble_members > num_outputs and truncation=1.0, every
+        inversion method should return the same result, whether it's an exact
+        method or an approximate method."""
+        assert num_emsemble > num_outputs
+        np.random.seed(num_outputs + num_emsemble)
+        num_inputs = num_outputs
+
+        # 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 = np.exp(np.random.randn())
+
+        # Create observations
+        D = np.random.randn(num_outputs, num_emsemble)
+        Y = np.random.randn(num_outputs, num_emsemble)
+        X = np.random.randn(num_inputs, num_emsemble)
+
+        K1 = inversion_exact_naive(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X, truncation=1.0)
+        K2 = function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X, truncation=1.0)
+
+        assert np.allclose(K1, K2)
+
+    @pytest.mark.parametrize("function", TRUNCATION_INVERSIONS)
+    @pytest.mark.parametrize("ensemble_over_outputs", [2, 1, 0.5])
+    def test_that_approximations_get_better_when_truncation_is_increased(
+        self, function, ensemble_over_outputs
+    ):
+        num_outputs = 16
+        num_emsemble = int(ensemble_over_outputs * num_outputs)
+        num_inputs = 8
+        np.random.seed(num_outputs + num_emsemble + 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 = np.exp(np.random.randn())
+
+        # Create observations
+        D = np.random.randn(num_outputs, num_emsemble)
+        Y = np.random.randn(num_outputs, num_emsemble)
+        X = np.random.randn(num_inputs, num_emsemble)
+
+        # Exact answer
+        exact = inversion_exact_naive(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X)
+
+        # Approximate answers
+        truncations = np.linspace(1e-8, 1, num=64)
+        approximations = [
+            function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X, truncation=trunaction)
+            for trunaction in truncations
+        ]
+
+        norms = np.array(
+            [np.linalg.norm(approx_i - exact) for approx_i in approximations]
+        )
+
+        # As we increase `truncation`, the difference between the approximation
+        # and the true answer should decrease
+        assert np.all(np.diff(norms) <= 0)
+
+        # Test that there is significant difference between endpoints
+        assert abs(norms[0] - norms[-1]) > 1
+
+    @pytest.mark.parametrize("ratio_ensemble_members_over_outputs", [0.5, 1, 2])
+    @pytest.mark.parametrize("num_outputs", [10, 50, 100])
+    @pytest.mark.parametrize(
+        "function",
+        EXACT_INVERSIONS + APPROX_INVERSIONS,
+    )
+    def test_that_inversion_methods_work_with_covariance_matrix_and_variance_vector(
+        self, ratio_ensemble_members_over_outputs, num_outputs, function
+    ):
+        """The result should be the same whether C_D is diagonal or not."""
+        ensemble_members = int(num_outputs / ratio_ensemble_members_over_outputs)
+        num_inputs = num_outputs
+        np.random.seed(num_outputs + ensemble_members)
+
+        # Diagonal covariance matrix
+        C_D_2D = np.diag(np.exp(np.random.randn(num_outputs)))
+        C_D_diag = np.diag(C_D_2D)
+
+        # Set alpha to something other than 1 to test that it works
+        alpha = np.exp(np.random.randn())
+
+        # 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)
+
+        assert C_D_2D.ndim == 2
+        assert C_D_diag.ndim == 1
+
+        result_diag = function(alpha=alpha, C_D=C_D_diag, D=D, Y=Y, X=X)
+        result_2D = function(alpha=alpha, C_D=C_D_2D, D=D, Y=Y, X=X)
+
+        assert np.allclose(result_diag, result_2D)
+
+    @pytest.mark.parametrize(
+        "function",
+        EXACT_INVERSIONS + APPROX_INVERSIONS,
+    )
+    def test_that_inversion_methods_do_do_not_mutate_input_args(self, function):
+        """Inversion functions should not mutate input arguments."""
+        num_outputs, num_inputs, ensemble_members = 100, 50, 10
+
+        np.random.seed(42)
+
+        # Diagonal covariance matrix
+        C_D = np.diag(np.exp(np.random.randn(num_outputs)))
+
+        # Set alpha to something other than 1 to test that it works
+        alpha = np.exp(np.random.randn())
+
+        # 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)
+
+        args = [alpha, C_D, D, Y, X]
+        args_copy = [np.copy(arg) for arg in args]
+
+        function(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X)
+
+        for arg, arg_copy in zip(args, args_copy):
+            assert np.allclose(arg, arg_copy)
+
+
+def test_timing(num_outputs=100, num_inputs=50, ensemble_members=25):
+    E = np.random.randn(num_outputs, num_outputs)
+    C_D = E.T @ E
+    C_D = np.diag(np.exp(np.random.randn(num_outputs)))  # Diagonal covariance matrix
+    C_D_diag = np.diag(C_D)
+    assert C_D_diag.ndim == 1
+    assert C_D.ndim == 2
+
+    # 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)
+
+    exact_inversion_funcs = [
+        inversion_exact_naive,
+        inversion_exact_cholesky,
+        # inversion_exact_rescaled,
+        # inversion_lstsq,
+    ]
+
+    from time import perf_counter
+
+    print("-" * 32)
+
+    for func in exact_inversion_funcs:
+        start_time = perf_counter()
+        result_matrix = func(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X)
+        elapsed_time = round(perf_counter() - start_time, 4)
+        print(f"Function: {func.__name__} on dense covariance: {elapsed_time} s")
+
+        start_time = perf_counter()
+        result_vector = func(alpha=alpha, C_D=C_D_diag, D=D, Y=Y, X=X)
+        elapsed_time = round(perf_counter() - start_time, 4)
+        print(f"Function: {func.__name__} on diagonal covariance: {elapsed_time} s")
+        assert np.allclose(result_matrix, result_vector)
+
+        print("-" * 32)
+
+    subspace_inversion_funcs = [
+        inversion_subspace,
+        inversion_exact_subspace_woodbury,
+        inversion_rescaled_subspace,
+    ]
+
+    from time import perf_counter
+
+    for func in subspace_inversion_funcs:
+        start_time = perf_counter()
+        result_matrix = func(alpha=alpha, C_D=C_D, D=D, Y=Y, X=X)
+        elapsed_time = round(perf_counter() - start_time, 4)
+        print(f"Function: {func.__name__} on dense covariance: {elapsed_time} s")
+
+        start_time = perf_counter()
+        result_vector = func(alpha=alpha, C_D=C_D_diag, D=D, Y=Y, X=X)
+        elapsed_time = round(perf_counter() - start_time, 4)
+        print(f"Function: {func.__name__} on diagonal covariance: {elapsed_time} s")
+        assert np.allclose(result_matrix, result_vector)
+
+        print("-" * 32)
+
+
+if __name__ == "__main__":
+    import pytest
+
+    pytest.main(
+        args=[
+            __file__,
+            "-v",
+            "-k test_that_exact_inverions_are_equal_with_few_ensemble_members",
+        ]
+    )