diffbank.utils#

Utilities for calculating properties of binaries, sampling their parameters, comparing waveforms and generating template banks.

Warning

Much of this module will be refactored into its own package in the future.

diffbank.utils.est_ratio_max(key, density_fun, sample_base, density_fun_base=<function <lambda>>, n_iter=1000, n_init=200, show_progress=True)[source]#

Estimate maximum of the ratio of target to base density using empirical supremum rejection sampling.

Note

Not jit-able.

References

https://bookdown.org/rdpeng/advstatcomp/rejection-sampling.html

Parameters

  • density_fun (Callable[[ndarray], ndarray]) – jit-able target density function

  • sample_base (Callable[[PRNGKeyArray, int], ndarray]) – jit-able base sampler taking a key and number of samples as arguments

  • density_fun_base (Callable[[ndarray], ndarray]) – jit-able density function for sample_base. Need not be normalized.

  • n_iter (int) – number of iterations of rejection sampling to perform to estimate the ratio of densities

  • n_init (int) – as an initial guess, the maximum ratio will be computed over by sampling n_init points from sample_base

  • show_progress (bool) – displays a tqdm progress bar if True

Return type

Tuple[ndarray, ndarray]

Returns

The estimated maximum ratio between density_fun and density_fun_base, and point at which it was attained. This can be passed as ratio_max to other sampling functions.

diffbank.utils.gen_bank_random(key, minimum_match, eta, match_fun, ratio_max, density_fun, sample_base, density_fun_base=<function <lambda>>, eff_pt_sampler=None, n_eff=1000, show_progress=True)[source]#

Generates a random bank using the method introduced in Coogan et al 2022.

References

See also Messenger, Prix & Papa 2008

Note

Not jit-able

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • minimum_match (Union[float, ndarray]) – the target minimum match match of the bank

  • eta (Union[float, ndarray]) – the target fraction of parameters space to cover with the bank

  • match_fun (Callable[[ndarray, ndarray], ndarray]) – jit-able function to compute the match between two points

  • ratio_max (ndarray) – maximum value of the ratio of density_fun to density_fun_base

  • density_fun (Callable[[ndarray], ndarray]) – jit-able target density function

  • sample_base (Callable[[PRNGKeyArray, int], ndarray]) – jit-able base sampler taking a key and number of samples as arguments

  • density_fun_base (Callable[[ndarray], ndarray]) – jit-able density function for sample_base. Need not be normalized.

  • eff_pt_sampler (Optional[Callable[[PRNGKeyArray], ndarray]]) – if provided, this function will be used to sample effectualness points instead of sampling them with the same density as templates

  • n_eff (int) – the number of effectualness points used to monitor convergence

  • show_progress (bool) – displays a tqdm progress bar if True

Return type

Tuple[ndarray, ndarray]

Returns

The array of template positions and effectualness points used to monitor convergence

diffbank.utils.gen_bank_stochastic(key, minimum_match, eta, match_fun, propose_template, eff_pt_sampler, n_eff=1000, show_progress=True, n_acc_monitoring=1)[source]#

Generates a stochastic bank by adding an accept/reject step to the random bank generation method introduced in Coogan et al 2022. This step rejects the template if its effectualness is greater than the target minimum match with respect to the bank at the time it is sampled.

References

See this paper, for example

Note

Not jit-able

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • minimum_match (Union[float, ndarray]) – the target minimum match match of the bank

  • eta (Union[float, ndarray]) – the target fraction of parameters space to cover with the bank

  • match_fun (Callable[[ndarray, ndarray], ndarray]) – jit-able function to compute the match between two points

  • propose_template (Callable[[PRNGKeyArray], ndarray]) – jit-able function for sampling new templates

  • eff_pt_sampler (Callable[[PRNGKeyArray], ndarray]) – if provided, this function will be used to sample effectualness points instead of sampling them with the same density as templates

  • n_eff (int) – the number of effectualness points used to monitor convergence

  • show_progress (bool) – displays a tqdm progress bar if True

  • n_acc_monitoring (int) – number of iterations to use for moving average calculation of the acceptance rate. Only relevant if show_progress is True.

Return type

Tuple[ndarray, ndarray]

Returns

The array of template positions and effectualness points used to monitor convergence

diffbank.utils.gen_template_rejection(key, ratio_max, density_fun, sample_base, density_fun_base=<function <lambda>>)[source]#

Generates a single template using rejection sampling.

Note

While this function is used to generate templates in diffbank, it can be used for any rejection sampling task.

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • ratio_max (Union[float, ndarray]) – maximum value of the ratio of density_fun to density_fun_base

  • density_fun (Callable[[ndarray], ndarray]) – jit-able target density function

  • sample_base (Callable[[PRNGKeyArray, int], ndarray]) – jit-able base sampler taking a key and number of samples as arguments

  • density_fun_base (Callable[[ndarray], ndarray]) – jit-able density function for sample_base. Need not be normalized.

Return type

ndarray

Returns

A single sample from the distribution with density density_fun

diffbank.utils.gen_templates_rejection(key, n_templates, ratio_max, density_fun, sample_base, density_fun_base=<function <lambda>>)[source]#

Generates multiple templates using rejection sampling.

Note

While this function is used to generate templates in diffbank, it can be used for any rejection sampling task.

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • n_templates (int) – the number of templates to sample

  • ratio_max (Union[float, ndarray]) – maximum value of the ratio of density_fun to density_fun_base

  • density_fun (Callable[[ndarray], ndarray]) – jit-able target density function

  • sample_base (Callable[[PRNGKeyArray, int], ndarray]) – jit-able base sampler taking a key and number of samples as arguments

  • density_fun_base (Callable[[ndarray], ndarray]) – jit-able density function for sample_base. Need not be normalized.

Return type

ndarray

Returns

n_templates samples from the distribution with density density_fun

diffbank.utils.get_M_eta_sampler(M_range, eta_range)[source]#

Uniformly values of the chirp mass and samples over the specified ranges. This function may be removed in the future since it is trivial.

Return type

Callable[[PRNGKeyArray, int], ndarray]

diffbank.utils.get_bank_effectualness(key, minimum_match, templates, match_fun, eff_pt_sampler, n=100, show_progress=True)[source]#

Computes effectualness of a bank at random points.

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • minimum_match (Union[float, ndarray]) – the target minimum match match of the bank

  • templates (ndarray) – the bank’s templates’ locations

  • match_fun (Callable[[ndarray, ndarray], ndarray]) – jit-able function to compute the match between two points

  • eff_pt_sampler (Callable[[PRNGKeyArray], ndarray]) – function for sampling points at which to compute the bank’s effectualness

  • n (int) – the number of points at which to compute the bank’s effectualness

  • show_progress (bool) – displays a tqdm progress bar if True

Return type

Tuple[ndarray, ndarray, ndarray, ndarray]

Returns

Effectualness, the corresponding points, and the Monte Carlo estimate and error for the fraction of parameter space covered by the bank

diffbank.utils.get_eff_pads(fs)[source]#

Gets arrays of zeros to pad a function evaluated on a frequency grid so the function values can be passed to jax.numpy.fft.ifft.

Parameters

fs (ndarray) – uniformly-spaced grid of frequencies. It is assumed that the first element in the grid must be an integer multiple of the grid spacing (i.e., fs[0] % df == 0, where df is the grid spacing).

Return type

Tuple[ndarray, ndarray]

Returns

The padding arrays of zeros. The first is of length fs[0] / df and the second is of length fs[-1] / df - 2.

diffbank.utils.get_f_isco(m)[source]#

Computes the ISCO frequency for a black hole.

Parameters

m – the black hole’s mass in kg

Returns

The ISCO frequency in Hz

diffbank.utils.get_m1_m2_sampler(m1_range, m2_range)[source]#

Creates a function to uniformly sample two parameters, with the restriction that the first is larger than the second.

Note

While this function is particularly useful for sampling masses in a binary, nothing in it is specific to that context.

Parameters

  • m1_range (Tuple[float, float]) – the minimum and maximum values of the first parameter

  • m2_range (Tuple[float, float]) – the minimum and maximum values of the second parameter

Return type

Callable[[PRNGKeyArray, int], ndarray]

Returns

The sampling function

diffbank.utils.get_match(theta1, theta2, amp1, Psi1, amp2, Psi2, fs, Sn, pad_low, pad_high)[source]#

Calculates the match between two waveforms with different parameters and of distinct types. The match is defined as the noise-weighted inner product maximized over the difference in time and phase at coalescence. The maximizations are performed using the absolute value of the inverse Fourier transform trick.

Parameters

  • theta1 (ndarray) – parameters for the first waveform

  • theta2 (ndarray) – parameters for the second waveform

  • amp1 (Callable[[ndarray, ndarray], ndarray]) – amplitude function for first waveform

  • Psi1 (Callable[[ndarray, ndarray], ndarray]) – phase function for first waveform

  • amp2 (Callable[[ndarray, ndarray], ndarray]) – amplitude function for second waveform

  • Psi2 (Callable[[ndarray, ndarray], ndarray]) – phase function for second waveform

  • fs (ndarray) – uniformly-spaced grid of frequencies used to perform the integration

  • Sn (Callable[[ndarray], ndarray]) – power spectral density of the detector noise

  • pad_low (ndarray) – array of zeros to pad the left side of the integrand before it is passed to jax.numpy.fft.ifft

  • pad_right – array of zeros to pad the right side of the integrand before it is passed to jax.numpy.fft.ifft

Return type

ndarray

Returns

The match \(m[\theta_1, \theta_2]\)

diffbank.utils.get_match_arr(h1, h2, Sns, fs, pad_low, pad_high)[source]#

Calculates the match between two frequency-domain complex strains. The maximizations over the difference in time and phase at coalescence are performed by taking the absolute value of the inverse Fourier transform.

Parameters

  • h1 (ndarray) – the first set of strains

  • h2 (ndarray) – the second set of strains

  • Sns (ndarray) – the noise power spectral densities

  • fs (ndarray) – frequencies at which the strains and noise PSDs were evaluated

  • pad_low (ndarray) – array of zeros to pad the left side of the integrand before it is passed to jax.numpy.fft.ifft

  • pad_right – array of zeros to pad the right side of the integrand before it is passed to jax.numpy.fft.ifft

Return type

ndarray

Returns

The match.

diffbank.utils.get_phase_maximized_inner_product(theta1, theta2, del_t, amp1, Psi1, amp2, Psi2, fs, Sn)[source]#

Calculates the inner product between two waveforms, maximized over the difference in phase at coalescence. This is just the absolute value of the noise-weighted inner product.

Parameters

  • theta1 (ndarray) – parameters for the first waveform

  • theta2 (ndarray) – parameters for the second waveform

  • del_t (ndarray) – difference in the time at coalescence for the waveforms

  • amp1 (Callable[[ndarray, ndarray], ndarray]) – amplitude function for first waveform

  • Psi1 (Callable[[ndarray, ndarray], ndarray]) – phase function for first waveform

  • amp2 (Callable[[ndarray, ndarray], ndarray]) – amplitude function for second waveform

  • Psi2 (Callable[[ndarray, ndarray], ndarray]) – phase function for second waveform

  • fs (ndarray) – uniformly-spaced grid of frequencies used to perform the integration

  • Sn (Callable[[ndarray], ndarray]) – power spectral density of the detector noise

Return type

ndarray

Returns

The noise-weighted inner product between the waveforms, maximized over the phase at coalescence

diffbank.utils.get_template_frac_in_bounds(key, theta, get_g, m_star, is_in_bounds, n)[source]#

Perform a Monte Carlo estimate of the fraction of a template’s metric ellipse lying inside the parameter space.

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • theta (ndarray) – the template’s location

  • get_g (Callable[[ndarray], ndarray]) – a function to compute the metric at a point

  • m_star – the minimum match, used to scale the metric ellipse by sqrt(m_star)

  • is_in_bounds (Callable[[ndarray], ndarray]) – callable that takes a point and returns 1 if it is in the parameter space and 0 if not

  • n (int) – number of points to sample for the Monte Carlo estimate

Return type

Tuple[ndarray, ndarray]

Returns

Estimate and error for the fraction of the template ellipse at theta lying inside the parameter space

diffbank.utils.get_top_matches_templates(templates, match_fun, points, n_top=1, show_progress=True)[source]#

Find templates that have highest match with respect to a set of points.

Parameters

  • templates (ndarray) – a template bank

  • match_fun (Callable[[ndarray, ndarray], ndarray]) – jit-able function to compute the match taking the template as the first argument and point as the second argument.

  • points (ndarray) – a single point (1D array) or a set of points (2D array)

  • n_top (int) – number of best-matching templates to return for each point. If this is 1, the best template and its match (i.e., the bank’s effectualness) are returned for each point. If this is greater than 1, the n_top best-matching templates are returned for each point.

  • show_progress (bool) – displays a tqdm progress bar if True

Raises

ValueError – if n_top < 1 or if points is not 1D or 2D

Return type

Tuple[ndarray, ndarray]

Returns

The top matches and corresponding templates in the bank for each point. If n_top is 1, these will have shapes (n,) and (n, p), where n is the length of points and p is the dimension of each of its elements. If n_top is greater than 1, their returned arrays will have shapes (n, n_top) and (n, n_top, p).

diffbank.utils.ms_to_Mc_eta(m)[source]#

Converts binary component masses to chirp mass and symmetric mass ratio.

Parameters

m – the binary component masses (m1, m2)

Returns

\((\mathcal{M}, \eta)\), with the chirp mass in the same units as the component masses

diffbank.utils.sample_uniform_ball(key, dim, shape=(1,))[source]#

Uniformly sample from the unit ball.

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • dim (int) – dimensionality of the ball

  • shape (Tuple[int]) – shape of samples to draw

Return type

ndarray

Returns

Samples from the dim-dimensional unit ball of shape shape

diffbank.utils.sample_uniform_metric_ellipse(key, g, n)[source]#

Uniformly sample inside a metric ellipse centered at the origin.

Parameters

  • key (PRNGKeyArray) – key to pass to sampler

  • g (ndarray) – the metric

  • n (int) – the number of samples to draw

Return type

ndarray

Returns

n samples from the metric ellipse centered at the origin