sage.data.waveform.sampler

Filename : distributions.py Description : Short description of the file

Created on 2026-02-16 10:35:39

__author__ = Narenraju Nagarajan __copyright__ = Copyright 2026, ProjectName __license__ = MIT Licence __version__ = 0.0.1 __maintainer__ = Narenraju Nagarajan __affiliation__ = N/A __email__ = N/A __status__ = [‘inProgress’, ‘Archived’, ‘inUsage’, ‘Debugging’]

GitHub Repository: NULL

Documentation:

sampler = read_from_config(“gw_config.yaml”, device=”cuda”)

batch = sampler.sample(4096)

print(batch[“mass1”].shape) # (4096,) print(batch[“spin1x”].shape) # derived param print(batch[“distance”].shape) # transformed param

Classes

NamedConstraint	Reference to a named constraint function defined in
ExpressionConstraint	Constraint defined as an inline Python expression string.
DistributionSampler	Batched waveform-parameter sampler driven by a YAML prior configuration.

Functions

spherical_to_cartesian(radial, polar, azimuthal)	Convert spherical spin components to Cartesian coordinates.
read_from_config(path, seed)	Construct a DistributionSampler from a YAML prior configuration.

Module Contents

spherical_to_cartesian(radial, polar, azimuthal)

Convert spherical spin components to Cartesian coordinates.

Parameters:

• radial (torch.Tensor) – Spin magnitude.

• polar (torch.Tensor) – Polar angle (inclination) in radians.

• azimuthal (torch.Tensor) – Azimuthal angle in radians.

Returns:

(x, y, z) Cartesian spin components.

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

read_from_config(path, seed)

Construct a DistributionSampler from a YAML prior configuration.

Parameters:

• path (str) – Path to the YAML file describing the prior distributions, transforms, and constraints.

• seed (int) – Seed for the internal torch.Generator.

Returns:

Configured sampler ready to draw waveform parameters.

Return type:

DistributionSampler

class NamedConstraint(name, params=None)

Reference to a named constraint function defined in sage.data.waveform.constraints.

Parameters:

• name (str) – Name of the constraint function (must appear in constraints._NAMED_CONSTRAINTS).

• params (list or None) – Additional parameter names passed to the constraint function.

name

params = []

class ExpressionConstraint(expr)

Constraint defined as an inline Python expression string.

The expression is evaluated with eval against the current parameter dictionary, so any parameter name can be referenced directly.

Parameters:

expr (str) – A Python expression that evaluates to True (accept) or False (reject) given the current parameter dict. Example: "mass1 >= mass2".

name = 'custom'

expr

check(params)

Return True if the parameter dict satisfies the stored expression.

class DistributionSampler(config, device, dtype, generator)

Bases: torch.nn.Module

Batched waveform-parameter sampler driven by a YAML prior configuration.

Reads a dictionary that describes:

• priors — per-parameter distributions (uniform, sin-angle, solid-angle, sky, uniform-radius, …).

• waveform_transforms — deterministic reparametrisations applied after sampling (spherical → Cartesian spins, m1/m2 → mchirp/q, chirp-distance → luminosity distance).

• constraints — named or custom Boolean filters applied via rejection sampling (e.g. mass_order to enforce m1 ≥ m2).

After construction, _compile_batch_standardiser() must be called once to pre-compute mean and standard deviation buffers used by standardise_from_batch() / unstandardise_from_batch().

Parameters:

• config (dict) – Parsed YAML dictionary with keys "priors", "waveform_transforms" (optional), and "constraints" (optional).

• device (str) – Torch device for all sampled tensors.

• dtype (torch.dtype) – Floating-point dtype for all sampled tensors.

• generator (torch.Generator) – Seeded generator used for all random draws.

param_names

Sorted list of all parameter names produced by this sampler (includes derived quantities from transforms).

Type:

list[str]

param_index

Mapping from parameter name to column index in the output tensor.

Type:

dict[str, int]

num_params

Total number of parameters in the output batch.

Type:

int

bounds

Theoretical (min, max) bounds for every parameter.

Type:

dict[str, tuple]

normalisers

Pre-built Normalise objects for each param.

Type:

dict[str, Normalise]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

sage_cfg = None

generator

cfg

device

dtype

variable_params

distributions

transforms = []

constraints = []

param_names = []

param_index

num_params

req_idx = None

bounds

normalisers

static get_named_constraints()

Return the list of built-in named constraint identifiers.

theoretical_bounds()

Compute analytic lower/upper bounds for all parameters based on YAML priors + constraints + deterministic transforms.

build_normalisers()

Construct Normalise objects for all parameters using theoretical bounds.

Returns:

Mapping parameter name -> Normalise object

Return type:

dict[str, Normalise]

norm_from_batch(batch)

Min-max normalise selected parameters in a full parameter batch.

Parameters:

batch (torch.Tensor, shape (B, total_params)) – Full parameter batch produced by forward().

Returns:

Selected columns normalised to [0, 1] using theoretical bounds.

Return type:

torch.Tensor, shape (B, selected_params)

unnorm_from_batch(normed_batch)

Invert min-max normalisation to recover physical parameter values.

Parameters:

normed_batch (torch.Tensor, shape (B, selected_params)) – Normalised batch from norm_from_batch().

Returns:

Parameters in their original physical units/range.

Return type:

torch.Tensor, shape (B, selected_params)

standardise_from_batch(batch)

Standardise batch using precomputed mean/std buffers.

Parameters:

batch (torch.Tensor) – Shape (B, total_params)

Returns:

Standardised batch (B, selected_params)

Return type:

torch.Tensor

unstandardise_from_batch(standardised_batch)

Convert standardised batch back to original parameter scale.

Parameters:

standardised_batch (torch.Tensor) – Shape (B, selected_params)

Returns:

Unstandardised batch

Return type:

torch.Tensor

forward(N)

Sample a batch of N waveform parameters from the configured prior.

Applies constraints (rejection sampling) and deterministic transforms (e.g. mchirp/q, Cartesian spins, luminosity distance) in order.

Parameters:

N (int) – Number of samples to draw.

Returns:

Complete parameter batch with all transforms applied.

Return type:

torch.Tensor, shape (N, num_params)