amplfi.train module

class amplfi.train.augmentations.WaveformProjector(ifos, sample_rate)[source]

Bases: Module

forward(dec, psi, phi, **polarizations)[source]
Return type:

Tensor

class amplfi.train.augmentations.PsdEstimator(length, sample_rate, fftlength, overlap=None, average='mean', fast=True)[source]

Bases: Module

Module that takes a sample of data, splits it into two unequal-length segments, calculates the PSD of the first section, then returns this PSD along with the second section.

Parameters:

  • length (float) – The length, in seconds, of timeseries data to be returned for whitening. Note that the length of time used for the PSD will then be whatever remains along first part of the time axis of the input.

  • sample_rate (float) – Rate at which input data has been sampled in Hz

  • fftlength (float) – Length of FFTs to use when computing the PSD

  • overlap (Optional[float]) – Amount of overlap between FFT windows when computing the PSD. Default value of None uses fftlength / 2

  • average (str) – Method for aggregating spectra from FFT windows, either “mean” or “median”

  • fast (bool) – If True, use a slightly faster PSD algorithm that is inaccurate for the lowest two frequency bins. If you plan on highpassing later, this should be fine.

forward(X)[source]
Return type:

Tuple[Tensor, Tensor]

class amplfi.train.losses.VICRegLoss(lambda_param=25.0, mu_param=25.0, nu_param=1.0, eps=0.0001, max_std=1.0)[source]

Bases: Module

Implementation of the VICReg loss [0].

This implementation is based on the code published by the authors [1].

lambda_param

Scaling coefficient for the invariance term of the loss.

mu_param

Scaling coefficient for the variance term of the loss.

nu_param

Scaling coefficient for the covariance term of the loss.

eps

Epsilon for numerical stability.

forward(z_a, z_b)[source]

Returns VICReg loss.

Parameters:

  • z_a (Tensor) – Tensor with shape (batch_size, …, dim).

  • z_b (Tensor) – Tensor with shape (batch_size, …, dim).

Return type:

Tensor

Returns:

The computed VICReg loss.

amplfi.train.losses.invariance_loss(x, y)[source]

Returns VICReg invariance loss.

Parameters:

  • x (Tensor) – Tensor with shape (batch_size, …, dim).

  • y (Tensor) – Tensor with shape (batch_size, …, dim).

Return type:

Tensor

Returns:

The computed VICReg invariance loss.

amplfi.train.losses.variance_loss(x, eps=0.0001, max_std=1.0)[source]

Returns VICReg variance loss.

Parameters:

  • x (Tensor) – Tensor with shape (batch_size, …, dim).

  • eps (float) – Epsilon for numerical stability.

Return type:

Tensor

Returns:

The computed VICReg variance loss.

amplfi.train.losses.covariance_loss(x)[source]

Returns VICReg covariance loss.

Generalized version of the covariance loss with support for tensors with more than two dimensions. Adapted from VICRegL: https://github.com/facebookresearch/VICRegL/blob/803ae4c8cd1649a820f03afb4793763e95317620/main_vicregl.py#L299

Parameters:

x (Tensor) – Tensor with shape (batch_size, …, dim).

Return type:

Tensor

Returns:

The computed VICReg covariance loss.

class amplfi.train.prior.AmplfiPrior(priors, conversion_function=None)[source]

Bases: object

log_prob(samples)[source]

Calculate the log probability of samples under the prior

Parameters:

samples (dict[str, Tensor]) – Dictionary where key is parameter and value is tensor of samples

Return type:

Tensor

class amplfi.train.prior.ParameterTransformer(**transforms)[source]

Bases: Module

Helper class for applying preprocessing transformations to inference parameters

forward(parameters)[source]