gwkokab.models.mass

Classes

BrokenPowerlaw	Base class for probability distributions in NumPyro. The design largely
BrokenPowerlawTwoPeak	Broken Powerlaw + 2 Peak is defined as a mixture of one Broken Powerlaw and two
GaussianPrimaryMassRatio	Base class for probability distributions in NumPyro. The design largely
GenericSmoothedPowerlawMassRatio	Base class for probability distributions in NumPyro. The design largely
PowerlawPrimaryMassRatio	Power law model for two-dimensional mass distribution, modelling primary mass and
SmoothedBrokenPowerlawMassRatioPowerlaw	Base class for probability distributions in NumPyro. The design largely
SmoothedGaussianPrimaryMassRatio	Normal with smoothing kernel on the
SmoothedPowerlawPrimaryMassRatio	Base class for probability distributions in NumPyro. The design largely
SmoothedTwoComponentPrimaryMassRatio	Base class for probability distributions in NumPyro. The design largely
Wysocki2019MassModel	It is a double side truncated power law distribution, as described in

Package Contents

class gwkokab.models.mass.BrokenPowerlaw(alpha1: jaxtyping.ArrayLike, alpha2: jaxtyping.ArrayLike, mbreak: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

Base class for probability distributions in NumPyro. The design largely follows from torch.distributions.

Parameters:

• batch_shape – The batch shape for the distribution. This designates independent (possibly non-identical) dimensions of a sample from the distribution. This is fixed for a distribution instance and is inferred from the shape of the distribution parameters.

• event_shape – The event shape for the distribution. This designates the dependent dimensions of a sample from the distribution. These are collapsed when we evaluate the log probability density of a batch of samples using .log_prob.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

As an example:

log_prob(value: jaxtyping.ArrayLike) → jaxtyping.ArrayLike

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.BrokenPowerlawTwoPeak(alpha1: jaxtyping.ArrayLike, alpha2: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, loc1: jaxtyping.ArrayLike, loc2: jaxtyping.ArrayLike, scale1: jaxtyping.ArrayLike, scale2: jaxtyping.ArrayLike, delta_m1: jaxtyping.ArrayLike, delta_m2: jaxtyping.ArrayLike, lambda_0: jaxtyping.ArrayLike, lambda_1: jaxtyping.ArrayLike, m1min: jaxtyping.ArrayLike, m2min: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, mbreak: jaxtyping.ArrayLike, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

Broken Powerlaw + 2 Peak is defined as a mixture of one Broken Powerlaw and two left truncated Normal distributions. For more details, see appendix B.3 of GWTC-4.0: Population Properties of Merging Compact Binaries.

It is defined as follows:

\[\begin{split}p_{\mathrm{BP}}(m_1 \mid \alpha_1, \alpha_2, m_{1,\mathrm{min}}, m_{\mathrm{max}}, m_{\mathrm{break}}) \propto \begin{cases} \left(\frac{m_1}{m_{\mathrm{break}}}\right)^{-\alpha_1} & m_{1,\mathrm{min}} \leq m_1 < m_{\text{break}} \\ \left(\frac{m_1}{m_{\mathrm{break}}}\right)^{-\alpha_2} & m_{\text{break}} \leq m_1 \leq m_{\mathrm{max}} \\ 0 & \text{otherwise} \end{cases}\end{split}\]

\[p(m_1 \mid \alpha_1, \alpha_2, m_{1,\mathrm{min}}, m_{\mathrm{max}}, m_{\mathrm{break}}, \lambda_0, \lambda_1, \mu_1, \sigma_1, \mu_2, \sigma_2) \propto \left( \lambda_0 p_{\mathrm{BP}}(m_1 \mid \alpha_1, \alpha_2, m_{1,\mathrm{min}}, m_{\mathrm{max}}, m_{\mathrm{break}}) + \lambda_1 \mathcal{N}_{[m_{1,\mathrm{min}}, \infty)}(m_1 \mid \mu_1, \sigma_1) + (1 - \lambda_0 - \lambda_1) \mathcal{N}_{[m_{1,\mathrm{min}}, \infty)}(m_1 \mid \mu_2, \sigma_2) \right) S\left(\frac{m_1 - m_{1,\mathrm{min}}}{\delta_{m_1}}\right)\]

\[p(q \mid \beta, m_1, m_{2,\mathrm{min}}, \delta_{m_2}) \propto q^{\beta} S\left(\frac{m_1 q - m_{2,\mathrm{min}}}{\delta_{m_2}}\right)\]

\[p(m_1, q \mid \alpha_1, \alpha_2, m_{1,\mathrm{min}}, m_{2,\mathrm{min}}, m_{\mathrm{max}}, m_{\mathrm{break}}, \lambda_0, \lambda_1, \mu_1, \sigma_1, \mu_2, \sigma_2, \beta, \delta_{m_1}, \delta_{m_2}) = p(m_1 \mid \alpha_1, \alpha_2, m_{1,\mathrm{min}}, m_{\mathrm{max}}, m_{\mathrm{break}}, \lambda_0, \lambda_1, \mu_1, \sigma_1, \mu_2, \sigma_2) p(q \mid \beta, m_1, m_{2,\mathrm{min}}, \delta_{m_2})\]

where \(\mathcal{N}_{[m_{1,\mathrm{min}}, \infty)}(m \mid \mu, \sigma)\) is a left truncated normal distribution with mean \(\mu\) and standard deviation \(\sigma\), \(\lambda_0\) and \(\lambda_1\) are the mixing fractions of the broken powerlaw and first Gaussian component respectively, \(\delta_{m_1}\) is the smoothing scale for the primary mass, and \(S(x)\) is the exponential of log_planck_taper_window().

log_prob(value: jaxtyping.ArrayLike) → jaxtyping.ArrayLike

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.GaussianPrimaryMassRatio(loc: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, *, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

Base class for probability distributions in NumPyro. The design largely follows from torch.distributions.

Parameters:

• batch_shape – The batch shape for the distribution. This designates independent (possibly non-identical) dimensions of a sample from the distribution. This is fixed for a distribution instance and is inferred from the shape of the distribution parameters.

• event_shape – The event shape for the distribution. This designates the dependent dimensions of a sample from the distribution. These are collapsed when we evaluate the log probability density of a batch of samples using .log_prob.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

As an example:

Parameters:

• loc (ArrayLike) – Location parameter for primary mass

• scale (ArrayLike) – Scale parameter for primary mass

• beta (ArrayLike) – Power law index for mass ratio

• mmin (ArrayLike) – Minimum mass

• mmax (ArrayLike) – Maximum mass

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.GenericSmoothedPowerlawMassRatio(primary_mass_distribution: numpyro.distributions.Distribution, beta: jaxtyping.ArrayLike, delta_m1: jaxtyping.ArrayLike, delta_m2: jaxtyping.ArrayLike, m2min: jaxtyping.ArrayLike, *, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

Base class for probability distributions in NumPyro. The design largely follows from torch.distributions.

Parameters:

• batch_shape – The batch shape for the distribution. This designates independent (possibly non-identical) dimensions of a sample from the distribution. This is fixed for a distribution instance and is inferred from the shape of the distribution parameters.

• event_shape – The event shape for the distribution. This designates the dependent dimensions of a sample from the distribution. These are collapsed when we evaluate the log probability density of a batch of samples using .log_prob.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

As an example:

log_prob(value: jaxtyping.ArrayLike) → jaxtyping.ArrayLike

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

args_constraints

primary_mass_distribution

class gwkokab.models.mass.PowerlawPrimaryMassRatio(alpha: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, *, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

Power law model for two-dimensional mass distribution, modelling primary mass and conditional mass ratio distribution.

\[p(m_1,q\mid\alpha,\beta) = p(m_1\mid\alpha)p(q \mid m_1, \beta)\]

\[\begin{split}\begin{align*} p(m_1\mid\alpha)& \propto m_1^{-\alpha},\qquad m_{\text{min}}\leq m_1\leq m_{\max}\\ p(q\mid m_1,\beta)& \propto q^{\beta},\qquad \frac{m_{\text{min}}}{m_1}\leq q\leq 1 \end{align*}\end{split}\]

Parameters:

• alpha (ArrayLike) – Power law index for primary mass

• beta (ArrayLike) – Power law index for mass ratio

• mmin (ArrayLike) – Minimum mass

• mmax (ArrayLike) – Maximum mass

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.key) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.SmoothedBrokenPowerlawMassRatioPowerlaw(alpha1: jaxtyping.ArrayLike, alpha2: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, mbreak: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, m1min: jaxtyping.ArrayLike, m2min: jaxtyping.ArrayLike, delta_m1: jaxtyping.ArrayLike, delta_m2: jaxtyping.ArrayLike, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

Base class for probability distributions in NumPyro. The design largely follows from torch.distributions.

Parameters:

• batch_shape – The batch shape for the distribution. This designates independent (possibly non-identical) dimensions of a sample from the distribution. This is fixed for a distribution instance and is inferred from the shape of the distribution parameters.

• event_shape – The event shape for the distribution. This designates the dependent dimensions of a sample from the distribution. These are collapsed when we evaluate the log probability density of a batch of samples using .log_prob.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

As an example:

log_prob(value: jaxtyping.ArrayLike) → jaxtyping.ArrayLike

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.SmoothedGaussianPrimaryMassRatio(loc, scale, beta, m1min, m2min, mmax, delta_m1, delta_m2, *, validate_args=None)

Bases: numpyro.distributions.Distribution

Normal with smoothing kernel on the lower edge.

\[p(m_1,q\mid\mu,\sigma^2,\beta,m_{\text{min}},m_{\text{max}},\delta) = \mathcal{N}(m_1\mid\mu,\sigma^2)S\left(\frac{m_1 - m_{\text{min}}}{\delta}\right)p(q \mid m_1,\beta,m_{\text{min}},\delta)\]

\[p(q\mid m_1,\beta) \propto q^{\beta}S\left(\frac{m_1q - m_{\text{min}}}{\delta}\right),\qquad \frac{m_{\text{min}}}{m_1}\leq q\leq 1\]

Logarithm of smoothing kernel is log_planck_taper_window().

Attention

If low or high are not provided to the TruncatedNormal, they default to \(-\infty\) or \(+\infty\), respectively. This class relies on this behavior to produce the desired distribution when bounds are unspecified.

log_prob(value: jaxtyping.ArrayLike) → jaxtyping.ArrayLike

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.SmoothedPowerlawPrimaryMassRatio(alpha: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, delta_m1: jaxtyping.ArrayLike, delta_m2: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, m1min: jaxtyping.ArrayLike, m2min: jaxtyping.ArrayLike, *, validate_args=None)

Bases: numpyro.distributions.Distribution

Base class for probability distributions in NumPyro. The design largely follows from torch.distributions.

Parameters:

• batch_shape – The batch shape for the distribution. This designates independent (possibly non-identical) dimensions of a sample from the distribution. This is fixed for a distribution instance and is inferred from the shape of the distribution parameters.

• event_shape – The event shape for the distribution. This designates the dependent dimensions of a sample from the distribution. These are collapsed when we evaluate the log probability density of a batch of samples using .log_prob.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

As an example:

log_prob(value: jaxtyping.ArrayLike) → jaxtyping.ArrayLike

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.SmoothedTwoComponentPrimaryMassRatio(alpha: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, delta: jaxtyping.ArrayLike, lambda_peak: jaxtyping.ArrayLike, loc: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, *, validate_args=None)

Bases: numpyro.distributions.Distribution

Base class for probability distributions in NumPyro. The design largely follows from torch.distributions.

Parameters:

• batch_shape – The batch shape for the distribution. This designates independent (possibly non-identical) dimensions of a sample from the distribution. This is fixed for a distribution instance and is inferred from the shape of the distribution parameters.

• event_shape – The event shape for the distribution. This designates the dependent dimensions of a sample from the distribution. These are collapsed when we evaluate the log probability density of a batch of samples using .log_prob.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

As an example:

log_prob(value: jaxtyping.ArrayLike) → jaxtyping.ArrayLike

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.

class gwkokab.models.mass.Wysocki2019MassModel(alpha_m: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, *, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

It is a double side truncated power law distribution, as described in equation 7 of the Reconstructing phenomenological distributions of compact binaries via gravitational wave observations.

\[p(m_1,m_2\mid\alpha,m_{\text{min}},m_{\text{max}},M_{\text{max}})\propto \frac{m_1^{-\alpha}}{m_1-m_{\text{min}}}\]

Parameters:

• alpha_m (ArrayLike) – index of the power law distribution

• mmin (ArrayLike) – lower mass limit

• mmax (ArrayLike) – upper mass limit

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=()) → jaxtyping.Array

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.key) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

support() → numpyro.distributions.constraints.Constraint

The support of this distribution. Subclasses can override this as a class attribute or as a property.