Compute Utils#

utils #

FUNCTION DESCRIPTION
check_random_state

Set random seed across all random number generators.
outer

Compute outer product between tensors using broadcasting.
safe_xlogy

Compute x * log(y) safely, handling x=0 case.
safe_broadcast

Broadcast parameter tensor to matrix shape.
stable_cumsum

Compute cumulative sum with high precision and stability validation.

Functions#

check_random_state(random_state: int = 111) -> None #

Set random seed across all random number generators.

This function sets seeds for Python's random module, NumPy, PyTorch (CPU/CUDA), and PyTorch Lightning to ensure reproducible results across the entire stack.

PARAMETER DESCRIPTION
random_state

Random seed value, by default 111.

TYPE: int DEFAULT: 111

Source code in spectre/compute/utils.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def check_random_state(random_state: int = 111) -> None:
    """
    Set random seed across all random number generators.

    This function sets seeds for Python's random module, NumPy, PyTorch
    (CPU/CUDA), and PyTorch Lightning to ensure reproducible results across the
    entire stack.

    Parameters
    ----------
    random_state : int, optional
        Random seed value, by default 111.
    """
    # Python's random module
    random.seed(random_state)

    # NumPy random
    np.random.seed(random_state)

    # PyTorch CPU random
    torch.manual_seed(random_state)

    # PyTorch CUDA random (for all devices)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)

    # Set deterministic algorithms for reproducibility
    torch.use_deterministic_algorithms(True, warn_only=True)

    # Additional PyTorch settings for reproducibility
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # PyTorch Lightning seed (if available)
    try:
        from pytorch_lightning import seed_everything

        seed_everything(random_state, workers=True)
    except ImportError:
        # PyTorch Lightning not available, skip
        pass

outer(X: torch.Tensor, Y: torch.Tensor | None = None, flatten: bool = False) -> torch.Tensor #

Compute outer product between tensors using broadcasting.

PARAMETER DESCRIPTION
X

First tensor with shape (n,).

TYPE: Tensor

Y

Second tensor with shape (m,). If None, uses X, by default None.

TYPE: Tensor | None DEFAULT: None

flatten

Whether the result must be flattened.

TYPE: bool, by default False DEFAULT: False

RETURNS DESCRIPTION
Tensor

Outer product with shape (n, m).

Examples:

>>> X = torch.tensor([1, 2, 3])
>>> Y = torch.tensor([4, 5])
>>> result = outer(X, Y)
>>> result.shape
torch.Size([3, 2])
Source code in spectre/compute/utils.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def outer(
    X: torch.Tensor, Y: torch.Tensor | None = None, flatten: bool = False
) -> torch.Tensor:
    """
    Compute outer product between tensors using broadcasting.

    Parameters
    ----------
    X : torch.Tensor
        First tensor with shape (n,).

    Y : torch.Tensor | None, optional
        Second tensor with shape (m,). If None, uses X, by default None.

    flatten : bool, by default False
        Whether the result must be flattened.

    Returns
    -------
    torch.Tensor
        Outer product with shape (n, m).

    Examples
    --------
    >>> X = torch.tensor([1, 2, 3])
    >>> Y = torch.tensor([4, 5])
    >>> result = outer(X, Y)
    >>> result.shape
    torch.Size([3, 2])
    """
    if Y is None:
        Y = X

    out = X[:, None] * Y[None, :]

    return out.flatten() if flatten else out

safe_xlogy(x: torch.Tensor, y: torch.Tensor | None = None) -> torch.Tensor #

Compute x * log(y) safely, handling x=0 case.

When x=0, returns 0 regardless of y value, avoiding NaN from 0*log(0).

PARAMETER DESCRIPTION
x

First tensor.

TYPE: Tensor

y

Second tensor. If None, uses x, by default None.

TYPE: Tensor | None DEFAULT: None

RETURNS DESCRIPTION
Tensor

Safe x * log(y) computation with same shape as inputs.

Examples:

>>> x = torch.tensor([0.0, 1.0, 2.0])
>>> y = torch.tensor([0.0, 1.0, 2.0])
>>> result = safe_xlogy(x, y)
>>> # Returns [0.0, 0.0, log(4)] instead of [NaN, 0.0, log(4)]
Source code in spectre/compute/utils.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def safe_xlogy(x: torch.Tensor, y: torch.Tensor | None = None) -> torch.Tensor:
    """
    Compute x * log(y) safely, handling x=0 case.

    When x=0, returns 0 regardless of y value, avoiding NaN from 0*log(0).

    Parameters
    ----------
    x : torch.Tensor
        First tensor.
    y : torch.Tensor | None, optional
        Second tensor. If None, uses x, by default None.

    Returns
    -------
    torch.Tensor
        Safe x * log(y) computation with same shape as inputs.

    Examples
    --------
    >>> x = torch.tensor([0.0, 1.0, 2.0])
    >>> y = torch.tensor([0.0, 1.0, 2.0])
    >>> result = safe_xlogy(x, y)
    >>> # Returns [0.0, 0.0, log(4)] instead of [NaN, 0.0, log(4)]
    """
    if y is None:
        y = x
    return torch.where(x == 0, torch.zeros_like(x), x * torch.log(y))

safe_broadcast(param: torch.Tensor, D: torch.Tensor, param_name: str = 'parameter') -> torch.Tensor #

Broadcast parameter tensor to matrix shape.

Supports scalar, row-wise (n, 1), column-wise (1, m), and full matrix (n, m) parameter shapes.

PARAMETER DESCRIPTION
param

Parameter tensor to broadcast.

TYPE: Tensor

D

Matrix of shape (n_samples, m_samples).

TYPE: Tensor

param_name

Name of parameter for error messages.

TYPE: str, optional, by default "parameter" DEFAULT: 'parameter'

RETURNS DESCRIPTION
Tensor

Broadcast parameter tensor of shape (n_samples, m_samples).
RAISES DESCRIPTION
ValueError

If parameter shape cannot be broadcast to matrix shape.
Source code in spectre/compute/utils.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
def safe_broadcast(
    param: torch.Tensor, D: torch.Tensor, param_name: str = "parameter"
) -> torch.Tensor:
    """
    Broadcast parameter tensor to matrix shape.

    Supports scalar, row-wise (n, 1), column-wise (1, m), and full matrix
    (n, m) parameter shapes.

    Parameters
    ----------
    param : torch.Tensor
        Parameter tensor to broadcast.

    D : torch.Tensor
        Matrix of shape (n_samples, m_samples).

    param_name : str, optional, by default "parameter"
        Name of parameter for error messages.

    Returns
    -------
    torch.Tensor
        Broadcast parameter tensor of shape (n_samples, m_samples).

    Raises
    ------
    ValueError
        If parameter shape cannot be broadcast to matrix shape.
    """
    try:
        return torch.broadcast_to(param, size=D.shape)
    except RuntimeError as e:
        raise ValueError(
            f"{param_name.capitalize()} tensor with shape {param.shape} "
            f"cannot be broadcast to distance matrix shape {D.shape}."
        ) from e

stable_cumsum(arr: torch.Tensor, dim: int | None = None, rtol: float = 1e-05, atol: float = 1e-08, high_precision: bool = True) -> torch.Tensor #

Compute cumulative sum with high precision and stability validation.

Uses high precision (float64) for cumsum computation and validates that the final value matches the sum to ensure numerical stability.

PARAMETER DESCRIPTION
arr

Tensor to be cumulatively summed.

TYPE: Tensor

dim

Dimension along which cumulative sum is computed. If None, computes over flattened array, by default None.

TYPE: int | None DEFAULT: None

rtol

Relative tolerance for stability check, by default 1e-05.

TYPE: float DEFAULT: 1e-05

atol

Absolute tolerance for stability check, by default 1e-08.

TYPE: float DEFAULT: 1e-08

high_precision

Whether to use float64 for cumsum computation. If False, uses input dtype, by default True.

TYPE: bool DEFAULT: True

RETURNS DESCRIPTION
Tensor

Cumulative sum with same shape as input (or flattened if dim=None).
WARNS DESCRIPTION
RuntimeWarning

If cumsum is found to be unstable (final element does not match sum).

Examples:

>>> arr = torch.tensor([1.0, 2.0, 3.0])
>>> cumsum = stable_cumsum(arr)
>>> cumsum
tensor([1., 3., 6.], dtype=torch.float64)
Source code in spectre/compute/utils.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def stable_cumsum(
    arr: torch.Tensor,
    dim: int | None = None,
    rtol: float = 1e-05,
    atol: float = 1e-08,
    high_precision: bool = True,
) -> torch.Tensor:
    """
    Compute cumulative sum with high precision and stability validation.

    Uses high precision (float64) for cumsum computation and validates that the
    final value matches the sum to ensure numerical stability.

    Parameters
    ----------
    arr : torch.Tensor
        Tensor to be cumulatively summed.
    dim : int | None, optional
        Dimension along which cumulative sum is computed. If None, computes
        over flattened array, by default None.
    rtol : float, optional
        Relative tolerance for stability check, by default 1e-05.
    atol : float, optional
        Absolute tolerance for stability check, by default 1e-08.
    high_precision : bool, optional
        Whether to use float64 for cumsum computation. If False, uses input
        dtype, by default True.

    Returns
    -------
    torch.Tensor
        Cumulative sum with same shape as input (or flattened if dim=None).

    Warns
    -----
    RuntimeWarning
        If cumsum is found to be unstable (final element does not match sum).

    Examples
    --------
    >>> arr = torch.tensor([1.0, 2.0, 3.0])
    >>> cumsum = stable_cumsum(arr)
    >>> cumsum
    tensor([1., 3., 6.], dtype=torch.float64)
    """
    import warnings

    if dim is None:
        arr = arr.flatten()
        dim = 0

    if high_precision:
        out = torch.cumsum(arr, dim=dim, dtype=torch.float64)
        expected = torch.sum(arr, dim=dim, dtype=torch.float64)
    else:
        out = torch.cumsum(arr, dim=dim)
        expected = torch.sum(arr, dim=dim)
    if not torch.all(
        torch.isclose(
            out.take(torch.tensor([-1]).long().to(arr.device)),
            expected,
            rtol=rtol,
            atol=atol,
            equal_nan=True,
        )
    ):
        warnings.warn(
            "`cumsum` was found to be unstable: its last element does not "
            "correspond to sum.",
            RuntimeWarning,
        )
    return out