sage.core.graph

Sage pipeline building blocks.

Preprocessor

Sequential chain of preprocessing modules (whitening, multirate, etc.). Passes data through each module in order. Modules may accept and return either raw tensors (legacy) or GWBatch objects — nn.Sequential is type-agnostic, so the chain works in both cases transparently.

TorchChoice

Probabilistic per-sample module selection. Useful for augmentation or domain randomisation.

Classes

Preprocessor	Sequential preprocessing pipeline for gravitational-wave data.
TorchChoice	Choose one module per sample according to provided probabilities.

Module Contents

class Preprocessor(modules)

Bases: torch.nn.Module

Sequential preprocessing pipeline for gravitational-wave data.

Chains an ordered list of nn.Module transforms. Each module receives the output of the previous one. A typical pipeline for TD training is:

Preprocessor([FiducialWhitening(), MultirateSampler(...)])

and for FD_COARSE (worst-case multibanding):

Preprocessor([FiducialWhitening()])

The pipeline is grid-aware when modules return GWBatch objects (which is the default for updated modules). The training loop is responsible for wrapping the combined signal+noise tensor in a GWBatch before calling this module, and for extracting batch.to_network_input() afterward.

Parameters:

• modules (list of nn.Module) – Ordered preprocessing steps.

• state (Initialize internal Module)

• ScriptModule. (shared by both nn.Module and)

seq

forward(x)

Run the full preprocessing pipeline.

Parameters:

x (torch.Tensor or GWBatch) – Input data. Shape and type depend on the first module.

Returns:

Preprocessed data. Type matches whatever the last module returns.

Return type:

torch.Tensor or GWBatch

class TorchChoice(modules, probabilities)

Bases: torch.nn.Module

Choose one module per sample according to provided probabilities.

Useful for data augmentation that should apply different transforms to different samples within a batch.

Parameters:

• modules (list of nn.Module) – Candidate modules.

• probabilities (list of float) – Probability of selecting each module (automatically normalised).

• state (Initialize internal Module)

• ScriptModule. (shared by both nn.Module and)

modules_list

forward(x, generator=None)