Stats Crosscov

crosscov

CLASS	DESCRIPTION
RunningCrossCovariance	Numerically stable running cross-covariance between two signals.

Classes

RunningCrossCovariance(state: dict | None = None, split_batch: bool = True)

Bases: Stats

Numerically stable running cross-covariance between two signals.

Computes cross-covariance between two aligned data sources (X and Y) incrementally using the Chan-Golub-LeVeque algorithm. More memory-efficient than RunningCovariance when only off-diagonal covariance blocks are needed.

Cross-covariance is computed using incremental updates:

\(\mu_{X,new} = \mu_{X,old} + \frac{\sum_i (x_i - \mu_{X,old})}{n_{total}}\), \(\mu_{Y,new} = \mu_{Y,old} + \frac{\sum_i (y_i - \mu_{Y,old})}{n_{total}}\), \(C_{XY,new} = C_{XY,old} + \delta_X^T \delta'_Y\),

where:

• \(\delta_X = x_i - \mu_{X,old}\) (X deviation from old mean)
• \(\delta'_Y = y_i - \mu_{Y,new}\) (Y deviation from new mean)
• \(C_{XY}\) is the cross-covariance second moment

PARAMETER	DESCRIPTION
state	Previous state to restore from. TYPE: dict | None, optional, by default None DEFAULT: None
split_batch	Reserved for future batch splitting functionality. TYPE: bool, optional, by default True DEFAULT: True

Properties

count : int Total number of samples processed.

mean : list[torch.Tensor, torch.Tensor] Current mean estimates for X and Y.

cmom2 : torch.Tensor Cross-covariance second central moment matrix (n_features_X, n_features_Y).

v_cmom2 : list[torch.Tensor, torch.Tensor] Variance second moments for X and Y (for diagonal covariances).

Examples:

Basic usage with two aligned signals:

>>> from spectre.stats import RunningCrossCovariance
>>> import torch
>>>
>>> stat = RunningCrossCovariance()
>>> for _ in range(10):
...     X = torch.randn(100, 5)  # Signal 1
...     Y = torch.randn(100, 3)  # Signal 2
...     stat.add(X, Y)
>>>
>>> means = stat.mean()  # [mean_X, mean_Y]
>>> cross_cov = stat.covariance()  # Shape: (5, 3)
>>> correlation = stat.correlation()

Memory-efficient cross-covariance:

>>> # When full covariance (X+Y, X+Y) doesn't fit in GPU memory
>>> stat = RunningCrossCovariance()
>>> X = torch.randn(1000, 1000)  # Large feature space
>>> Y = torch.randn(1000, 500)  # Different feature space
>>> stat.add(X, Y)
>>> cross_cov = stat.covariance()  # Only (1000, 500) instead of (1500, 1500)

METHOD	DESCRIPTION
add	Add aligned batch of X and Y samples to cross-covariance computation.
mean	Get current mean estimates for both signals.
variance	Get variance estimates for both signals.
stdev	Get standard deviation estimates for both signals.
covariance	Get cross-covariance matrix between X and Y.
correlation	Get cross-correlation matrix between X and Y.
to	Move internal tensors to specified device.
state_dict	Save state to dictionary.
load_state_dict	Load state from dictionary.

ATTRIBUTE	DESCRIPTION
count	Get total number of samples processed. TYPE: int
cmom2	Get cross-covariance second central moment matrix. TYPE: Tensor
v_cmom2	Get variance second central moments for X and Y. TYPE: Tensor | list[Tensor] | None

Source code in spectre/stats/crosscov.py

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

    self.split_batch = split_batch

    self._count = 0
    self._mean = None
    self._cmom2 = None
    self._v_cmom2 = None

Attributes

count: int property writable

Get total number of samples processed.

cmom2: torch.Tensor property

Get cross-covariance second central moment matrix.

v_cmom2: torch.Tensor | list[torch.Tensor] | None property

Get variance second central moments for X and Y.

Functions

add(X: torch.Tensor, Y: torch.Tensor) -> None

Add aligned batch of X and Y samples to cross-covariance computation.

Updates means, variances, and cross-covariance incrementally using numerically stable Chan algorithm. X and Y must have the same batch size.

PARAMETER	DESCRIPTION
X	First batch with shape (n_samples, n_features_X, ...). TYPE: Tensor
Y	Second batch with shape (n_samples, n_features_Y, ...). Must have same batch size as X. TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If X and Y have different batch sizes (len(X) != len(Y)).

Source code in spectre/stats/crosscov.py

def add(self, X: torch.Tensor, Y: torch.Tensor) -> None:  # ty: ignore
    """
    Add aligned batch of X and Y samples to cross-covariance computation.

    Updates means, variances, and cross-covariance incrementally using
    numerically stable Chan algorithm. X and Y must have the same batch size.

    Parameters
    ----------
    X : torch.Tensor
        First batch with shape (n_samples, n_features_X, ...).

    Y : torch.Tensor
        Second batch with shape (n_samples, n_features_Y, ...).
        Must have same batch size as X.

    Raises
    ------
    ValueError
        If X and Y have different batch sizes (len(X) != len(Y)).
    """
    if len(X.shape) == 1:
        X = X[None, :]
        Y = Y[None, :]
    check_same_len(X, Y)

    if X.ndim > 2:
        X, Y = [
            d.view(d.shape[0], d.shape[1], -1)
            .permute(0, 2, 1)
            .reshape(-1, d.shape[1])
            for d in [X, Y]
        ]
    size = X.shape[0]

    # Initial batch.
    if self._mean is None:
        self._count = size
        self._mean = [d.sum(0) / size for d in [X, Y]]
        centered = [d - bm for d, bm in zip([X, Y], self._mean)]
        self._v_cmom2 = [c.pow(2).sum(0) for c in centered]
        self._cmom2 = centered[0].t().mm(centered[1])
        return

    # Update a batch using Chan-style.
    self._count += size
    # Update the mean according to the batch deviation from the old mean.
    delta = [(d - bm) for d, bm in zip([X, Y], self._mean)]
    for m, d in zip(self._mean, delta):
        m.add_(d.sum(0) / self._count)
    delta2 = [(d - bm) for d, bm in zip([X, Y], self._mean)]
    # Update the cross-covariance using the batch deviation
    if self._cmom2 is not None:
        self._cmom2.addmm_(mat1=delta[0].t(), mat2=delta2[1])

    # Update the variance using the batch deviation
    for vc2, d, d2 in zip(self._v_cmom2, delta, delta2):
        vc2.add_((d * d2).sum(0))

mean() -> list[torch.Tensor]

Get current mean estimates for both signals.

RETURNS	DESCRIPTION
list[Tensor, Tensor]	List containing [mean_X, mean_Y].

Source code in spectre/stats/crosscov.py

def mean(self) -> list[torch.Tensor]:
    """
    Get current mean estimates for both signals.

    Returns
    -------
    list[torch.Tensor, torch.Tensor]
        List containing [mean_X, mean_Y].
    """
    return self._mean

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

Get variance estimates for both signals.

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
list[Tensor, Tensor] | None	List containing [variance_X, variance_Y].

Source code in spectre/stats/crosscov.py

def variance(self, unbiased: bool = True) -> list[torch.Tensor] | None:
    """
    Get variance estimates for both signals.

    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
    -------
    list[torch.Tensor, torch.Tensor] | None
        List containing [variance_X, variance_Y].
    """
    if self._cmom2 is not None:
        return [
            vc2 / (self._count - (1 if unbiased else 0))
            for vc2 in self._v_cmom2  # ty: ignore[not-iterable]
        ]
    return None

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

Get standard deviation estimates for both signals.

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
list[Tensor, Tensor] | None	List containing [stdev_X, stdev_Y].

Source code in spectre/stats/crosscov.py

def stdev(self, unbiased: bool = True) -> list[torch.Tensor] | None:
    """
    Get standard deviation estimates for both signals.

    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
    -------
    list[torch.Tensor, torch.Tensor] | None
        List containing [stdev_X, stdev_Y].
    """
    variances = self.variance(unbiased=unbiased)
    if variances is not None:
        return [v.sqrt() for v in variances]
    return None

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

Get cross-covariance matrix between X and Y.

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	Cross-covariance matrix with shape (n_features_X, n_features_Y).

Source code in spectre/stats/crosscov.py

def covariance(self, unbiased: bool = True) -> torch.Tensor:
    """
    Get cross-covariance matrix between X and Y.

    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
        Cross-covariance matrix with shape (n_features_X, n_features_Y).
    """
    return self._cmom2 / (self._count - (1 if unbiased else 0))

correlation() -> torch.Tensor

Get cross-correlation matrix between X and Y.

Computes Pearson correlation coefficients by normalizing cross-covariance by the product of standard deviations.

RETURNS	DESCRIPTION
Tensor	Cross-correlation matrix with shape (n_features_X, n_features_Y). Values range from -1 to 1. NaN values are replaced with 0.

Source code in spectre/stats/crosscov.py

def correlation(self) -> torch.Tensor:
    """
    Get cross-correlation matrix between X and Y.

    Computes Pearson correlation coefficients by normalizing cross-covariance
    by the product of standard deviations.

    Returns
    -------
    torch.Tensor
        Cross-correlation matrix with shape (n_features_X, n_features_Y).
        Values range from -1 to 1. NaN values are replaced with 0.
    """
    covariance = self.covariance(unbiased=False)
    rstdev = [s.reciprocal() for s in self.stdev(unbiased=False)]  # ty: ignore[not-iterable]
    cor = rstdev[0][:, None] * covariance * rstdev[1][None, :]
    # Remove NaNs
    cor[torch.isnan(cor)] = 0
    return cor

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/crosscov.py

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

    Parameters
    ----------
    device : str
        Target device ("cpu", "cuda", "cuda:0", etc.).
    """
    self._mean = [m.to(device) for m in self._mean]  # ty: ignore[not-iterable]
    self._v_cmom2 = [vcs.to(device) for vcs in self._v_cmom2]  # ty: ignore[not-iterable]
    if self._cmom2 is not None:
        self._cmom2 = self._cmom2.to(device)

state_dict() -> dict

Save state to dictionary.

RETURNS	DESCRIPTION
dict	State dictionary containing count, means, variances, and cross-covariance matrix.

Source code in spectre/stats/crosscov.py

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

    Returns
    -------
    dict
        State dictionary containing count, means, variances, and
        cross-covariance matrix.
    """
    return dict(
        constructor=f"{self.__module__}.{self.__class__.__name__}()",
        count=self._count,
        mean_X=self._mean[0],
        mean_Y=self._mean[1],
        cmom2_X=self._v_cmom2[0],
        cmom2_Y=self._v_cmom2[1],
        cmom2=self._cmom2,
    )

load_state_dict(state: dict) -> None

Load state from dictionary.

PARAMETER	DESCRIPTION
state	State dictionary from state_dict() containing: "count": Total sample count "mean_X", "mean_Y": Mean estimates for X and Y "cmom2_X", "cmom2_Y": Variance second moments for X and Y "cmom2": Cross-covariance second moment matrix TYPE: dict

Source code in spectre/stats/crosscov.py

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

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

        - "count": Total sample count
        - "mean_X", "mean_Y": Mean estimates for X and Y
        - "cmom2_X", "cmom2_Y": Variance second moments for X and Y
        - "cmom2": Cross-covariance second moment matrix
    """
    self._count = int(state["count"])
    self._mean = [state[f"mean_{k}"] for k in "XY"]
    self._v_cmom2 = [state[f"cmom2_{k}"] for k in "XY"]
    self._cmom2 = state["cmom2"]

Functions