Estimators (data assimilation)

Summary

File: src/estimators/__init__.py

Abstract base class (ABC) that wraps a Model and an optional Bias into a full DA loop. Provides the shared forecast_step() that advances both the model and the bias in time.

Key attributes: est_phi, est_alpha, est_bias, inflation_factor, num_DA_blind, num_SE_only

Key concrete methods: forecast_step(), _init_bias(), _MA() (applies measurement operator M)

Subclasses must implement: - analysis_step(d, Cdd) — Bayesian update given observation vector d and noise covariance Cdd

Estimator  (ABC)
├── EnsembleEstimator            → Stochastic filters
│   ├── EnKF
│   ├── EnSRKF
│   └── rBA_EnKF
└── DeterministicEstimator       → Deterministic filters
    └── KalmanFilter

The relationship between the three main classes — attributes, not subclasses:

Estimator instance
  .model  →  Model instance   (used for the forecast)
  .bias   →  Bias instance    (optional; None if bias-unaware)
             .forecaster  →  Model instance   (usually an ESN_model)

See Stochastic filters and Deterministic filters for the concrete implementations.

The real-time assimilation cycle: forecast, analysis, repeat.

---

romda.estimators.Estimator(**kwargs)

Bases: ABC

Abstract base class shared by all state/parameter estimators.

Every concrete estimator (EnKF, EnSRKF, rBA_EnKF, KalmanFilter, ...) owns a Model instance for the forecast step and, optionally, a Bias instance for bias-aware assimilation. The observation operator maps state to observation space, \(\mathbf{y} = \mathbf{M}\boldsymbol{\psi}\) if \(\mathbf{M}\) is a matrix, or \(\mathbf{y} = \mathbf{M}(\boldsymbol{\psi})\) if \(\mathbf{M}\) is callable; if not supplied explicitly it is read from model.M.

Subclasses must implement analysis_step(d, Cdd, **kwargs): the Bayesian update given observation \(\mathbf{d}\) and observation-noise covariance \(\mathbf{C}_{dd}\). Implementations build the forecast state internally, run the filter update, validate the result, and update the model history in-place.

Notation

Symbol	Meaning	Shape
\(N\)	augmented state dimension (\(N_\phi{+}N_\alpha\), or \(+N_q\) once observables are appended)
\(N_\phi\)	model state dimension
\(N_\alpha\)	number of estimated parameters
\(N_q\)	number of observables
\(m\)	ensemble size
\(\boldsymbol{\psi}\)	state vector / ensemble	\((N,)\) or \((N, m)\)
\(\mathbf{M}\)	measurement operator (matrix or callable)	\((N_q, N)\)
\(\mathbf{d}\)	observation vector	\((N_q,)\)
\(\mathbf{C}_{dd}\)	observation-noise covariance	\((N_q, N_q)\)
\(\mathbf{C}_{\psi\psi}\)	forecast (prior) covariance	\((N, N)\)
\(\mathbf{K}\)	Kalman gain	\((N, N_q)\)

Attributes:

Name	Type	Description
est_phi	bool	Whether to estimate the model state (default True).
est_alpha	list of str	Names of model parameters to estimate (default []).
est_bias	bool	Whether to estimate an observation bias (default False).
start_param	int	Analysis step at which parameter estimation starts; parameters are frozen in earlier analyses (0 = active from the first analysis).
start_bias	int	Analysis step at which a bias-aware filter starts its bias-aware update; a plain EnKF update is applied in earlier analyses (0 = active from the first analysis).
inflation_factor	float	Covariance/ensemble inflation factor (default 1.0).
inflation_factor_rejection	float	Inflation applied after a rejected analysis (default 1.002).
results_folder	str or None	Optional path for saving results.

References

Kalman (1960). A new approach to linear filtering and prediction problems. J. Basic Eng., 82(1), 35-45.

Evensen (2009). Data Assimilation: The Ensemble Kalman Filter. Springer.

Nóvoa, Racca & Magri (2023). Inferring unknown unknowns: Regularized bias-aware ensemble Kalman filter. Comput. Methods Appl. Mech. Eng., 418, 116502.

Source code in src/estimators/base.py

def __init__(self, **kwargs) -> None:
    # Record what was consumed so subclasses can exclude these keys from
    # whatever they forward the remaining kwargs to (e.g. the model ctor).
    self._consumed_kwargs: set = set()
    keys = list(kwargs.keys())
    for attr in keys:
         # Derived read-only properties (e.g. EnsembleEstimator.m, which reads
         # the model's history) must not be assigned; they are consumed
         # downstream, e.g. by Model.init_ensemble.
         class_attr = getattr(self.__class__, attr, None)
         if isinstance(class_attr, property) and class_attr.fset is None:
             continue
         # hasattr, not __class__.__dict__: the config fields are declared on
         # Estimator, so a leaf class (EnSRKF, EnKF, ...) has none of them in its
         # own __dict__ and every config kwarg was silently dropped.
         if hasattr(self.__class__, attr):
             setattr(self, attr, kwargs.pop(attr))
             self._consumed_kwargs.add(attr)

model property

The Model instance used for the forecast step.

bias property writable

Bias instance, or None if not configured.

current_state property

Current state (delegates to model).

current_time property

Current time (delegates to model).

current_bias_estimate property

Current bias estimate, or None if no bias model.

Nphi property

Size of the model state vector.

Nq property

Number of observable dimensions.

is_bias_aware property

True when the estimator carries an explicit bias correction.

assimilated_data property writable

Namedtuple with fields data and times of all assimilated obs.

update_history(psi, t=None, b=None, modify_saved_states=False, reset=False)

Update model (and bias) history.

Source code in src/estimators/base.py

def update_history(
    self,
    psi: np.ndarray,
    t=None,
    b=None,
    modify_saved_states: bool = False,
    reset: bool = False,
) -> None:
    """Update model (and bias) history."""
    self.model.update_history(psi, t, reset=reset, modify_saved_states=modify_saved_states)
    # b is None when the caller only advances the model state; there is no
    # new bias sample to record, so leave the bias history untouched.
    if self.bias is not None and b is not None:
        self.bias.update_history(b, t, reset=reset, modify_saved_states=modify_saved_states)

forecast_step(t_end=None, reset=False, close=False, output_forecast=False, **kwargs)

Advance model (and bias, if present) in time.

Parameters:

Name	Type	Description	Default
t_end	float	Target time; Nt is derived from it if provided.	None
reset	bool	Whether to reset model history on update.	False
close	bool	Close the integrator after stepping.	False
output_forecast	bool	If True, return the raw forecast array.	False

Returns:

Type	Description
ndarray or None

Source code in src/estimators/base.py

def forecast_step(
    self,
    t_end=None,
    reset: bool = False,
    close: bool = False,
    output_forecast: bool = False,
    **kwargs,
) -> np.ndarray | None:
    """Advance model (and bias, if present) in time.

    Parameters
    ----------
    t_end : float, optional
        Target time; Nt is derived from it if provided.
    reset : bool
        Whether to reset model history on update.
    close : bool
        Close the integrator after stepping.
    output_forecast : bool
        If True, return the raw forecast array.

    Returns
    -------
    ndarray or None
    """
    pm = self.model

    if t_end is not None:
        t_end = round(t_end, pm.precision_t)
        # round, not truncate: the division is one short whenever it is inexact
        # in binary (1.18 / 0.01 = 117.999...), same fix as legacy ensemble.py
        Nt = int(np.round((t_end - pm.current_time) / pm.dt))
    else:
        Nt = kwargs.get("Nt", -1)

    # print(f"\nForecasting from t={pm.current_time:.3f} to t={t_end:.3f} (Nt={Nt})...")

    if Nt == 0:
        return  # Already at requested time; nothing to advance.
    assert Nt > 0, "Must specify positive Nt or t_end for forecast_step."

    psi, t = pm.time_integrate(Nt=Nt, averaged=kwargs.get("averaged", False))

    if t_end is not None:
        assert abs(t[-1] - t_end) < pm.dt, (
            f"Final time {t[-1]} does not match requested t_end {t_end}."
        )

    try:
        pm.update_history(psi, t, reset=reset)
    except ValueError as e:
        print("Solver didn't return a homogeneous psi. Check initial conditions.")
        raise e

    # Advance bias model if present
    if self.bias is not None:
        pb = self.bias
        b, t_b = pb.time_integrate(Nt=Nt)
        pb.update_history(b, t_b, reset=reset)
        if pm.current_time != pb.current_time:
            raise AssertionError(
                f"Time mismatch: model {pm.current_time} vs bias {pb.current_time}"
            )

    if close:
        pm.close()

    if output_forecast:
        return psi

analysis_step(d, Cdd, return_analysis=False) abstractmethod

Perform the Bayesian analysis step given observation d.

Implementations are responsible for building the augmented forecast state, running the filter update, validating the result, and updating the model history in-place. Optionally return the analysed state when return_analysis=True.

Source code in src/estimators/base.py

@abstractmethod
def analysis_step(
    self,
    d: np.ndarray,
    Cdd: np.ndarray,
    return_analysis: bool = False,
) -> np.ndarray | None:
    """Perform the Bayesian analysis step given observation *d*.

    Implementations are responsible for building the augmented forecast
    state, running the filter update, validating the result, and updating
    the model history in-place.  Optionally return the analysed state when
    ``return_analysis=True``.
    """