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) -> -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.
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) -> 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.
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
.- arg_constraints = {}
- log_prob(value)
Returns the log of the probability density/mass function evaluated at value.
- Parameters:
value (Tensor) --
- rsample(sample_shape=torch.Size([]))
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=torch.Size([]))
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, .. math:
p(x) pprox x^{lpha}.
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 .. math:
p(
ho) = 3* ho_0^3 / ho^4
where :math:`
- ho_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.Sine(low=0.0, high=3.141592653589793, validate_args=None)
Bases:
TransformedDistribution
Sine distribution based on
torch.distributions.TransformedDistribution
.
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:
https://arxiv.org/abs/1811.02042
Specifically the code here: https://github.com/lscsoft/bilby/blob/master/bilby/gw/detector/interferometer.py
- ml4gw.gw.breathing(m, n)
- Return type:
Tensor
- 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 (
Tensor
) -- Angle of each source in radians relative to the celestial equatorpsi (
Tensor
) -- Angle in radians between each source's natural polarization basis and the basis which has the 0th unit vector pointing along the celestial equatorphi (
Tensor
) -- Angle in radians between each source's right ascension and the right ascension of the geocenterdetector_tensors (
Tensor
) -- Detector tensor for each of the interferometers for which a response is being calculated, stacked along the 0th axismodes (
List
[str
]) -- Which polarization modes to compute the response for
- Return type:
Tensor
- 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)
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 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 (
Tensor
) -- A batch of interferometer responses to a batch of raw gravitational waveformspsd (
Tensor
) -- 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
,Tensor
,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.
- Return type:
Tensor
- Returns:
Batch of SNRs computed for each interferometer
- ml4gw.gw.compute_network_snr(responses, psd, sample_rate, highpass=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 (
Tensor
) -- A batch of interferometer responses to a batch of raw gravitational waveformsbackgrounds -- 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
,Tensor
,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.
- Return type:
Tensor
- 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 (
Tensor
) -- Declination of each source in radians relative to the celestial northpsi (
Tensor
) -- Angle in radians between each source's natural polarization basis and the basis which has the 0th unit vector pointing along the celestial equatorphi (
Tensor
) -- Angle in radians between each source's right ascension and the right ascension of the geocenterdetector_tensors (
Tensor
) -- Detector tensor for each of the interferometers for which observed strain is being calculated, stacked along the 0th axisdetector_vertices (
Tensor
) -- 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 sampledpolarziations -- Timeseries for each waveform polarization which contributes to the interferometer response. Allowed polarizations are cross, plus, and breathing.
- Return type:
Tensor
- Returns:
Tensor representing the observed strain at each interferometer for each waveform.
- ml4gw.gw.cross(m, n)
- Return type:
Tensor
- 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
[Tensor
,Tensor
]- 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:
Tensor
- ml4gw.gw.plus(m, n)
- Return type:
Tensor
- ml4gw.gw.reweight_snrs(responses, target_snrs, psd, sample_rate, highpass=None)
Scale interferometer responses such that they have a desired SNR
- Parameters:
responses (
Tensor
) -- A batch of interferometer responses to a batch of raw gravitational waveformstarget_snrs (
Union
[float
,Tensor
]) -- Either a tensor of desired SNRs for each waveform, or a single SNR to which all waveforms should be scaled.psd (
Tensor
) -- 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
,Tensor
,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.
- Return type:
Tensor
- Returns:
Rescaled interferometer responses
- ml4gw.gw.shift_responses(responses, theta, phi, vertices, sample_rate)
- Return type:
Tensor
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 (
Tensor
) -- The timeseries tensor whose power spectral density to compute, or for cross spectral density the timeseries whose fft will be conjugated. Can have shape either (batch_size, num_channels, length * sample_rate) or (num_channels, length * sample_rate).nperseg (
int
) -- Number of samples included in each FFT windownstride (
int
) -- Stride between FFT windowswindow (
Tensor
) -- Window array to multiply by each FFT window before FFT computation. Should have length nperseg // 2 + 1.scale (
Tensor
) -- 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'.y (
Optional
[Tensor
]) -- 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:
Tensor
- 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. see https://github.com/scipy/scipy/blob/main/scipy/signal/_spectral_py.py#L2066 # noqa
- 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, or the cross power spectral density of two such timeseries. This implementation is exact for all frequency bins, but slower than the fast implementation.
- Parameters:
x (
Tensor
) -- The timeseries tensor whose power spectral density to compute, or for cross spectral density the timeseries whose fft will be conjugated. Can have shape either (batch_size, num_channels, length * sample_rate) or (num_channels, length * sample_rate).nperseg (
int
) -- Number of samples included in each FFT windownstride (
int
) -- Stride between FFT windowswindow (
Tensor
) -- Window array to multiply by each FFT window before FFT computation. Should have length nperseg // 2 + 1.scale (
Tensor
) -- 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:
Tensor
- ml4gw.spectral.truncate_inverse_power_spectrum(psd, fduration, sample_rate, highpass=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 https://github.com/vivinousi/gw-detection-deep-learning/blob/203966cc2ee47c32c292be000fb009a16824b7d9/modules/whiten.py#L8 # noqa
- Parameters:
psd (
Tensor
) -- The one-sided power spectraul density used to construct a whitening filter.fduration (
Union
[Tensor
,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.
- Return type:
Tensor
- Returns:
- The PSD with its time domain response truncated
to fduration and any highpassed frequencies tapered.
- ml4gw.spectral.whiten(X, psd, fduration, sample_rate, highpass=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 filter.
- Parameters:
X (
Tensor
) -- batch of multichannel timeseries to whitenpsd (
Tensor
) -- 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
[Tensor
,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 sampledhighpass (
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.
- Return type:
Tensor
- Returns:
- Batch of whitened multichannel timeseries with
fduration / 2 seconds trimmed from each side.