Stats Base

base

CLASS	DESCRIPTION
RunningStats	Composite statistic that bundles multiple Stats objects together.

Classes

Stats(state: dict)

Abstract base class for numerically stable running statistics.

Provides infrastructure for computing statistics incrementally over batches of data without storing the entire dataset in memory. All subclasses implement memory-efficient algorithms for their respective statistics.

Usage Pattern:

1. Create the desired stat object: stat = RunningMean()
2. Feed batches via add method: stat.add(batch)
3. Repeat step 2 any number of times
4. Read out the statistic: result = stat.mean()

Available Statistics:

• RunningMean: Computes mean()
• RunningVariance: Computes mean(), variance(), stdev()
• RunningCovariance: Computes mean(), covariance(), correlation(), variance(), stdev()
• RunningTopK: Computes topk() returning (values, indexes)
• RunningBincount: Computes bincount() for histogram of integral data
• RunningHistory: Computes history() returning concatenation of all data
• RunningCrossCovariance: Cross-covariance between two signals
• RunningStats: Aggregates multiple stat objects together

Data Format:

Statistics are vectorized along dim >= 1, so add() expects:

• Dimension 0: Batch/sampling dimension (n_samples)
• Dimension 1: Feature dimensions (n_features, ...)

The data type and device used matches data passed to add(). For higher-precision statistics, convert to torch.float64 before calling add().

State Persistence:

All statistics support serialization via state_dict() and load_state_dict() for saving/loading state from disk (e.g., as .npz files).

Notes

Adapted from: https://gist.github.com/davidbau/00a9b6763a260be8274f6ba22df9a145

Examples:

Computing running mean over batches:

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

Combining multiple statistics:

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

Initialize state from saved state.

All Stat subclasses can be initialized by passing state parameter, which calls load_state_dict() to restore the statistic.

PARAMETER	DESCRIPTION
state	Dictionary containing saved state from state_dict(). TYPE: dict

RAISES	DESCRIPTION
TypeError	If state is not a dictionary.

METHOD	DESCRIPTION
add	Add batch of samples to the running statistic.
load_state_dict	Load statistic state from dictionary.
state_dict	Save statistic state to dictionary.
to	Move statistic to specified device.

Source code in spectre/stats/base.py

def __init__(self, state: dict):
    """
    Initialize state from saved state.

    All Stat subclasses can be initialized by passing `state` parameter,
    which calls `load_state_dict()` to restore the statistic.

    Parameters
    ----------
    state : dict
        Dictionary containing saved state from `state_dict()`.

    Raises
    ------
    TypeError
        If `state` is not a dictionary.
    """
    if not isinstance(state, dict):
        raise TypeError(f"Expected `state` dict, got {type(state)}.")

    self.load_state_dict(state)

Functions

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

Add batch of samples to the running statistic.

Dimension 0 is the batch dimension (n_samples), and remaining dimensions are feature dimensions. Subclasses implement the actual update logic.

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. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
*args	Additional positional arguments for specific stat implementations. DEFAULT: ()
**kwargs	Additional keyword arguments for specific stat implementations. DEFAULT: {}

Notes

This method should be implemented by all subclasses to perform the incremental statistic update.

Source code in spectre/stats/base.py

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

    Dimension 0 is the batch dimension (n_samples), and remaining dimensions
    are feature dimensions. Subclasses implement the actual update logic.

    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.

    *args
        Additional positional arguments for specific stat implementations.

    **kwargs
        Additional keyword arguments for specific stat implementations.

    Notes
    -----
    This method should be implemented by all subclasses to perform the
    incremental statistic update.
    """
    pass

load_state_dict(state: dict) -> None

Load statistic state from dictionary.

Restores the internal state of this Stat from a dictionary as saved by state_dict(). Enables resuming computation from a saved checkpoint.

PARAMETER	DESCRIPTION
state	Dictionary containing state tensors and metadata. TYPE: dict

Notes

Subclasses must implement this method to restore all internal state.

Source code in spectre/stats/base.py

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

    Restores the internal state of this Stat from a dictionary as saved
    by `state_dict()`. Enables resuming computation from a saved checkpoint.

    Parameters
    ----------
    state : dict
        Dictionary containing state tensors and metadata.

    Notes
    -----
    Subclasses must implement this method to restore all internal state.
    """
    pass

state_dict() -> dict

Save statistic state to dictionary.

Returns a dictionary containing all internal state (tensors and metadata) that can be saved to disk (e.g., as .npz) and reloaded later using load_state_dict().

RETURNS	DESCRIPTION
dict	Dictionary containing state tensors and metadata.

Notes

Subclasses should implement this to save all necessary state for complete reconstruction via load_state_dict().

Source code in spectre/stats/base.py

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

    Returns a dictionary containing all internal state (tensors and metadata)
    that can be saved to disk (e.g., as `.npz`) and reloaded later using
    `load_state_dict()`.

    Returns
    -------
    dict
        Dictionary containing state tensors and metadata.

    Notes
    -----
    Subclasses should implement this to save all necessary state for
    complete reconstruction via `load_state_dict()`.
    """
    return {}

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

Move statistic to specified device.

Transfers all internal tensors to the target device (CPU or GPU).

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

Notes

Subclasses should implement this to move all internal tensors to the specified device.

Source code in spectre/stats/base.py

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

    Transfers all internal tensors to the target device (CPU or GPU).

    Parameters
    ----------
    device : str, optional
        Target device ("cpu", "cuda", "cuda:0", etc.), by default "cpu".

    Notes
    -----
    Subclasses should implement this to move all internal tensors to the
    specified device.
    """
    pass

RunningStats(state: dict | None = None, **kwargs)

Bases: Stats

Composite statistic that bundles multiple Stats objects together.

Convenient for computing multiple statistics simultaneously over the same data, with unified state persistence and device management. All bundled statistics are updated together when add() is called.

PARAMETER	DESCRIPTION
state	Previous state to restore from. TYPE: dict | None, optional, by default None DEFAULT: None
**kwargs	Named Stats objects to bundle together. Each key becomes an attribute for accessing the corresponding statistic. DEFAULT: {}

ATTRIBUTE	DESCRIPTION
kwargs	Dictionary of bundled Stats objects. TYPE: dict

RAISES	DESCRIPTION
TypeError	If any value in kwargs is not a Stats object, or if multiple stats of the same type are provided.

Examples:

Computing mean and variance together:

>>> from spectre.stats import RunningStats, RunningMean, RunningVariance
>>> import torch
>>>
>>> stat = RunningStats(mean=RunningMean(), var=RunningVariance())
>>>
>>> for _ in range(10):
...     batch = torch.randn(100, 5)
...     stat.add(batch)  # Updates both mean and variance
>>>
>>> mean = stat.mean.mean()
>>> variance = stat.var.variance()
>>> stdev = stat.var.stdev()

Using with loader:

>>> from spectre.stats import RunningStats, RunningMean, RunningCovariance, loader
>>>
>>> X = torch.randn(1000, 10)
>>> stat = RunningStats(mean=RunningMean(), cov=RunningCovariance())
>>>
>>> for batch in loader(stat, X, batch_size=100):
...     stat.add(batch)
>>>
>>> mean = stat.mean.mean()
>>> covariance = stat.cov.covariance()
>>> correlation = stat.cov.correlation()

State persistence:

>>> # Save state
>>> state_dict = stat.state_dict()
>>> torch.save(state_dict, "stats.pt")
>>>
>>> # Load state later
>>> loaded_state = torch.load("stats.pt")
>>> stat2 = RunningStats(state=loaded_state)
>>> # Continue computation
>>> stat2.add(new_batch)

METHOD	DESCRIPTION
add	Add batch to all bundled statistics.
load_state_dict	Load state for all bundled statistics from dictionary.
state_dict	Save state for all bundled statistics to dictionary.
to	Switches internal storage to specified device.

Source code in spectre/stats/base.py

def __init__(self, state: dict | None = None, **kwargs) -> None:
    for key, val in kwargs.items():
        if not isinstance(val, Stats):
            raise TypeError(
                f"Object under key {key} is not a Stats object but {type(val)}."
            )

    kwargs_types = [type(val) for val in kwargs.values()]
    if len(set(kwargs_types)) != len(kwargs_types):
        raise TypeError("Multiple running Stats instances of the same kind.")

    self.kwargs = kwargs

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

Functions

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

Add batch to all bundled statistics.

Source code in spectre/stats/base.py

def add(
    self, X: torch.Tensor, weights: torch.Tensor | None = None, *args, **kwargs
) -> None:
    """Add batch to all bundled statistics."""
    for obj in self.kwargs.values():
        obj.add(X, weights=weights, *args, **kwargs)

load_state_dict(state: dict) -> None

Load state for all bundled statistics from dictionary.

Source code in spectre/stats/base.py

def load_state_dict(self, state: dict) -> None:
    """Load state for all bundled statistics from dictionary."""
    for prefix, obj in self.kwargs.items():
        obj.load_state_dict(_pull_key_prefix(prefix, state))

state_dict() -> dict

Save state for all bundled statistics to dictionary.

Source code in spectre/stats/base.py

def state_dict(self) -> dict:
    """Save state for all bundled statistics to dictionary."""
    result = {}
    for prefix, obj in self.kwargs.items():
        result.update(_push_key_prefix(prefix, obj.state_dict()))
    return result

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

Switches internal storage to specified device.

Source code in spectre/stats/base.py

def to(self, device: str = "cpu") -> None:
    """Switches internal storage to specified device."""
    for v in self.kwargs.values():
        v.to(device)