Stats Mean

mean

CLASS	DESCRIPTION
RunningMean	Numerically stable running mean computation.

Classes

RunningMean(state: dict | None = None)

Bases: Stats

Numerically stable running mean computation.

Uses the Chan-Golub-LeVeque algorithm for incremental mean updates, which avoids catastrophic cancellation and maintains numerical precision and stability even with large numbers of samples.

The mean is updated incrementally using:

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

where:

• n_{batch} is the size of the current batch
• n_{total} is the cumulative count of all samples
• mu_{batch} is the mean of the current batch
• mu_{old} is the previous running mean

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.

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

Examples:

Basic usage:

>>> from spectre.stats import RunningMean
>>> import torch
>>>
>>> stat = RunningMean()
>>> for _ in range(10):
...     batch = torch.randn(100, 5)  # 100 samples, 5 features
...     stat.add(batch)
>>>
>>> mean = stat.mean()
>>> mean.shape
torch.Size([5])

Higher precision computation:

>>> stat = RunningMean()
>>> for _ in range(10):
...     batch = torch.randn(100, 5, dtype=torch.float64)
...     stat.add(batch)
>>> mean = stat.mean()  # float64 precision maintained

Multi-dimensional features:

>>> stat = RunningMean()
>>> for _ in range(10):
...     batch = torch.randn(100, 3, 32, 32)  # Image-like data
...     stat.add(batch)
>>> mean = stat.mean()
>>> mean.shape
torch.Size([3, 32, 32])

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

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

Source code in spectre/stats/mean.py

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

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

    self._mean = 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.

Functions

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

Add batch of samples to running mean.

Updates the mean estimate incrementally using the Chan algorithm. Handles multi-dimensional features by flattening and restoring shapes. Supports weighted samples for robust statistics.

PARAMETER	DESCRIPTION
X	Batch of data with shape (n_samples, n_features, ...). First dimension is batch, remaining dimensions are features. TYPE: Tensor
weights	Sample weights with shape (n_samples,). If None, all samples have equal weight (equivalent to weights=1). By default None. TYPE: Tensor | None DEFAULT: None

Source code in spectre/stats/mean.py

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

    Updates the mean estimate incrementally using the Chan algorithm.
    Handles multi-dimensional features by flattening and restoring shapes.
    Supports weighted samples for robust statistics.

    Parameters
    ----------
    X : torch.Tensor
        Batch of data with shape (n_samples, n_features, ...).
        First dimension is batch, remaining dimensions are features.

    weights : torch.Tensor | None, optional
        Sample weights with shape (n_samples,). If None, all samples
        have equal weight (equivalent to weights=1). By default None.
    """
    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(dim=0) / batch_weight_sum
    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

        return

    self._count += size
    self._weight_sum += float(batch_weight_sum)

    # Weighted Chan-style update
    delta = (batch_mean - self._mean) * float(batch_weight_sum) / self._weight_sum

    self._mean = self._mean + delta

mean() -> torch.Tensor

Get current mean estimate.

RETURNS	DESCRIPTION
Tensor	Mean with original feature shape.

Source code in spectre/stats/mean.py

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

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

load_state_dict(state: dict) -> None

Load state from dictionary.

Restores the running mean computation from a saved state, enabling resumption of computation.

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 "data_shape": Original feature shape TYPE: dict

Source code in spectre/stats/mean.py

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

    Restores the running mean computation from a saved state,
    enabling resumption of computation.

    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
        - "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._data_shape = (
        None if state["data_shape"] is None else tuple(state["data_shape"])
    )

state_dict() -> dict

Save state to dictionary.

Returns a dictionary containing all internal state that can be saved to disk and later restored via load_state_dict().

RETURNS	DESCRIPTION
dict	State dictionary containing: "constructor": String for recreating the object "count": Total sample count "batch_count": Total batch count "weight_sum": Cumulative sum of weights "data_shape": Original feature shape (or None) "mean": Current mean estimate tensor

Source code in spectre/stats/mean.py

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

    Returns a dictionary containing all internal state that can be
    saved to disk and later restored via `load_state_dict()`.

    Returns
    -------
    dict
        State dictionary containing:

        - "constructor": String for recreating the object
        - "count": Total sample count
        - "batch_count": Total batch count
        - "weight_sum": Cumulative sum of weights
        - "data_shape": Original feature shape (or None)
        - "mean": Current mean estimate tensor
    """
    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,
    )

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/mean.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)

Functions