Basic 2 parameter estimation

Here is a simple example where amplitude and change factor parameters of n exponential function are estimated:

import numpy as np
from pyesmda import ESMDA, ESMDA_RS, ESMDA_DMC, ESMDAInversionType, FixedLocalization
import logging
import covmats

# set logging level
logging.basicConfig(level=logging.INFO)


def exponential(p, x):
    """
    Simple exponential function with an amplitude and change factor.

    Parameters
    ----------
    p : tuple, list
        Parameters vector: amplitude i.e. initial value and change factor.
    x : np.array
        Independent variable (e.g. time).

    Returns
    -------
    np.array
        Result.

    """
    return p[0] * np.exp(x * p[1])


def forward_model(m_ensemble, x):
    """
    Wrap the non-linear observation model (forward model).

    Function calling the non-linear observation model (forward model).
    for all ensemble members and returning the predicted data for
    each ensemble member.

    Parameters
    ----------
    m_ensemble : np.array
        Initial ensemble of N_{e} parameters vector..
    x : np.array
        Independent variable (e.g. time).

    Returns
    -------
    d_pred: np.array
        Predicted data for each ensemble member.
    """
    # Initiate an array of predicted results.
    d_pred = np.zeros([x.shape[0], m_ensemble.shape[1]])
    for j in range(m_ensemble.shape[1]):
        # Calling the forward model for each member of the ensemble
        d_pred[:, j] = exponential(m_ensemble[:, j], x)
    return d_pred

# seed for the reproductibility
seed = 0
rng = np.random.default_rng(seed=seed)

a = 10.0
b = -0.0020
# timesteps
x = np.arange(500)
obs = exponential((a, b), x) + rng.normal(0.0, 1.0, 500)
# Initiate an ensemble of (a, b) parameters
n_ensemble = 100  # size of the ensemble
# Uniform law for the parameter a ensemble
ma = rng.uniform(low=-10.0, high=50.0, size=n_ensemble)
# ma = rng.normal(loc=20.0, scale=20.0, size=n_ensemble)
# Uniform law for the parameter b ensemble
mb = rng.uniform(low=-0.001, high=0.01, size=n_ensemble)
# mb = rng.normal(loc=-.002, scale=.005, size=n_ensemble)
# Prior ensemble
m_ensemble = np.stack((ma, mb), axis=0)

# Observation error covariance matrix
cov_obs = covmats.CovViaDiagonal(np.ones(obs.size) * 2.0)
# cov_obs = np.diag([0.5] * obs.shape[0])

# cov_obs = np.ones(obs.size)

# Bounds on parameters (size m * 2)
m_bounds = np.array([[0.0, 50.0], [-1.0, 1.0]])

# Number of assimilations
n_assimilations = 4

# Use a geometric suite (see procedure un evensen 2018) to compte alphas.
# Also explained in Torrado 2021 (see her PhD manuscript.)
cov_obs_inflation_geo = 1.2
cov_obs_inflation_factors: list[float] = [1.1]
for l in range(1, n_assimilations):
    cov_obs_inflation_factors.append(
        cov_obs_inflation_factors[l - 1] / cov_obs_inflation_geo
    )
scaling_factor: float = np.sum(1 / np.array(cov_obs_inflation_factors))
cov_obs_inflation_factors = [
    alpha * scaling_factor for alpha in cov_obs_inflation_factors
]

np.testing.assert_almost_equal(sum(1.0 / np.array(cov_obs_inflation_factors)), 1.0)

# This is just for the test
cov_mm_inflation_factor = 1.2

solver = ESMDA(
    obs,
    m_ensemble,
    cov_obs,
    forward_model,
    forward_model_args=(x,),
    forward_model_kwargs={},
    n_assimilations=n_assimilations,
    cov_obs_inflation_factors=cov_obs_inflation_factors,
    cov_mm_inflation_factor=cov_mm_inflation_factor,
    m_bounds=m_bounds,
    save_ensembles_history=True,
    inversion_type=ESMDAInversionType.EXACT_CHOLESKY,
    seed=seed,
    truncation=0.99,
    logger=logging.getLogger("ESMDA"),
)
# Call the ES-MDA solver
solver.solve()

# Assert that the parameters are found with a 5% accuracy.
assert np.isclose(np.average(solver.m_prior, axis=1), np.array([a, b]), rtol=5e-2).all()

# Get the approximated parameters
a_approx, b_approx = np.average(solver.m_prior, axis=1)

# Get the uncertainty on the parameters
a_std, b_std = np.sqrt(np.diagonal(solver.cov_mm))


solver.logger.info(f"a = {a_approx:.5f} +/- {a_std:.4E}")
solver.logger.info(f"b = {b_approx:.5f} +/- {b_std: 4E}")

INFO:ESMDA:Assimilation # 1
INFO:ESMDA:Assimilation # 2
INFO:ESMDA:Assimilation # 3
INFO:ESMDA:Assimilation # 4
INFO:ESMDA:Forecast for the final ensemble
INFO:ESMDA:a = 10.03181 +/- 1.3964E-01
INFO:ESMDA:b = -0.00202 +/-  7.653077E-05

• Same example but with ESMDA_RS

# This is to defined

# A priori estimated parameters covariance
std_m_prior = np.array([30.0, 0.01])

# Note: no inflation factor provided nor number of assimilations

# This is just for the test
cov_mm_inflation_factor: float = 0.9

solver = ESMDA_RS(
    obs,
    m_ensemble,
    cov_obs,
    forward_model,
    forward_model_args=(x,),
    forward_model_kwargs={},
    cov_mm_inflation_factor=cov_mm_inflation_factor,
    m_bounds=m_bounds,
    C_MD_localization=FixedLocalization(np.ones((m_ensemble.shape[0], obs.size))),
    C_DD_localization=FixedLocalization(np.ones((obs.size, obs.size))),
    save_ensembles_history=True,
    std_m_prior=std_m_prior,
    random_state=123,
    batch_size=1,
    is_parallel_analyse_step=True,
    logger=logging.getLogger("ESMDA-RS"),
)
# Call the ES-MDA-RS solver
solver.solve()

# Assert that the parameters are found with a 5% accuracy.
assert np.isclose(np.average(solver.m_prior, axis=1), np.array([a, b]), rtol=1e-1).all()

# Get the approximated parameters
a_approx, b_approx = np.average(solver.m_prior, axis=1)

# Get the uncertainty on the parameters
a_std, b_std = np.sqrt(np.diagonal(solver.cov_mm))


solver.logger.info(f"a = {a_approx:.5f} +/- {a_std:.4E}")
solver.logger.info(f"b = {b_approx:.5f} +/- {b_std: 4E}")

INFO:ESMDA-RS:Assimilation # 1
INFO:ESMDA-RS:- Inflation factor = 4132.540
INFO:ESMDA-RS:Assimilation # 2
INFO:ESMDA-RS:- Inflation factor = 71.034
INFO:ESMDA-RS:Assimilation # 3
INFO:ESMDA-RS:- Inflation factor = 4.377
INFO:ESMDA-RS:Assimilation # 4
INFO:ESMDA-RS:- Inflation factor = 1.321
INFO:ESMDA-RS:Forecast for the final ensemble
INFO:ESMDA-RS:a = 9.85206 +/- 1.2995E-01
INFO:ESMDA-RS:b = -0.00192 +/-  6.263513E-05

# Note: no inflation factor provided nor number of assimilations
# This is just for the test
cov_mm_inflation_factor: float = 0.9

solver = ESMDA_DMC(
    obs,
    m_ensemble,
    cov_obs,
    forward_model,
    forward_model_args=(x,),
    forward_model_kwargs={},
    cov_mm_inflation_factor=cov_mm_inflation_factor,
    m_bounds=m_bounds,
    C_MD_localization=FixedLocalization(np.ones((m_ensemble.shape[0], obs.size))),
    C_DD_localization=FixedLocalization(np.ones((obs.size, obs.size))),
    save_ensembles_history=True,
    random_state=123,
    batch_size=1,
    is_parallel_analyse_step=True,
    logger=logging.getLogger("ESMDA-DMC"),
)
# Call the ES-MDA-RS solver
solver.solve()

# Assert that the parameters are found with a 5% accuracy.
assert np.isclose(np.average(solver.m_prior, axis=1), np.array([a, b]), rtol=5e-2).all()

# Get the approximated parameters
a_approx, b_approx = np.average(solver.m_prior, axis=1)

# Get the uncertainty on the parameters
a_std, b_std = np.sqrt(np.diagonal(solver.cov_mm))


solver.logger.info(f"a = {a_approx:.5f} +/- {a_std:.4E}")
solver.logger.info(f"b = {b_approx:.5f} +/- {b_std: 4E}")

INFO:ESMDA-DMC:Assimilation # 1
INFO:ESMDA-DMC:- Inflation factor = 33060.322
INFO:ESMDA-DMC:Assimilation # 2
INFO:ESMDA-DMC:- Inflation factor = 737.164
INFO:ESMDA-DMC:Assimilation # 3
INFO:ESMDA-DMC:- Inflation factor = 125.379
INFO:ESMDA-DMC:Assimilation # 4
INFO:ESMDA-DMC:- Inflation factor = 11.016
INFO:ESMDA-DMC:Assimilation # 5
INFO:ESMDA-DMC:- Inflation factor = 1.153
INFO:ESMDA-DMC:Assimilation # 6
INFO:ESMDA-DMC:- Inflation factor = 30.593
INFO:ESMDA-DMC:Forecast for the final ensemble
INFO:ESMDA-DMC:a = 9.85543 +/- 1.3871E-01
INFO:ESMDA-DMC:b = -0.00193 +/-  6.096230E-05