# --- System setup ---
import os
os.environ["JAX_ENABLE_X64"] = "True"

# --- Core scientific stack ---
import numpy as np           # Standard CPU-based numerical computing
import scipy                 # Stats & scientific routines (e.g. chi-square, mahalanobis)

# --- Optional GPU-accelerated backends ---
# The estimators below take the array module as a parameter (`xnp`), so they run
# on jax.numpy as well. Neither JAX nor TensorFlow is required for this notebook,
# and both are skipped silently when they are not installed.
try:
    import jax.numpy as jnp
except ImportError:
    jnp = None

# --- Visualization ---
import matplotlib.pyplot as plt

print("numpy", np.__version__, "| jax.numpy available:", jnp is not None)

numpy 2.1.3 | jax.numpy available: True

import numpy as np
from typing import Optional, Union

ArrayLike = Union[np.ndarray, "jnp.ndarray"]


class OnlineCovarianceMeanEstimator:
    """
    Incremental (online) estimator for mean and covariance matrix,
    including an online estimate of the inverse covariance matrix.

    This class allows updating the distribution mean, covariance,
    and inverse covariance incrementally as new samples arrive,
    without recomputing from scratch.

    Attributes:
        xnp: Backend numerical library (`numpy` or `jax.numpy`).
        dim: Dimensionality of the input data.
        λ: Forgetting factor in (0,1]; values < 1 give exponentially
           less weight to older samples (useful for non-stationary data).
        W_n: Running weight (scalar).
        W_n2: Running sum of squared weights (scalar).
        x̅_n: Current running mean vector of shape (dim, 1).
        M̅_n: Accumulated scatter matrix of shape (dim, dim).
        M̅_n_inv: Running estimate of the inverse scatter matrix of shape (dim, dim),
                 or None while the scatter is not yet invertible.

    Notes:
        - If only the inverse covariance is required, references to `M̅_n` can be removed
          to save memory and computation.
        - The unbiased covariance estimate uses both `W_n` and `W_n2`.
        - A zero scatter matrix has no inverse. The estimator therefore starts with
          `M̅_n_inv = None` and initialises it by a direct solve as soon as the scatter
          is positive definite and reasonably conditioned; from then on it is updated
          with the Sherman-Morrison formula. Initialising the inverse to a scaled
          identity while the scatter is still zero would mean that the two matrices
          describe different models: their product is not the identity, and the
          reported covariance would silently omit the implied ridge.
    """

    def __init__(self, xnp, dim: int, λ: float,
                 track_inverse: bool = True, cond_limit: float = 1e10):
        """
        Initialize the online estimator.

        Args:
            xnp: Numerical backend (`numpy` or `jax.numpy`).
            dim: Dimensionality of the data.
            λ: Forgetting factor (close to 1.0 for long memory,
               smaller values adapt faster to changes).
            track_inverse: Maintain the inverse scatter matrix as well.
            cond_limit: Condition number above which the scatter is not yet
               considered safe to invert.
        """
        if not 0.0 < λ <= 1.0:
            raise ValueError("λ must lie in (0, 1]")
        self.xnp = xnp
        self.dim = dim
        self.λ = λ
        self.track_inverse = track_inverse
        self.cond_limit = cond_limit

        self.W_n: float = 0.0
        self.W_n2: float = 0.0
        self.x̅_n: ArrayLike = self.xnp.zeros((dim, 1), dtype="float64")
        self.M̅_n: ArrayLike = self.xnp.zeros((dim, dim), dtype="float64")
        self.M̅_n_inv: Optional[ArrayLike] = None

    # -- normalization factors -------------------------------------------------

    def _unbiased_denominator(self) -> float:
        """Reliability-weight correction W_n - W_n^{(2)}/W_n."""
        if self.W_n <= 0.0:
            raise ValueError("no observations yet")
        denominator = self.W_n - self.W_n2 / self.W_n
        if denominator <= 0.0:
            raise ValueError("unbiased covariance needs more than one effective observation")
        return denominator

    def get_cov(self) -> ArrayLike:
        """
        Return the unbiased covariance estimate.

        Returns:
            Covariance matrix of shape (dim, dim).
        """
        return self.M̅_n / self._unbiased_denominator()

    def get_cov_biased(self) -> ArrayLike:
        """
        Return the biased (population-normalized) covariance estimate.

        Returns:
            Covariance matrix of shape (dim, dim).
        """
        if self.W_n <= 0.0:
            raise ValueError("no observations yet")
        return self.M̅_n / self.W_n

    def get_cov_inv(self) -> ArrayLike:
        """
        Return the inverse of the unbiased covariance estimate.

        Note that this is the inverse of an unbiased estimate, which is not itself
        an unbiased estimate of the population precision.

        Returns:
            Inverse covariance matrix of shape (dim, dim).
        """
        return self._require_inverse() * self._unbiased_denominator()

    def get_cov_inv_biased(self) -> ArrayLike:
        """
        Return the inverse of the biased covariance estimate.

        Returns:
            Inverse covariance matrix of shape (dim, dim).
        """
        if self.W_n <= 0.0:
            raise ValueError("no observations yet")
        return self._require_inverse() * self.W_n

    def get_mean(self) -> ArrayLike:
        """
        Return the current running mean estimate.

        Returns:
            Mean vector of shape (dim,).
        """
        return self.x̅_n.flatten()

    # -- inverse bookkeeping ---------------------------------------------------

    def _require_inverse(self) -> ArrayLike:
        if self.M̅_n_inv is None:
            raise ValueError(
                "the inverse is not available yet: the scatter matrix is still "
                "singular or badly conditioned (it has rank at most min(dim, n-1))")
        return self.M̅_n_inv

    def _try_start_inverse(self) -> None:
        """Initialize M̅_n_inv by a direct solve once that is numerically sensible."""
        M = np.asarray(self.M̅_n, dtype="float64")
        M = 0.5 * (M + M.T)
        try:
            np.linalg.cholesky(M)                      # fails fast while rank-deficient
        except np.linalg.LinAlgError:
            return
        if np.linalg.cond(M) > self.cond_limit:
            return
        self.M̅_n_inv = self.xnp.linalg.solve(self.M̅_n, self.xnp.eye(self.dim))

    def refactor_inverse(self) -> None:
        """Recompute the inverse from the current scatter, discarding accumulated drift."""
        self.M̅_n_inv = None
        self._try_start_inverse()

    # -- the update ------------------------------------------------------------

    def update(self, x_n: ArrayLike) -> None:
        """
        Update mean, covariance, and inverse covariance with a new sample.

        Args:
            x_n: New data vector of shape (dim,) or (dim, 1).
        """
        x_n = x_n.reshape(self.dim, 1)

        # Update weights. The second line accumulates squared weights, so the
        # decay factor is squared and the previous squared sum is propagated.
        self.W_n = self.λ * self.W_n + 1.0
        self.W_n2 = self.λ**2 * self.W_n2 + 1.0

        # Update mean
        Δ_n = x_n - self.x̅_n
        self.x̅_n = self.x̅_n + Δ_n / self.W_n

        # Update scatter matrix. Δ_n uses the OLD mean, Δ_n1 the NEW one.
        Δ_n1 = x_n - self.x̅_n
        self.M̅_n = self.λ * self.M̅_n + self.xnp.dot(Δ_n, Δ_n1.T)

        if not self.track_inverse:
            return

        if self.M̅_n_inv is None:
            # Nothing to update yet; try to start from the current scatter.
            self._try_start_inverse()
            return

        # Sherman-Morrison. The two matrix-vector products are formed first, so
        # that the only matrix-matrix operation is the rank-one outer product.
        P = self.M̅_n_inv
        P_u = P @ Δ_n                     # (dim, 1)
        v_P = Δ_n1.T @ P                  # (1, dim)
        den = self.λ + (v_P @ Δ_n).item()
        self.M̅_n_inv = (P - (P_u @ v_P) / den) / self.λ

# --- Create synthetic covariance matrix and mean vector ---

rng = np.random.default_rng(20250927)   # recorded seed: the original run had none

matrixSize = 200  # dimensionality of the random dataset

# Random base matrix with values centered around 0
A = rng.random((matrixSize, matrixSize)) - 0.5

# Construct covariance as A @ A.T
# (guaranteed symmetric positive semidefinite, but may be ill-conditioned)
cov = 5 * (A @ A.T)

# Add a small multiple of the identity to improve conditioning:
# - ensures positive definiteness
# - ridge term is scaled relative to trace(cov) / matrixSize
#   (so it's adaptive to the scale of the covariance)
cov += 1e-2 * np.trace(cov) / matrixSize * np.eye(matrixSize)

# Random mean vector
mean = rng.random(matrixSize)

import time

# Initialize OnlineCovarianceMeanEstimator
# λ = 1 → no forgetting (full history is considered).
# For streaming data with drift, set 0 < λ < 1 to introduce exponential forgetting.
ocme = OnlineCovarianceMeanEstimator(xnp=np, dim=mean.shape[0], λ=1)

# Generate synthetic dataset: 5000 samples from a Gaussian distribution
X = rng.multivariate_normal(mean, cov, size=5000)

# --- Online update loop ---
# Process each data point sequentially to update the running mean and covariance.
start = time.perf_counter()
for x in X:
    ocme.update(x)
elapsed = time.perf_counter() - start

print(f"Time needed: {elapsed:.2f} seconds")
print(f"Inverse scatter available: {ocme.M̅_n_inv is not None}")

Time needed: 1.74 seconds
Inverse scatter available: True

# --- Consistency check between online and batch estimators ---

# Mean absolute error between online covariance estimate and batch covariance
error_sigma = np.abs(ocme.get_cov() - np.cov(X.T)).mean()

# Mean absolute error between online inverse covariance estimate
# and the inverse of the batch covariance
error_sigma_inv = np.abs(ocme.get_cov_inv() - np.linalg.inv(np.cov(X.T))).mean()

# Mean absolute error between online mean estimate and batch mean
error_mu = np.abs(ocme.get_mean() - X.mean(axis=0)).mean()

# Print error values for inspection
print("error_mu", error_mu)
print("error_sigma", error_sigma)
print("error_sigma_inv", error_sigma_inv)

# Assert that all errors are within a small numerical tolerance
# (note: inverse covariance usually has higher sensitivity, so tolerance is relaxed)
assert error_mu < 1e-9
assert error_sigma < 1e-9
assert error_sigma_inv < 1e-5

error_mu 1.0507653774860515e-15
error_sigma 8.279804747097878e-15
error_sigma_inv 2.1016007240598895e-14

# --- The maintained inverse must invert the covariance that is reported ---
small = np.random.default_rng(20260426).normal(size=(40, 4)) * .01

est = OnlineCovarianceMeanEstimator(xnp=np, dim=4, λ=1.0)
for i, x in enumerate(small):
    est.update(x)
    if i < 4:
        # rank(M̅_n) <= min(d, n-1): no inverse can exist yet
        assert est.M̅_n_inv is None, f"inverse claimed after {i+1} observations"

reference_cov = np.cov(small.T)
residual = np.abs(est.get_cov() @ est.get_cov_inv() - np.eye(4)).max()

print("max |Σ̂ - np.cov|          :", np.abs(est.get_cov() - reference_cov).max())
print("max |Σ̂ · Σ̂⁻¹ - I|         :", residual)
assert np.abs(est.get_cov() - reference_cov).max() < 1e-12
assert residual < 1e-9

max |Σ̂ - np.cov|          : 2.710505431213761e-20
max |Σ̂ · Σ̂⁻¹ - I|         : 3.1289817268479293e-16

from typing import Any, Optional


class BatchCovarianceMeanEstimator:
    """
    Batch-incremental estimator of the mean, covariance matrix, and inverse covariance matrix.

    This class updates running estimates when a new *batch* of observations arrives,
    using a forgetting factor λ ∈ (0, 1] to adapt to non-stationary data. The inverse
    covariance is updated via the matrix inversion lemma to avoid full matrix inversion.

    Notes:
        - If you only need the inverse covariance, you can remove all references to `M̅_n`
          (the covariance accumulator) to save memory/compute.
        - Shapes follow the convention:
              X_n: (μ, dim) or (dim,)      # batch of μ row-vectors or a single vector
              x̅_n: (dim, 1)
              M̅_n, M̅_n_inv: (dim, dim)
        - As in the fully online estimator, `M̅_n_inv` starts as None and is initialised
          by a direct solve once the scatter is invertible. The forgetting factor applies
          once per *batch*: a batch of size μ with λ has limiting memory μ(1+λ)/(1-λ),
          so the same λ means different things for different batch sizes.
    """

    def __init__(self, xnp: Any, dim: int, λ: float,
                 track_inverse: bool = True, cond_limit: float = 1e10) -> None:
        """
        Initialize the batch estimator.

        Args:
            xnp: Numerical backend (e.g., `numpy` or `jax.numpy` module).
            dim: Data dimensionality.
            λ:   Forgetting factor in (0, 1]; values < 1 downweight older batches.
            track_inverse: Maintain the inverse scatter matrix as well.
            cond_limit: Condition number above which the scatter is not yet
                considered safe to invert.

        Returns:
            None
        """
        if not 0.0 < λ <= 1.0:
            raise ValueError("λ must lie in (0, 1]")
        self.xnp: Any = xnp
        self.dim: int = dim
        self.λ: float = λ
        self.track_inverse: bool = track_inverse
        self.cond_limit: float = cond_limit

        # Weighted counts for unbiased covariance correction
        self.W_n: float = 0.0
        self.W_n2: float = 0.0

        # Running statistics
        self.x̅_n = self.xnp.zeros((dim, 1), dtype="float64")   # mean
        self.M̅_n = self.xnp.zeros((dim, dim), dtype="float64") # scatter accumulator
        self.M̅_n_inv: Optional[Any] = None                     # inverse scatter, once available

    def _unbiased_denominator(self) -> float:
        """Reliability-weight correction W_n - W_n^{(2)}/W_n."""
        if self.W_n <= 0.0:
            raise ValueError("no observations yet")
        denominator = self.W_n - self.W_n2 / self.W_n
        if denominator <= 0.0:
            raise ValueError("unbiased covariance needs more than one effective observation")
        return denominator

    def get_cov(self):
        """
        Return the *unbiased* covariance estimate.

        Returns:
            (dim, dim) covariance matrix.
        """
        return self.M̅_n / self._unbiased_denominator()

    def get_cov_biased(self):
        """
        Return the *biased* covariance estimate (dividing by total weight).

        Returns:
            (dim, dim) covariance matrix.
        """
        if self.W_n <= 0.0:
            raise ValueError("no observations yet")
        return self.M̅_n / self.W_n

    def get_cov_inv(self):
        """
        Return the inverse of the *unbiased* covariance estimate.

        Warning:
            Multiplying `M̅_n_inv` by the unbiased weight factor may be numerically
            sensitive when the condition number is large.

        Returns:
            (dim, dim) inverse covariance matrix.
        """
        return self._require_inverse() * self._unbiased_denominator()

    def get_cov_inv_biased(self):
        """
        Return the inverse of the *biased* covariance estimate.

        Returns:
            (dim, dim) inverse covariance matrix.
        """
        if self.W_n <= 0.0:
            raise ValueError("no observations yet")
        return self._require_inverse() * self.W_n

    def get_mean(self):
        """
        Return the current running mean as a flat vector.

        Returns:
            (dim,) mean vector.
        """
        return self.x̅_n.flatten()

    def _require_inverse(self):
        if self.M̅_n_inv is None:
            raise ValueError(
                "the inverse is not available yet: the scatter matrix is still "
                "singular or badly conditioned (it has rank at most min(dim, n-1))")
        return self.M̅_n_inv

    def _try_start_inverse(self) -> None:
        """Initialize M̅_n_inv by a direct solve once that is numerically sensible."""
        M = np.asarray(self.M̅_n, dtype="float64")
        M = 0.5 * (M + M.T)
        try:
            np.linalg.cholesky(M)
        except np.linalg.LinAlgError:
            return
        if np.linalg.cond(M) > self.cond_limit:
            return
        self.M̅_n_inv = self.xnp.linalg.solve(self.M̅_n, self.xnp.eye(self.dim))

    def refactor_inverse(self) -> None:
        """Recompute the inverse from the current scatter, discarding accumulated drift."""
        self.M̅_n_inv = None
        self._try_start_inverse()

    def update(self, X_n) -> None:
        """
        Update mean, covariance accumulator, and inverse covariance with a new batch.

        Args:
            X_n: Batch of observations of shape (μ, dim); a single vector of shape (dim,)
                 is also accepted.

        Returns:
            None
        """
        # Ensure 2D: (μ, dim)
        if X_n.ndim == 1:
            X_n = X_n.reshape(1, self.dim)
        μ: int = X_n.shape[0]

        # Update weighted counts
        self.W_n = self.λ * self.W_n + μ
        self.W_n2 = self.λ**2 * self.W_n2 + μ

        # Mean update. Every Δ uses the mean from BEFORE the batch.
        Δ_n = X_n - self.x̅_n.T                                             # (μ, dim)
        self.x̅_n = self.x̅_n + Δ_n.sum(axis=0, keepdims=True).T / self.W_n  # (dim, 1)

        # Scatter accumulator update, with residuals about the NEW mean
        XX_n = X_n - self.x̅_n.T                                            # (μ, dim)
        self.M̅_n = self.λ * self.M̅_n + Δ_n.T @ XX_n                        # (dim, dim)

        if not self.track_inverse:
            return

        if self.M̅_n_inv is None:
            self._try_start_inverse()
            return

        # Woodbury update. The inner system is μ x μ and is solved rather than
        # inverted; only that smaller system depends on the batch size.
        P = self.M̅_n_inv
        P_D = P @ Δ_n.T                                   # (dim, μ)
        X_P = XX_n @ P                                    # (μ, dim)
        inner = self.λ * self.xnp.eye(μ) + XX_n @ P_D     # (μ, μ)
        self.M̅_n_inv = (P - P_D @ self.xnp.linalg.solve(inner, X_P)) / self.λ

μ_min_max = 10, 20 # minimum and maximum batch size

# Use the same cov and mean matrix as above, to allow some comparison
if cov is None or mean is None:
  # Create some sample data
  matrixSize = 100
  A = np.random.rand(matrixSize, matrixSize) - 0.5
  cov = 5 * (A @ A.T)                # SPD but can be ill-conditioned
  cov += 1e-2 * np.trace(cov)/matrixSize * np.eye(matrixSize)  # scale-aware ridge
  mean = np.random.rand(matrixSize)
else:
  print(f"Using the already defined covariance matrix with shape {cov.shape} and the already defined mean vector.")

Using the already defined covariance matrix with shape (200, 200) and the already defined mean vector.

# Change 0 < λ <= 1 to values smaller than 1 to introduce "forgetting"
#   λ = 1.0   → no forgetting (stationary data)
#   λ < 1.0   → older data is exponentially down-weighted (good for drift adaptation)
ocme = BatchCovarianceMeanEstimator(xnp=np, dim=matrixSize, λ=1.0)

# Store all generated samples for later comparison with offline estimates
all_X = []

start = time.perf_counter()
for i in range(5000):
    # Draw a random batch size between μ_min_max[0] and μ_min_max[1]
    batch_size = int(rng.integers(*μ_min_max))

    # Generate batch of samples from the same Gaussian distribution
    X = rng.multivariate_normal(mean, cov, size=batch_size)

    # Update online estimates with the new batch
    ocme.update(X)

    # Keep a copy for offline validation
    all_X.append(X)

print(f"Time needed: {time.perf_counter() - start:.2f} seconds")

# Stack all batches into a single array for offline covariance/mean comparison
all_X = np.vstack(all_X)

Time needed: 87.61 seconds

# --- Validate online batch estimator against offline ground truth ---

# Mean absolute error of the covariance estimate vs. offline covariance
error_sigma = np.abs(ocme.get_cov() - np.cov(all_X.T)).mean()

# Mean absolute error of the inverse covariance vs. offline inverse covariance
error_sigma_inv = np.abs(ocme.get_cov_inv() - np.linalg.inv(np.cov(all_X.T))).mean()

# Mean absolute error of the mean estimate vs. offline mean
error_mu = np.abs(ocme.get_mean() - all_X.mean(axis=0)).mean()

# Print errors for inspection
print("error_mu", error_mu)
print("error_sigma", error_sigma)
print("error_sigma_inv", error_sigma_inv)

# --- Sanity checks ---
# Very strict tolerances: mean and covariance must be nearly identical to offline values,
# but the inverse covariance is allowed a slightly higher tolerance due to numerical instability
assert error_mu < 1e-9
assert error_sigma < 1e-9
assert error_sigma_inv < 1e-5

error_mu 2.576174950447152e-15
error_sigma 7.957882661483676e-15
error_sigma_inv 2.6448878451212686e-16

# Create some sample data
def experiment(
    mu: np.ndarray,
    cov: np.ndarray,
    n_mem: int,
    λ: float,
    size: int = 500,
    generator: np.random.Generator = None,
) -> np.ndarray:
    """
    Run an experiment to compare the OnlineCovarianceMeanEstimator against
    the "true" mean and covariance computed from the last `n_mem` samples.

    Args:
        mu: Mean vector of the multivariate Gaussian distribution
            (shape: (n_features,)).
        cov: Covariance matrix of the distribution (shape: (n_features, n_features)).
        n_mem: Effective memory length to compare against (number of most recent samples).
        λ: Forgetting factor for the online estimator (0 < λ <= 1).
           - λ = 1: No forgetting, uses all past samples.
           - λ < 1: Introduces forgetting, so older samples gradually lose influence.
        size: Total number of samples to generate (default: 500).
        generator: NumPy random generator, so that a whole run is reproducible.

    Returns:
        A NumPy array of shape (2, d), where:
          - Row 0: Online estimates (mean and covariance, flattened).
          - Row 1: Ground-truth estimates from the last `n_mem` samples.

        Order: [mean_x1, mean_x2, Σ_11, Σ_12, Σ_21, Σ_22]
    """
    generator = np.random.default_rng() if generator is None else generator

    # Generate synthetic dataset of `size` samples
    # Note: The online estimator will "forget" old samples when λ < 1.
    X = generator.multivariate_normal(mu, cov, size=size)

    # Initialize online estimator (dim=2 since we assume 2D data here).
    # This experiment only reads the mean and covariance, so the inverse is
    # not maintained -- that is pure bookkeeping cost here.
    ocme = OnlineCovarianceMeanEstimator(xnp=np, dim=2, λ=λ, track_inverse=False)

    # Update estimator with samples one by one
    for x in X:
        ocme.update(x)

    # Compute ground-truth mean and covariance from the last `n_mem` samples
    mu_est = X[-n_mem:].mean(axis=0)
    cov_est = np.cov(X[-n_mem:].T)

    # Concatenate results for easier comparison
    return np.array([
        np.concatenate([ocme.get_mean(), ocme.get_cov().flatten()]),
        np.concatenate([mu_est, cov_est.flatten()])
    ])


# Labels for result vector elements (for plotting or reporting)
elem_names = ["$x̅_1$", "$x̅_2$", "$Σ_{1,1}$", "$Σ_{1,2}$", "$Σ_{2,1}$", "$Σ_{2,2}$"]

# --- Define Gaussian distribution parameters ---
cov_2d = np.array([[3, -3], [-3, 3.5]])   # Covariance matrix (2D, positive-definite)
mean_2d = [1, 2]                          # Mean vector

# --- Online estimator settings ---
lam = 0.99

# Effective memory length derived from λ. The limiting value is (1+λ)/(1-λ);
# after a finite number of observations it is slightly smaller.
n_mem_limit = round((1 + lam) / (1 - lam))
size = 500             # Number of samples per experiment
n_experiments = 50000  # Number of repeated experiments
# ⚠ Reduce this to e.g. 1000 for quicker runs (less accurate histograms though)

w_raw = lam ** np.arange(size)
n_mem_finite = w_raw.sum() ** 2 / (w_raw ** 2).sum()

print(f"Limiting memory of the fully online estimator with λ={lam}: n_mem={n_mem_limit}")
print(f"Finite-sample memory after n={size} observations:           n_mem={n_mem_finite:.3f}")
print(f"(the sum of weights W_n converges to {1/(1-lam):.0f}, which is a different quantity)")


def run_experiments(n_mem: int, seed: int) -> np.ndarray:
    """Repeat `experiment` n_experiments times and stack the results."""
    generator = np.random.default_rng(seed)
    start = time.perf_counter()
    results = [experiment(mean_2d, cov_2d, n_mem, lam, size, generator)
               for _ in range(n_experiments)]
    print(f"Time needed: {time.perf_counter() - start:.2f} seconds")
    return np.stack(results)


all_experiments = run_experiments(n_mem_limit, seed=20260412)
all_experiments.shape

Limiting memory of the fully online estimator with λ=0.99: n_mem=199
Finite-sample memory after n=500 observations:           n_mem=196.402
(the sum of weights W_n converges to 100, which is a different quantity)
Time needed: 230.20 seconds

(50000, 2, 6)

# Covariance of the weighted means
w_raw = lam ** np.arange(size)
W = w_raw.sum()
W2 = (w_raw ** 2).sum()

# Normalize FOR A MEAN: divide by W_n, so that the weights sum to one.
w = w_raw / W
assert np.isclose(w.sum(), 1.0), "mean weights must sum to one"

expected_Sigma_mean = cov_2d * (w ** 2).sum()
actual_Sigma_mean = np.cov(all_experiments[:, 0, 0:2].T)

print("expected_Sigma_mean:\n", expected_Sigma_mean)
print("\nactual_Sigma_mean:\n", actual_Sigma_mean)
print(f"\n1 / sum(w_i^2) = {1/(w**2).sum():.3f} = n_mem after {size} observations")

# For comparison only: normalizing with the unbiased COVARIANCE denominator
w_wrong = w_raw / (W - W2 / W)
print(f"\nUsing the covariance denominator instead would predict "
      f"{(cov_2d * (w_wrong**2).sum())[0, 0]:.8f}")
print(f"the correct value is                                     "
      f"{expected_Sigma_mean[0, 0]:.8f}")
print(f"the simulation gives                                     "
      f"{actual_Sigma_mean[0, 0]:.8f}")

assert np.abs(expected_Sigma_mean - actual_Sigma_mean).mean() < 1e-3

expected_Sigma_mean:
 [[ 0.01527479 -0.01527479]
 [-0.01527479  0.01782059]]

actual_Sigma_mean:
 [[ 0.01520006 -0.01515556]
 [-0.01515556  0.01767541]]

1 / sum(w_i^2) = 196.402 = n_mem after 500 observations

Using the covariance denominator instead would predict 0.01543153
the correct value is                                     0.01527479
the simulation gives                                     0.01520006

# --- Compare online vs. offline estimators: parameter distributions ---
def plot_distributions(all_experiments, n_mem, expected_Sigma_mean):
    fig, axs = plt.subplots(2, 3, figsize=(15, 10), constrained_layout=True)

    # Loop over each estimated parameter (mean components + covariance elements)
    for idx in range(all_experiments.shape[-1]):
        ax_row = idx // 3
        ax_col = idx % 3
        ax = axs[ax_row, ax_col]

        # Extract offline ("real") vs. online estimates across experiments
        reals = all_experiments[:, 1, idx]    # offline/batch reference
        onlines = all_experiments[:, 0, idx]  # online estimates

        # Plot histograms (only label once to avoid duplicate legend entries)
        ax.hist(
            reals,
            label="offline estimator ($n_{mem}=" + str(round(n_mem)) + "$)" if idx < 1 else None,
            bins=50,
            alpha=0.5,
        )
        ax.hist(
            onlines,
            label="online estimator" if idx < 1 else None,
            bins=50,
            alpha=0.5,
        )

        # --- Annotate offline estimator statistics (mean & std) ---
        mu, sigma = np.mean(reals), np.std(reals)
        props = dict(boxstyle="round", facecolor="wheat", alpha=0.4)
        text_box_str = (
            "$\\hat{\\mu}_{x̅_{offline}}=" + str(round(mu, 2)) + "$\n"
            + "$\\hat{\\sigma}_{x̅_{offline}}=" + str(round(sigma, 2)) + "$"
        )
        ax.text(0.025, 0.95, text_box_str, transform=ax.transAxes, fontsize=14,
                verticalalignment="top", bbox=props)

        # --- Annotate online estimator statistics (mean & std) ---
        mu, sigma = np.mean(onlines), np.std(onlines)
        text_box_str = (
            "$\\hat{\\mu}_{x̅_{online}}=" + str(round(mu, 2)) + "$\n"
            + "$\\hat{\\sigma}_{x̅_{online}}=" + str(round(sigma, 2)) + "$"
        )
        ax.text(0.7, 0.95, text_box_str, transform=ax.transAxes, fontsize=14,
                verticalalignment="top", bbox=props)

        # --- Add theoretical expected std of the online mean (mean components only) ---
        # This is the predicted value sqrt(Σ_jj · Σ_i w_i²), not a second copy of
        # the empirical standard deviation.
        if idx < 2:
            std = np.sqrt(expected_Sigma_mean[idx, idx])
            text_box_str = "Expected:\n" + "$\\sigma_{x̅_{online}}=" + str(round(std, 2)) + "$"
            ax.text(0.7, 0.7, text_box_str, transform=ax.transAxes, fontsize=14,
                    verticalalignment="top", bbox=props)

        # Title for each subplot (parameter being estimated)
        ax.set_title(f"Distribution for Estimations of {elem_names[idx]}")
        ax.grid()

    # --- Shared legend across subplots ---
    fig.legend(loc="upper center", bbox_to_anchor=(0.5, -0.0),
               fancybox=True, shadow=True, ncol=2)

    # --- Global title ---
    fig.suptitle(
        "Distributions of the est. parameters $\\mathbf{\\overline{x}}$ & $\\mathbf{\\overline{\\Sigma}}$\n"
        "for a forgetting fully-online estimator vs. an offline estimator using the last "
        f"$n_={round(n_mem)}$ samples",
        fontsize=16,
    )
    return fig

# The sum of weights W_n converges to 1/(1-λ) = 100, so n_mem = 100 is the
# tempting guess. The online estimates come out visibly narrower than the
# offline ones, so 100 samples is not the right comparison.
experiments_100 = run_experiments(100, seed=20260413)
expected_100 = cov_2d * (w ** 2).sum()

fig = plot_distributions(experiments_100, 100, expected_100)
plt.show(fig)

Time needed: 224.52 seconds

# With n_mem = (1+λ)/(1-λ) = 199 the two distributions agree closely, for the
# means as well as for the covariance entries. The remaining small mismatch is
# consistent with the finite-sample memory being 196.4 rather than 199.
fig = plot_distributions(all_experiments, n_mem_limit, expected_Sigma_mean)
plt.show(fig)

print("empirical std of the online mean :", all_experiments[:, 0, 0:2].std(axis=0))
print("predicted std of the online mean :", np.sqrt(np.diag(expected_Sigma_mean)))
print("empirical std of the offline mean:", all_experiments[:, 1, 0:2].std(axis=0))

empirical std of the online mean : [0.1232873  0.13294756]
predicted std of the online mean : [0.12359123 0.13349379]
empirical std of the offline mean: [0.12279982 0.13262518]

Incremental / Online Estimation of the (Inverse) Covariance Matrix and Mean¶

Example using the Fully Online Mean-Covariance Estimator¶

A note on starting the inverse¶

Batch-incremental of (inverse) Covariance Matrix and Mean¶

Memory of the Fully Online Mean–Covariance Estimator for $\lambda < 1$¶

Variance of Weighted Means¶

Results when (wrongly) assuming that $n_{mem}=100$¶

Results when assuming that $n_{mem}=\frac{1+\lambda}{1-\lambda}$¶

What this notebook corrects¶