Stats Cov

cov

CLASS	DESCRIPTION
RunningCovariance	Numerically stable running covariance matrix computation.

Classes

RunningCovariance(state: dict | None = None)

Bases: Stats

Numerically stable running covariance matrix computation.

Computes mean, covariance, correlation, variance, and standard deviation incrementally using the Chan-Golub-LeVeque algorithm. Maintains full covariance matrix between all feature dimensions.

The covariance is computed using incremental updates:

\(\mu_{new} = \mu_{old} + \frac{\sum_i (x_i - \mu_{old})}{n_{total}}\), \(C_{new} = C_{old} + \delta^\top \delta'\)

where:

• \(\delta = x_i - \mu_{old}\) (deviation from old mean)
• \(\delta' = x_i - \mu_{new}\) (deviation from new mean)
• \(C\) is the second central moment matrix

The covariance matrix is then:

\(\Sigma = \frac{C}{n - 1} \quad \text{(unbiased)}\)

PARAMETER	DESCRIPTION
state	Previous state to restore from. TYPE: dict | None, optional, by default None DEFAULT: None

Properties

count : int Total number of samples processed.

data_shape : tuple | None Original shape of feature dimensions.

mean : torch.Tensor | None Current mean estimate (n_features,).

cmom2 : torch.Tensor | None Second central moment matrix (n_features, n_features).

weight_sum : float Cumulative sum of weights (for weighted samples).

Examples:

Basic usage:

>>> from spectre.stats import RunningCovariance
>>> import torch
>>>
>>> stat = RunningCovariance()
>>> for _ in range(10):
...     batch = torch.randn(100, 5)
...     stat.add(batch)
>>>
>>> mean = stat.mean()
>>> cov = stat.covariance()
>>> corr = stat.correlation()

Extract variance from covariance diagonal:

>>> variance = stat.variance()  # Diagonal of covariance matrix
>>> stdev = stat.stdev()

METHOD	DESCRIPTION
add	Add batch of samples to running covariance computation.
to	Move internal tensors to specified device.
mean	Get current mean estimate.
covariance	Get current covariance matrix estimate.
correlation	Get correlation matrix from covariance.
variance	Get variance (diagonal of covariance matrix).
stdev	Get standard deviation (square root of variance).
state_dict	Save state to dictionary.
load_state_dict	Load state from dictionary.

ATTRIBUTE	DESCRIPTION
count	Get total number of samples processed. TYPE: int
data_shape	Get original shape of feature dimensions. TYPE: tuple | None
weight_sum	Get total weight of samples processed. TYPE: float

Source code in spectre/stats/cov.py

def __init__(self, state: dict | None = None):
    if state is not None:
        return super().__init__(state=state)

    self._count = 0
    self._weight_sum = 0.0
    self._mean = None
    self._cmom2 = None
    self._data_shape = None

Attributes

count: int property writable

Get total number of samples processed.

data_shape: tuple | None property writable

Get original shape of feature dimensions.

weight_sum: float property writable

Get total weight of samples processed.

Functions

add(X: torch.Tensor, weights: torch.Tensor | None = None, *args, **kwargs) -> None

Add batch of samples to running covariance computation.

Updates both mean and covariance matrix incrementally using numerically stable Chan algorithm. Supports weighted samples for robust statistics.

PARAMETER	DESCRIPTION
X	Batch of data with shape (n_samples, n_features, ...). TYPE: Tensor
weights	Sample weights with shape (n_samples,). If None, all samples have equal weight (equivalent to weights=1). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

Source code in spectre/stats/cov.py

def add(
    self, X: torch.Tensor, weights: torch.Tensor | None = None, *args, **kwargs
) -> None:
    """
    Add batch of samples to running covariance computation.

    Updates both mean and covariance matrix incrementally using
    numerically stable Chan algorithm. Supports weighted samples
    for robust statistics.

    Parameters
    ----------
    X : torch.Tensor
        Batch of data with shape (n_samples, n_features, ...).

    weights : torch.Tensor | None, optional, by default None
        Sample weights with shape (n_samples,). If None, all samples
        have equal weight (equivalent to weights=1).
    """
    X = self._flatten_shape(X)
    if X.numel() == 0:
        return

    size = X.shape[0]

    # Handle weights
    if weights is None:
        weights = torch.ones(size, dtype=X.dtype, device=X.device)
    else:
        check_same_shape(X, weights, axis=0)
        check_sample_weights(weights)

    weights = weights.view(-1, 1)  # (n_samples, 1) for broadcasting

    batch_weight_sum = weights.sum()
    if batch_weight_sum == 0:
        return  # Skip batch with zero total weight

    # Initial batch
    if self._mean is None:
        self._count = size
        self._weight_sum = float(batch_weight_sum)
        self._mean = (X * weights).sum(dim=0) / batch_weight_sum
        centered = X - self._mean
        # Weighted outer product: C = sum_i w_i * delta_i * delta_i^T
        self._cmom2 = (centered.t() * weights.t()).mm(centered)
        return

    # Weighted Chan-style update for numerical stability
    self._count += size
    self._weight_sum += float(batch_weight_sum)

    # Update the mean according to weighted batch deviation
    delta = X - self._mean
    weighted_delta_sum = (delta * weights).sum(dim=0)

    self._mean = self._mean + weighted_delta_sum / self._weight_sum
    delta2 = X - self._mean

    # Update covariance using weighted batch deviation
    # Add weighted outer products: sum_i w_i * delta_i * delta2_i^T
    if self._cmom2 is not None:
        self._cmom2.addmm_(mat1=(delta.t() * weights.t()), mat2=delta2)

to(device: str = 'cpu') -> None

Move internal tensors to specified device.

PARAMETER	DESCRIPTION
device	Target device ("cpu", "cuda", "cuda:0", etc.). TYPE: str DEFAULT: 'cpu'

Source code in spectre/stats/cov.py

def to(self, device: str = "cpu") -> None:
    """
    Move internal tensors to specified device.

    Parameters
    ----------
    device : str
        Target device ("cpu", "cuda", "cuda:0", etc.).
    """
    if self._mean is not None:
        self._mean = self._mean.to(device)

    if self._cmom2 is not None:
        self._cmom2 = self._cmom2.to(device)

mean() -> torch.Tensor

Get current mean estimate.

RETURNS	DESCRIPTION
Tensor	Mean with original feature shape.

Source code in spectre/stats/cov.py

def mean(self) -> torch.Tensor:
    """
    Get current mean estimate.

    Returns
    -------
    torch.Tensor
        Mean with original feature shape.
    """
    return self._restore_shape(self._mean)

covariance(unbiased: bool = True) -> torch.Tensor

Get current covariance matrix estimate.

Computes covariance from the second central moment matrix. For weighted data, uses reliability weights correction when unbiased=True.

PARAMETER	DESCRIPTION
unbiased	If True, use Bessel's correction for unweighted or reliability weights correction for weighted data. If False, use biased estimate (divide by weight_sum). TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
Tensor	Covariance matrix with shape (n_features, n_features).

Notes

Unbiased weighted covariance uses reliability weights correction: \(\Sigma = C / (W - W^2/W)\) where \(W\) is sum of weights. For equal weights this reduces to Bessel's correction (n-1).

Biased covariance: \(\Sigma = C / W\)

Source code in spectre/stats/cov.py

def covariance(self, unbiased: bool = True) -> torch.Tensor:
    """
    Get current covariance matrix estimate.

    Computes covariance from the second central moment matrix. For weighted
    data, uses reliability weights correction when unbiased=True.

    Parameters
    ----------
    unbiased : bool, optional, by default True
        If True, use Bessel's correction for unweighted or reliability
        weights correction for weighted data.
        If False, use biased estimate (divide by weight_sum).

    Returns
    -------
    torch.Tensor
        Covariance matrix with shape (n_features, n_features).

    Notes
    -----
    Unbiased weighted covariance uses reliability weights correction:
    $\\Sigma = C / (W - W^2/W)$
    where $W$ is sum of weights. For equal weights this reduces to Bessel's
    correction (n-1).

    Biased covariance: $\\Sigma = C / W$
    """
    if unbiased:
        # Reliability weights correction for weighted covariance
        denom = self._weight_sum - (self._weight_sum / self._count)
    else:
        denom = self._weight_sum

    return self._cmom2 / denom

correlation(unbiased: bool = True) -> torch.Tensor

Get correlation matrix from covariance.

Computes Pearson correlation coefficients by normalizing covariance by standard deviations. Properly handles weighted data.

PARAMETER	DESCRIPTION
unbiased	If True, use unbiased covariance for computation. TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
Tensor	Correlation matrix with shape (n_features, n_features). Values range from -1 to 1.

Notes

Correlation is computed as: \(\rho_{ij} = \frac{\Sigma_{ij}}{\sqrt{\Sigma_{ii} \Sigma_{jj}}}\)

Source code in spectre/stats/cov.py

def correlation(self, unbiased: bool = True) -> torch.Tensor:
    """
    Get correlation matrix from covariance.

    Computes Pearson correlation coefficients by normalizing covariance
    by standard deviations. Properly handles weighted data.

    Parameters
    ----------
    unbiased : bool, optional, by default True
        If True, use unbiased covariance for computation.

    Returns
    -------
    torch.Tensor
        Correlation matrix with shape (n_features, n_features).
        Values range from -1 to 1.

    Notes
    -----
    Correlation is computed as:
    $\\rho_{ij} = \\frac{\\Sigma_{ij}}{\\sqrt{\\Sigma_{ii} \\Sigma_{jj}}}$
    """
    if unbiased:
        denom = self._weight_sum - (self._weight_sum / self._count)
    else:
        denom = self._weight_sum

    cov = self._cmom2 / denom
    rstdev = cov.diag().sqrt().reciprocal()

    return rstdev[:, None] * cov * rstdev[None, :]

variance(unbiased: bool = True) -> torch.Tensor | None

Get variance (diagonal of covariance matrix).

More efficient than computing full covariance when only variance is needed. Properly handles weighted data.

PARAMETER	DESCRIPTION
unbiased	If True, use Bessel's correction for unweighted or reliability weights correction for weighted data. If False, use biased estimate (divide by weight_sum). TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
Tensor | None	Variance for each feature with original shape.

Source code in spectre/stats/cov.py

def variance(self, unbiased: bool = True) -> torch.Tensor | None:
    """
    Get variance (diagonal of covariance matrix).

    More efficient than computing full covariance when only variance
    is needed. Properly handles weighted data.

    Parameters
    ----------
    unbiased : bool, optional, by default True
        If True, use Bessel's correction for unweighted or reliability
        weights correction for weighted data.
        If False, use biased estimate (divide by weight_sum).

    Returns
    -------
    torch.Tensor | None
        Variance for each feature with original shape.
    """
    if unbiased:
        denom = self._weight_sum - (self._weight_sum / self._count)
    else:
        denom = self._weight_sum

    if self._cmom2 is not None:
        return self._restore_shape(self._cmom2.diag() / denom)
    return None

stdev(unbiased: bool = True) -> torch.Tensor | None

Get standard deviation (square root of variance).

PARAMETER	DESCRIPTION
unbiased	If True, use Bessel's correction (divide by n-1). If False, use biased estimate (divide by n). TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
Tensor | None	Standard deviation for each feature with original shape.

Source code in spectre/stats/cov.py

def stdev(self, unbiased: bool = True) -> torch.Tensor | None:
    """
    Get standard deviation (square root of variance).

    Parameters
    ----------
    unbiased : bool, optional, by default True
        If True, use Bessel's correction (divide by n-1).
        If False, use biased estimate (divide by n).

    Returns
    -------
    torch.Tensor | None
        Standard deviation for each feature with original shape.
    """
    var = self.variance(unbiased=unbiased)
    if var is not None:
        return var.sqrt()
    return None

state_dict() -> dict

Save state to dictionary.

RETURNS	DESCRIPTION
dict	State dictionary containing count, weight_sum, mean, second moment matrix, and shape info.

Source code in spectre/stats/cov.py

def state_dict(self) -> dict:
    """
    Save state to dictionary.

    Returns
    -------
    dict
        State dictionary containing count, weight_sum, mean, second moment
        matrix, and shape info.
    """
    return dict(
        constructor=f"{self.__module__}.{self.__class__.__name__}()",
        count=self._count,
        data_shape=self._data_shape and tuple(self._data_shape),
        weight_sum=self._weight_sum,
        mean=self._mean,
        cmom2=self._cmom2,
    )

load_state_dict(state: dict) -> None

Load state from dictionary.

Restores the running covariance computation from a saved state.

PARAMETER	DESCRIPTION
state	State dictionary from state_dict() containing: "count": Total sample count "weight_sum": Cumulative sum of weights "mean": Current mean estimate "cmom2": Second central moment matrix "data_shape": Original feature shape TYPE: dict

Source code in spectre/stats/cov.py

def load_state_dict(self, state: dict) -> None:
    """
    Load state from dictionary.

    Restores the running covariance computation from a saved state.

    Parameters
    ----------
    state : dict
        State dictionary from `state_dict()` containing:

        - "count": Total sample count
        - "weight_sum": Cumulative sum of weights
        - "mean": Current mean estimate
        - "cmom2": Second central moment matrix
        - "data_shape": Original feature shape
    """
    self._count = state["count"]
    self._weight_sum = state.get("weight_sum", float(state["count"]))
    self._mean = state["mean"]
    self._cmom2 = state["cmom2"]
    self._data_shape = (
        None if state["data_shape"] is None else tuple(state["data_shape"])
    )

Functions