ml4gw package

Subpackages

Submodules

ml4gw.augmentations module

class ml4gw.augmentations.SignalInverter(prob=0.5)

Bases: Module

Takes a tensor of timeseries of arbitrary dimension and randomly inverts i.e. \(h(t) \rightarrow -h(t)\) each timeseries with probability prob.

Parameters:

prob (float) -- Probability that a timeseries is inverted

forward(X)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Float[Tensor, '*batch time']

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml4gw.augmentations.SignalReverser(prob=0.5)

Bases: Module

Takes a tensor of timeseries of arbitrary dimension and randomly reverses i.e., \(h(t) \rightarrow h(-t)\). each timeseries with probability prob.

Parameters:

prob (float) -- Probability that a kernel is reversed

forward(X)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Float[Tensor, '*batch time']

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

ml4gw.distributions module

Module containing callables classes for generating samples from specified distributions. Each callable should map from an integer N to a 1D torch Tensor containing N samples from the corresponding distribution.

class ml4gw.distributions.Cosine(low=-1.5707963267948966, high=1.5707963267948966, validate_args=None)

Bases: Distribution

Cosine distribution based on torch.distributions.TransformedDistribution (see documentation).

arg_constraints = {}
log_prob(value)

Returns the log of the probability density/mass function evaluated at value.

Parameters:

value (Tensor)

Return type:

Float[Tensor, '']

rsample(sample_shape=None)

Generates a sample_shape shaped reparameterized sample or sample_shape shaped batch of reparameterized samples if the distribution parameters are batched.

Return type:

Tensor

class ml4gw.distributions.DeltaFunction(peak=0.0, validate_args=None)

Bases: Distribution

arg_constraints = {}
rsample(sample_shape=None)

Generates a sample_shape shaped reparameterized sample or sample_shape shaped batch of reparameterized samples if the distribution parameters are batched.

Return type:

Tensor

class ml4gw.distributions.LogNormal(mean, std, low=None, validate_args=None)

Bases: LogNormal

support()

Constrain to a real half line (lower_bound, inf].

class ml4gw.distributions.LogUniform(low, high, validate_args=None)

Bases: TransformedDistribution

Sample from a log uniform distribution

class ml4gw.distributions.PowerLaw(minimum, maximum, index, validate_args=None)

Bases: TransformedDistribution

Sample from a power law distribution,

\[p(x) \approx x^{\alpha}.\]

Index alpha cannot be 0, since it is equivalent to a Uniform distribution. This could be used, for example, as a universal distribution of signal-to-noise ratios (SNRs) from uniformly volume distributed sources

\[p(\rho) = 3\;\rho_0^3 / \rho^4\]

where \(\rho_0\) is a representative minimum SNR considered for detection. See, for example, Schutz (2011). Or, for example, index=2 for uniform in Euclidean volume.

support = GreaterThanEq(lower_bound=0.0)
class ml4gw.distributions.RateEvolution(rate_function, *args, **kwargs)

Bases: UniformComovingVolume

Wrapper around UniformComovingVolume() to allow for arbitrary rate evolution functions. E.g., if rate_function = lambda z: 1 / (1 + z), then the distribution will sample values such that they occur uniform in source frame time.

Parameters:

  • rate_function (Callable) -- Callable that takes redshift as input and returns the rate evolution factor.

  • *args -- Arguments passed to UniformComovingVolume() constructor.

  • **kwargs -- Keyword arguments passed to UniformComovingVolume() constructor.

class ml4gw.distributions.Sine(low=0.0, high=3.141592653589793, validate_args=None)

Bases: TransformedDistribution

Sine distribution based on torch.distributions.TransformedDistribution.

class ml4gw.distributions.UniformComovingVolume(minimum, maximum, distance_type='redshift', h0=67.66, omega_m=0.30966, z_grid_max=5, grid_size=10000, validate_args=None)

Bases: Distribution

Sample either redshift, comoving distance, or luminosity distance such that they are uniform in comoving volume, assuming a flat lambda-CDM cosmology. Default H0 and Omega_M values match Planck18 parameters in Astropy.

Parameters:

  • minimum (float) -- Minimum distance in the specified distance type

  • maximum (float) -- Maximum distance in the specified distance type

  • distance_type (str) -- Type of distance to sample from. Can be redshift, comoving_distance, or luminosity_distance

  • h0 (float) -- Hubble constant in km/s/Mpc

  • omega_m (float) -- Matter density parameter

  • z_max -- Maximum redshift for the grid

  • grid_size (int) -- Number of points in the grid for interpolation

  • validate_args (Optional[bool]) -- Whether to validate arguments

arg_constraints = {}
log_prob(value)

Returns the log of the probability density/mass function evaluated at value.

Parameters:

value (Tensor)

Return type:

Tensor

rsample(sample_shape=None)

Generates a sample_shape shaped reparameterized sample or sample_shape shaped batch of reparameterized samples if the distribution parameters are batched.

Return type:

Tensor

support = GreaterThanEq(lower_bound=0.0)

ml4gw.gw module

Tools for manipulating raw gravitational waveforms and projecting them onto interferometer responses. Much of the projection code is an extension of the implementation made available in bilby. Specifically code from this module.

ml4gw.gw.breathing(m, n)
Return type:

Float[Tensor, 'batch space space']

ml4gw.gw.compute_antenna_responses(theta, psi, phi, detector_tensors, modes)

Compute the antenna pattern factors of a batch of waveforms as a function of the sky parameters of their sources as well as the detector tensors of the interferometers whose response is being calculated.

Parameters:

  • theta (Float[Tensor, 'batch']) -- Angle of each source in radians relative to the celestial equator

  • psi (Float[Tensor, 'batch']) -- Angle in radians between each source's natural polarization basis and the basis which has the 0th unit vector pointing along the celestial equator

  • phi (Float[Tensor, 'batch']) -- Angle in radians between each source's right ascension and the right ascension of the geocenter

  • detector_tensors (Float[Tensor, 'num_ifos 3 3']) -- Detector tensor for each of the interferometers for which a response is being calculated, stacked along the 0th axis

  • modes (List[str]) -- Which polarization modes to compute the response for

Return type:

Float[Tensor, 'batch polarizations num_ifos']

Returns:

A tensor representing interferometer antenna pattern

factors for each of the polarizations of each of the waveforms, for each interferometer.

ml4gw.gw.compute_ifo_snr(responses, psd, sample_rate, highpass=None, lowpass=None)

Compute the SNRs of a batch of interferometer responses

Compute the signal to noise ratio (SNR) of individual interferometer responses to gravitational waveforms with respect to a background PSD for each interferometer. The SNR of the \(i\) th waveform at the \(j\) th interferometer is computed as:

\[\rho_{ij} = 4 \int_{f_{\text{min}}}^{f_{\text{max}}} \frac{\tilde{h_{ij}}(f)\tilde{h_{ij}}^*(f)} {S_n^{(j)}(f)}df\]

Where \(f_{\text{min}}\) is a minimum frequency denoted by highpass, \(f_{\text{max}}\) is the maximum frequency denoted by lowpass, which defaults to the Nyquist frequency dictated by sample_rate; \(\tilde{h}_{ij}\) and \(\tilde{h}_{ij}^*\) indicate the fourier transform of the \(i\) th waveform at the \(j\) th inteferometer and its complex conjugate, respectively; and \(S_n^{(j)}\) is the backround PSD at the \(j\) th interferometer.

Parameters:

  • responses (Float[Tensor, 'batch num_ifos time']) -- A batch of interferometer responses to a batch of raw gravitational waveforms

  • psd (Float[Tensor, 'num_ifos frequency']) -- The one-sided power spectral density of the background noise at each interferometer to which a response in responses has been calculated. If 2D, each row of psd will be assumed to be the background PSD for each channel of _every_ batch element in responses. If 3D, this should contain a background PSD for each channel of each element in responses, and therefore the first two dimensions of psd and responses should match.

  • sample_rate (float) -- The frequency at which the waveform responses timeseries have been sampled. Upon fourier transforming, should match the frequency resolution of the provided PSDs.

  • highpass (Union[float, Float[Tensor, 'frequency'], None]) -- The minimum frequency above which to compute the SNR. If a tensor is provided, it will be assumed to be a pre-computed mask used to 0-out low frequency components. If a float, it will be used to compute such a mask. If left as None, all frequencies up to lowpass will contribute to the SNR calculation.

  • lowpass (Union[float, Float[Tensor, 'frequency'], None]) -- The maximum frequency below which to compute the SNR. If a tensor is provided, it will be assumed to be a pre-computed mask used to 0-out high frequency components. If a float, it will be used to compute such a mask. If left as None, all frequencies from highpass up to the Nyquist freqyency will contribute to the SNR calculation.

Return type:

Float[Tensor, 'batch num_ifos']

Returns:

Batch of SNRs computed for each interferometer

ml4gw.gw.compute_network_snr(responses, psd, sample_rate, highpass=None, lowpass=None)

Compute the total SNR from a gravitational waveform from a network of interferometers. The total SNR for the \(i\) th waveform is computed as

\[\rho_i = \sqrt{\sum_{j}^{N}\rho_{ij}^2}\]

where \(\rho_{ij}\) is the SNR for the \(i\) th waveform at the \(j\) th interferometer in the network and \(N\) is the total number of interferometers.

Parameters:

  • responses (Float[Tensor, 'batch num_ifos time']) -- A batch of interferometer responses to a batch of raw gravitational waveforms

  • backgrounds -- The one-sided power spectral density of the background noise at each interferometer to which a response in responses has been calculated. If 2D, each row of psd will be assumed to be the background PSD for each channel of every batch element in responses. If 3D, this should contain a background PSD for each channel of each element in responses, and therefore the first two dimensions of psd and responses should match.

  • sample_rate (float) -- The frequency at which the waveform responses timeseries have been sampled. Upon fourier transforming, should match the frequency resolution of the provided PSDs.

  • highpass (Union[float, Float[Tensor, 'frequency'], None]) -- The minimum frequency above which to compute the SNR. If a tensor is provided, it will be assumed to be a pre-computed mask used to 0-out low frequency components. If a float, it will be used to compute such a mask. If left as None, all frequencies up to sample_rate / 2 will contribute to the SNR calculation.

  • lowpass (Union[float, Float[Tensor, 'frequency'], None]) -- The maximum frequency below which to compute the SNR. If a tensor is provided, it will be assumed to be a pre-computed mask used to 0-out high frequency components. If a float, it will be used to compute such a mask. If left as None, all frequencies from highpass up to the Nyquist freqyency will contribute to the SNR calculation.

Return type:

Float[Tensor, 'batch']

Returns:

Batch of SNRs for each waveform across the interferometer network

ml4gw.gw.compute_observed_strain(dec, psi, phi, detector_tensors, detector_vertices, sample_rate, **polarizations)

Compute the strain timeseries $h(t)$ observed by a network of interferometers from the given polarization timeseries corresponding to gravitational waveforms from sources with the indicated sky parameters.

Parameters:

  • dec (Float[Tensor, 'batch']) -- Declination of each source in radians relative to the celestial north

  • psi (Float[Tensor, 'batch']) -- Angle in radians between each source's natural polarization basis and the basis which has the 0th unit vector pointing along the celestial equator

  • phi (Float[Tensor, 'batch']) -- Angle in radians between each source's right ascension and the right ascension of the geocenter

  • detector_tensors (Float[Tensor, 'num_ifos 3 3']) -- Detector tensor for each of the interferometers for which observed strain is being calculated, stacked along the 0th axis

  • detector_vertices (Float[Tensor, 'num_ifos 3']) -- Vertices for each interferometer's spatial location relative to the geocenter. Used to compute delay between the waveform observed at the geocenter and the one observed at the detector site. To avoid adding any delay between the two, reset your coordinates such that the desired interferometer is at (0., 0., 0.).

  • sample_rate (float) -- Rate at which the polarization timeseries have been sampled

  • polarziations -- Timeseries for each waveform polarization which contributes to the interferometer response. Allowed polarizations are cross, plus, and breathing.

Return type:

Float[Tensor, 'batch num_ifos time']

Returns:

Tensor representing the observed strain at each interferometer for each waveform.

ml4gw.gw.cross(m, n)
Return type:

Float[Tensor, 'batch space space']

ml4gw.gw.get_ifo_geometry(*ifos)

For a given list of interferometer names, retrieve and concatenate the associated detector tensors and vertices of those interferometers.

Parameters:

ifos (str) -- Names of the interferometers whose geometry to retrieve

Return type:

Tuple[Float[Tensor, 'num_ifos 3 3'], Float[Tensor, 'num_ifos 3']]

Returns:

A concatenation of the detector tensors of each interferometer A concatenation of the vertices of each interferometer

ml4gw.gw.outer(x, y)

Compute the outer product of two batches of vectors

Return type:

Float[Tensor, 'batch space space']

ml4gw.gw.plus(m, n)
Return type:

Float[Tensor, 'batch space space']

ml4gw.gw.reweight_snrs(responses, target_snrs, psd, sample_rate, highpass=None, lowpass=None)

Scale interferometer responses such that they have a desired SNR

Parameters:

  • responses (Float[Tensor, 'batch num_ifos time']) -- A batch of interferometer responses to a batch of raw gravitational waveforms

  • target_snrs (Union[float, Float[Tensor, 'batch']]) -- Either a tensor of desired SNRs for each waveform, or a single SNR to which all waveforms should be scaled.

  • psd (Float[Tensor, 'num_ifos frequency']) -- The one-sided power spectral density of the background noise at each interferometer to which a response in responses has been calculated. If 2D, each row of psd will be assumed to be the background PSD for each channel of every batch element in responses. If 3D, this should contain a background PSD for each channel of each element in responses, and therefore the first two dimensions of psd and responses should match.

  • sample_rate (float) -- The frequency at which the waveform responses timeseries have been sampled. Upon fourier transforming, should match the frequency resolution of the provided PSDs.

  • highpass (Union[float, Float[Tensor, 'frequency'], None]) -- The minimum frequency above which to compute the SNR. If a tensor is provided, it will be assumed to be a pre-computed mask used to 0-out low frequency components. If a float, it will be used to compute such a mask. If left as None, all frequencies up to sample_rate / 2 will contribute to the SNR calculation.

  • lowpass (Union[float, Float[Tensor, 'frequency'], None]) -- The maximum frequency below which to compute the SNR. If a tensor is provided, it will be assumed to be a pre-computed mask used to 0-out high frequency components. If a float, it will be used to compute such a mask. If left as None, all frequencies from highpass up to the Nyquist freqyency will contribute to the SNR calculation.

Return type:

Float[Tensor, 'batch num_ifos time']

Returns:

Rescaled interferometer responses

ml4gw.gw.shift_responses(responses, theta, phi, vertices, sample_rate)
Return type:

Float[Tensor, 'batch num_ifos time']

ml4gw.spectral module

Several implementation details are derived from the scipy csd and welch implementations. For more info, see

https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.welch.html

and

https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.csd.html

ml4gw.spectral.fast_spectral_density(x, nperseg, nstride, window, scale, average='median', y=None)

Compute the power spectral density of a multichannel timeseries or a batch of multichannel timeseries, or the cross power spectral density of two such timeseries. This implementation is non-exact for the two lowest frequency bins, since it implements centering using the mean for the entire timeseries rather than on a per- window basis. The benefit of this is a faster implementation, which might be beneficial for cases where the lowest frequency components are going to be discarded anyway.

Parameters:

  • x (Union[Float[Tensor, 'time'], Float[Tensor, 'channel time'], Float[Tensor, 'batch channel time']]) -- The timeseries tensor whose power spectral density to compute, or for cross spectral density the timeseries whose fft will be conjugated. Can have shape (batch_size, num_channels, length * sample_rate), (num_channels, length * sample_rate), or (length * sample_rate).

  • nperseg (int) -- Number of samples included in each FFT window

  • nstride (int) -- Stride between FFT windows

  • window (Float[Tensor, '{nperseg//2+1}']) -- Window array to multiply by each FFT window before FFT computation. Should have length nperseg // 2 + 1.

  • scale (float) -- Scale factor to multiply the FFT'd data by, related to desired units for output tensor (e.g. letting this equal 1 / (sample_rate * (window**2).sum()) will give output units of density, :math`text{Hz}^-1`.

  • average (str) -- How to aggregate the contributions of each FFT window to the spectral density. Allowed options are 'mean' and 'median'.

  • y (Union[Float[Tensor, 'time'], Float[Tensor, 'channel time'], Float[Tensor, 'batch channel time'], None]) -- Timeseries tensor to compute cross spectral density with x. If left as None, x's power spectral density will be returned. Otherwise, if x is 1D, y must also be 1D. If x is 2D, the assumption is that this represents a single multi-channel timeseries, and y must be either 2D or 1D. In the former case, the cross-spectral densities of each channel will be computed individually, so y must have the same shape as x. Otherwise, this will compute the CSD of each of x's channels with y. If x is 3D, this will be assumed to be a batch of multi-channel timeseries. In this case, y can either be 3D, in which case each channel of each batch element will have its CSD calculated or 2D, which has two different options. If y's 0th dimension matches x's 0th dimension, it will be assumed that y represents a batch of 1D timeseries, and for each batch element this timeseries will have its CSD with each channel of the corresponding batch element of x calculated. Otherwise, it sill be assumed that y represents a single multi-channel timeseries, in which case each channel of y will have its CSD calculated with the corresponding channel in x across _all_ of x's batch elements.

Return type:

Union[Float[Tensor, 'frequency'], Float[Tensor, 'channel frequency'], Float[Tensor, 'batch channel frequency']]

Returns:

Tensor of power spectral densities of x or its cross spectral density with the timeseries in y.

ml4gw.spectral.median(x, axis)

Implements a median calculation that matches numpy's behavior for an even number of elements and includes the same bias correction used by scipy's implementation.

Return type:

']

ml4gw.spectral.normalize_by_psd(X, psd, sample_rate, pad)
ml4gw.spectral.spectral_density(x, nperseg, nstride, window, scale, average='median')

Compute the power spectral density of a multichannel timeseries or a batch of multichannel timeseries. This implementation is exact for all frequency bins, but slower than the fast implementation.

Parameters:

  • x (Union[Float[Tensor, 'time'], Float[Tensor, 'channel time'], Float[Tensor, 'batch channel time']]) -- The timeseries tensor whose power spectral density to compute, or for cross spectral density the timeseries whose fft will be conjugated. Can have shape (batch_size, num_channels, length * sample_rate), (num_channels, length * sample_rate), or (length * sample_rate).

  • nperseg (int) -- Number of samples included in each FFT window

  • nstride (int) -- Stride between FFT windows

  • window (Float[Tensor, '{nperseg//2+1}']) -- Window array to multiply by each FFT window before FFT computation. Should have length nperseg // 2 + 1.

  • scale (float) -- Scale factor to multiply the FFT'd data by, related to desired units for output tensor (e.g. letting this equal 1 / (sample_rate * (window**2).sum()) will give output units of density, \(\text{Hz}^-1\).

  • average (str) -- How to aggregate the contributions of each FFT window to the spectral density. Allowed options are 'mean' and 'median'.

Return type:

Union[Float[Tensor, 'frequency'], Float[Tensor, 'channel frequency'], Float[Tensor, 'batch channel frequency']]

ml4gw.spectral.truncate_inverse_power_spectrum(psd, fduration, sample_rate, highpass=None, lowpass=None)

Truncate the length of the time domain response of a whitening filter built using the specified psd so that it has maximum length fduration seconds. This is meant to mitigate the impact of sharp features in the background PSD causing time domain responses longer than the segments to which the whitening filter will be applied.

Implementation details adapted from here.

Parameters:

  • psd (Float[Tensor, 'num_ifos frequency']) -- The one-sided power spectraul density used to construct a whitening filter.

  • fduration (Union[Float[Tensor, 'time'], float]) -- Desired length in seconds of the time domain response of a whitening filter built using this PSD, or a window of this length to taper the edges of the time domain response of the filter. If passed as a float, a Hann window of this length will be used.

  • sample_rate (float) -- Rate at which the time domain data to which the whitening filter will be applied has been sampled.

  • highpass (Optional[float]) -- If specified, will zero out the frequency response of all frequencies below this value in Hz. If left as None, no highpass filtering will be applied.

  • lowpass (Optional[float]) -- If specified, will zero out the frequency response of all frequencies above this value in Hz. If left as None, no lowpass filtering will be applied.

Return type:

Float[Tensor, 'num_ifos frequency']

Returns:

The PSD with its time domain response truncated

to fduration and any filtered frequencies tapered.

ml4gw.spectral.whiten(X, psd, fduration, sample_rate, highpass=None, lowpass=None)

Whiten a batch of timeseries using the specified background one-sided power spectral densities (PSDs), modified to have the desired time domain response length fduration and possibly to highpass/lowpass filter.

Parameters:

  • X (Float[Tensor, 'batch num_ifos time']) -- batch of multichannel timeseries to whiten

  • psd (Float[Tensor, 'num_ifos frequency']) -- PSDs use to whiten the data. The frequency response of the whitening filter will be roughly the inverse of the square root of this PSD, ensuring that data from the same distribution will have approximately uniform power after whitening. If 2D, each batch element in X will be whitened using the same PSDs. If 3D, each batch element will be whitened by the PSDs contained along the 0th dimenion of psd, and so the first two dimensions of X and psd should match.

  • fduration (Union[Float[Tensor, 'time'], float]) -- Desired length in seconds of the time domain response of a whitening filter built using this PSD, or a window of this length to taper the edges of the time domain response of the filter. If passed as a float, a Hann window of this length will be used. Moreover, half of this length will be removed from each edge of the whitened timeseries to account for filter settle-in time.

  • sample_rate (float) -- Rate at which the data in X has been sampled

  • highpass (Optional[float]) -- The frequency in Hz at which to highpass filter the data, setting the frequency response in the whitening filter to 0. If left as None, no highpass filtering will be applied.

  • lowpass (Optional[float]) -- The frequency in Hz at which to lowpass filter the data, setting the frequency response in the whitening filter to 0. If left as None, no lowpass filtering will be applied.

Return type:

Float[Tensor, 'batch num_ifos time']

Returns:

Batch of whitened multichannel timeseries with

fduration / 2 seconds trimmed from each side.

ml4gw.types module

Module contents