Stats Var

var

CLASS	DESCRIPTION
RunningVariance	Numerically stable running variance and mean computation.

Classes

RunningVariance(state=None)

Bases: Stats

Numerically stable running variance and mean computation.

Computes both mean and variance incrementally using the Chan-Golub-LeVeque algorithm. More memory-efficient than RunningCovariance when only variance (not full covariance matrix) is needed.

The variance is computed using the numerically stable two-pass algorithm:

\(\mu_{new} = \mu_{old}+\frac{n_{batch}}{n_{total}}(\mu_{batch}-\mu_{old})\),

\(M_{2,new} = M_{2,old} + \sum_i (x_i - \mu_{batch})^2 + \frac{n_{batch} \cdot n_{old}}{n_{total}} (\mu_{batch} - \mu_{old})^2\),

\(\sigma^2 = \frac{M_2}{n - 1} \quad \text{(unbiased)}\)

where \(M_2\) is the second central moment (sum of squared deviations).

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.

batch_count : int Number of batches processed.

data_shape : tuple | None Original shape of feature dimensions.

mean : torch.Tensor | None Current mean estimate.

v_cmom2 : torch.Tensor | None Second central moment (sum of squared deviations).

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

Examples:

Basic usage:

>>> from spectre.stats import RunningVariance
>>> import torch
>>>
>>> stat = RunningVariance()
>>> for _ in range(10):
...     batch = torch.randn(100, 5)
...     stat.add(batch)
>>>
>>> mean = stat.mean()
>>> variance = stat.variance(unbiased=True)
>>> stdev = stat.stdev(unbiased=True)

Biased vs unbiased variance:

>>> stat = RunningVariance()
>>> X = torch.randn(1000, 5)
>>> stat.add(X)
>>>
>>> var_biased = stat.variance(unbiased=False)  # Divide by n
>>> var_unbiased = stat.variance(unbiased=True)  # Divide by n-1

METHOD	DESCRIPTION
add	Add batch of samples to running variance computation.
mean	Get current mean estimate.
variance	Get current variance estimate.
stdev	Get current standard deviation estimate.
to	Move internal tensors to specified device.
load_state_dict	Load state from dictionary.
state_dict	Save state to dictionary.

ATTRIBUTE	DESCRIPTION
count	Get total number of samples processed. TYPE: int
batch_count	Get number of batches processed. TYPE: int
data_shape	Get original feature shape (or None if unknown). TYPE: tuple | None
weight_sum	Get cumulative sum of weights. TYPE: float
v_cmom2	Get second central moment (sum of squared deviations). TYPE: Tensor | None

Source code in spectre/stats/var.py

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

    self._count = 0
    self._batch_count = 0
    self._weight_sum = 0.0
    self._mean = None
    self._v_cmom2 = None
    self._data_shape = None

Attributes

count: int property writable

Get total number of samples processed.

batch_count: int property writable

Get number of batches processed.

data_shape: tuple | None property writable

Get original feature shape (or None if unknown).

weight_sum: float property writable

Get cumulative sum of weights.

v_cmom2: torch.Tensor | None property writable

Get second central moment (sum of squared deviations).

Functions

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

Add batch of samples to running variance computation.

Updates both mean and second central moment 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/var.py

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

    Updates both mean and second central moment 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]

    if weights is None:
        weights = torch.ones(size, dtype=X.dtype, device=X.device)
    else:
        check_same_shape(X, weights, axis=0)
        weights = weights.view(-1)

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

    # Weighted batch mean
    batch_mean = (X * weights.unsqueeze(1)).sum(0) / batch_weight_sum
    centered = X - batch_mean
    self._batch_count += 1

    # Initial batch
    if self._mean is None:
        self._count = size
        self._weight_sum = float(batch_weight_sum)
        self._mean = batch_mean
        self._v_cmom2 = (centered.pow(2) * weights.unsqueeze(1)).sum(dim=0)
        return

    # Update using weighted Chan-style update for numerical stability
    old_weight_sum = self._weight_sum
    self._count += size
    self._weight_sum += float(batch_weight_sum)

    # Update the mean according to the batch deviation from the old mean
    delta = float(batch_weight_sum) / self._weight_sum * (batch_mean - self._mean)
    self._mean = self._mean + delta

    # Update the variance using the batch deviation
    self._v_cmom2 = self._v_cmom2 + (centered.pow(2) * weights.unsqueeze(1)).sum(
        dim=0
    )
    self._v_cmom2 = self._v_cmom2 + delta.pow_(2).mul_(
        float(batch_weight_sum) / self.weight_sum * old_weight_sum
    )

mean() -> torch.Tensor

Get current mean estimate.

RETURNS	DESCRIPTION
Tensor	Mean with original feature shape.

Source code in spectre/stats/var.py

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

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

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

Get current variance estimate.

Computes variance from the second central moment. 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	Variance with original feature shape.

Source code in spectre/stats/var.py

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

    Computes variance from the second central moment. 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
        Variance with original feature shape.
    """
    if unbiased:
        # Reliability weights correction for weighted variance
        # For equal weights, this reduces to n-1
        denom = self._weight_sum - (self._weight_sum / self._count)
    else:
        denom = self._weight_sum

    return self._restore_shape(self._v_cmom2 / denom)

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

Get current standard deviation estimate.

Computes standard deviation as 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	Standard deviation with original feature shape.

Source code in spectre/stats/var.py

def stdev(self, unbiased: bool = True) -> torch.Tensor:
    """
    Get current standard deviation estimate.

    Computes standard deviation as 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
        Standard deviation with original feature shape.
    """
    return self.variance(unbiased=unbiased).sqrt()

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/var.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._v_cmom2 is not None:
        self._v_cmom2 = self._v_cmom2.to(device)

load_state_dict(state: dict) -> None

Load state from dictionary.

Restores the running variance computation from a saved state.

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

Source code in spectre/stats/var.py

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

    Restores the running variance computation from a saved state.

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

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

state_dict() -> dict

Save state to dictionary.

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

Source code in spectre/stats/var.py

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

    Returns
    -------
    dict
        State dictionary containing count, weight_sum, mean, second moment,
        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),
        batch_count=self._batch_count,
        weight_sum=self._weight_sum,
        mean=self._mean,
        cmom2=self._v_cmom2,
    )

Functions