gwkokab.models.mass

Two-dimensional mass models for compact binary populations.

Each model is a joint distribution over a pair of mass coordinates – \((m_1, m_2)\) or \((m_1, q)\) – with the ordering \(m_2 \leq m_1\) and the mass bounds built into its support. The Smoothed* variants taper the density above the minimum mass with a Planck window rather than truncating it sharply, and are normalised numerically at construction time.

Classes

BrokenPowerlaw	One-dimensional broken power law over the primary mass.
BrokenPowerlawTwoPeak	Broken Powerlaw + 2 Peak is defined as a mixture of one Broken Powerlaw and two
GaussianPrimaryMassRatio	Truncated normal primary mass with a conditional mass ratio power law.
GenericSmoothedPowerlawMassRatio	Smoothed conditional mass ratio power law over an arbitrary primary mass model.
PowerlawPrimaryMassRatio	Power law model for two-dimensional mass distribution, modelling primary mass and
SmoothedBrokenPowerlawMassRatioPowerlaw	Smoothed broken power law primary mass with a conditional mass ratio power law.
SmoothedGaussianPrimaryMassRatio	Normal with smoothing kernel on the
SmoothedPowerlawPrimaryMassRatio	Smoothed power law primary mass with a conditional mass ratio power law.
SmoothedTwoComponentPrimaryMassRatio	Smoothed power-law-plus-peak primary mass with a conditional mass ratio power
Wysocki2019MassModel	Double side truncated power law mass model.

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

One-dimensional broken power law over the primary mass.

Two power laws with indices \(\alpha_1\) and \(\alpha_2\) joined continuously at \(m_{\mathrm{break}}\) and truncated to \([m_{\min}, m_{\max}]\). See _broken_powerlaw_log_prob() for the density.

Parameters:

• alpha1 (ArrayLike) – Power law index below the break.

• alpha2 (ArrayLike) – Power law index above the break.

• mbreak (ArrayLike) – Break mass \(m_{\mathrm{break}}\).

• mmax (ArrayLike) – Maximum mass.

• mmin (ArrayLike) – Minimum mass.

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Notes

Sampling is not implemented; this model evaluates densities only.

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

Log probability density at value.

Parameters:

value (ArrayLike) – Primary masses.

Returns:

The log density, of shape batch_shape.

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The interval \([m_{\min}, m_{\max}]\).

Return type:

constraints.Constraint

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().

Parameters:

• alpha1 (ArrayLike) – Power law index below the break.

• alpha2 (ArrayLike) – Power law index above the break.

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

• loc1 (ArrayLike) – Location \(\mu_1\) of the first Gaussian peak.

• loc2 (ArrayLike) – Location \(\mu_2\) of the second Gaussian peak.

• scale1 (ArrayLike) – Scale \(\sigma_1\) of the first Gaussian peak.

• scale2 (ArrayLike) – Scale \(\sigma_2\) of the second Gaussian peak.

• delta_m1 (ArrayLike) – Width of the smoothing window above \(m_{1,\mathrm{min}}\).

• delta_m2 (ArrayLike) – Width of the smoothing window above \(m_{2,\mathrm{min}}\).

• lambda_0 (ArrayLike) – Mixing fraction \(\lambda_0\) of the broken power law, in \([0, 1]\).

• lambda_1 (ArrayLike) – Mixing fraction \(\lambda_1\) of the first Gaussian, in \([0, 1]\). The second Gaussian takes the remainder, \(1 - \lambda_0 - \lambda_1\).

• m1min (ArrayLike) – Minimum primary mass.

• m2min (ArrayLike) – Minimum secondary mass.

• mmax (ArrayLike) – Maximum mass.

• mbreak (ArrayLike) – Break mass \(m_{\mathrm{break}}\).

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Notes

Sampling is not implemented; this model evaluates densities only.

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

Log probability density at value.

Parameters:

value (ArrayLike) – Values whose last axis is \((m_1, q)\).

Returns:

The normalised log density, of shape batch_shape. The conditional normalisation of \(q\) is interpolated from the grid built at construction.

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{2,\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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

Truncated normal primary mass with a conditional mass ratio power law.

\[p(m_1, q) = p(m_1\mid\mu,\sigma^2)\,p(q\mid m_1, \beta)\]

\[\begin{split}\begin{align*} p(m_1\mid\mu,\sigma^2)& \propto \mathcal{N}(m_1\mid\mu,\sigma^2),\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:

• 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.

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Notes

Sampling is not implemented; this model evaluates densities only.

log_prob(value)

Log probability density at value.

Parameters:

value (Array) – Values whose last axis is \((m_1, q)\).

Returns:

The log density, of shape batch_shape. It is \(-\infty\) where \(m_1 \leq m_{\min}\), since the conditional support of \(q\) is then empty.

Return type:

Array

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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

Smoothed conditional mass ratio power law over an arbitrary primary mass model.

Generalises SmoothedPowerlawPrimaryMassRatio by taking the primary mass distribution as an argument instead of fixing it to a power law. The mass bounds are read off that distribution’s interval support, and the same Planck taper is applied to both coordinates.

\[p(m_1) \propto p_{\text{prim}}(m_1) S\!\left(\frac{m_1 - m_{1,\text{min}}}{\delta_{m_1}}\right)\]

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

Parameters:

• primary_mass_distribution (Distribution) – Distribution of the primary mass. Must have an interval support, whose bounds supply \(m_{1,\text{min}}\) and \(m_{\max}\).

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

• delta_m1 (ArrayLike) – Width of the smoothing window above \(m_{1,\text{min}}\).

• delta_m2 (ArrayLike) – Width of the smoothing window above \(m_{2,\text{min}}\).

• m2min (ArrayLike) – Minimum secondary mass.

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Raises:

ValueError – If primary_mass_distribution does not have an interval support.

Notes

Only the conditional mass ratio density is normalised numerically; the primary mass term is taken as given by primary_mass_distribution, so the taper leaves it slightly sub-normalised. Sampling is not implemented.

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

Log probability density at value.

Parameters:

value (ArrayLike) – Values whose last axis is \((m_1, q)\).

Returns:

The log density, of shape batch_shape, with the conditional mass ratio term normalised by interpolation on the grid built at construction.

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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.

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

log_prob(value)

Log probability density at value.

Parameters:

value (Array) – Values whose last axis is \((m_1, q)\).

Returns:

The log density, of shape batch_shape. It is \(-\infty\) where \(m_1 \leq m_{\min}\), since the conditional support of \(q\) is then empty.

Return type:

Array

sample(key, sample_shape=())

Draw samples by inverse transform sampling.

\(m_1\) is drawn from its marginal power law and \(q\) from the power law conditional on it.

Parameters:

• key (PRNGKeyArray) – JAX random key.

• sample_shape (tuple[int, ...]) – Shape of the sample batch to draw. Defaults to ().

Returns:

Samples of shape sample_shape + (2,), holding \((m_1, q)\).

Return type:

Array

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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

Smoothed broken power law primary mass with a conditional mass ratio power law.

Same construction as SmoothedPowerlawPrimaryMassRatio, with the single power law over the primary mass replaced by a broken one:

\[p(m_1) \propto p_{\mathrm{BP}}(m_1\mid\alpha_1,\alpha_2,m_{1,\text{min}}, m_{\text{max}},m_{\text{break}}) S\!\left(\frac{m_1 - m_{1,\text{min}}}{\delta_{m_1}}\right)\]

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

Both normalisation constants are computed numerically at construction, the conditional one as a function of \(m_1\) that log_prob() interpolates.

Parameters:

• alpha1 (ArrayLike) – Power law index below the break.

• alpha2 (ArrayLike) – Power law index above the break.

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

• mbreak (ArrayLike) – Break mass \(m_{\mathrm{break}}\).

• mmax (ArrayLike) – Maximum mass.

• m1min (ArrayLike) – Minimum primary mass.

• m2min (ArrayLike) – Minimum secondary mass.

• delta_m1 (ArrayLike) – Width of the smoothing window above \(m_{1,\text{min}}\).

• delta_m2 (ArrayLike) – Width of the smoothing window above \(m_{2,\text{min}}\).

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Notes

Sampling is not implemented; this model evaluates densities only.

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

Log probability density at value.

Parameters:

value (ArrayLike) – Values whose last axis is \((m_1, q)\).

Returns:

The normalised log density, of shape batch_shape.

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{2,\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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.

Parameters:

• loc (ArrayLike) – Location \(\mu\) of the primary mass.

• scale (ArrayLike) – Scale \(\sigma\) of the primary mass.

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

• m1min (ArrayLike) – Minimum primary mass.

• m2min (ArrayLike) – Minimum secondary mass.

• mmax (ArrayLike) – Maximum mass.

• delta_m1 (ArrayLike) – Width of the smoothing window above \(m_{1,\text{min}}\).

• delta_m2 (ArrayLike) – Width of the smoothing window above \(m_{2,\text{min}}\).

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Notes

Both normalisation constants are computed numerically at construction, the conditional one as a function of \(m_1\) that log_prob() interpolates. Sampling is not implemented; this model evaluates densities only.

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

Log probability density at value.

Parameters:

value (ArrayLike) – Values whose last axis is \((m_1, q)\).

Returns:

The normalised log density, of shape batch_shape.

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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

Smoothed power law primary mass with a conditional mass ratio power law.

Both the primary mass and the mass ratio are tapered on their lower edges by a Planck window, each with its own width:

\[p(m_1) \propto m_1^{-\alpha} S\!\left(\frac{m_1 - m_{1,\text{min}}}{\delta_{m_1}}\right)\]

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

Both normalisation constants are computed numerically at construction, the conditional one as a function of \(m_1\) that log_prob() interpolates.

Parameters:

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

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

• delta_m1 (ArrayLike) – Width of the smoothing window above \(m_{1,\text{min}}\).

• delta_m2 (ArrayLike) – Width of the smoothing window above \(m_{2,\text{min}}\).

• mmax (ArrayLike) – Maximum mass.

• m1min (ArrayLike) – Minimum primary mass.

• m2min (ArrayLike) – Minimum secondary mass.

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Notes

Sampling is not implemented; this model evaluates densities only.

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

Log probability density at value.

Parameters:

value (ArrayLike) – Values whose last axis is \((m_1, q)\).

Returns:

The normalised log density, of shape batch_shape.

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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

Smoothed power-law-plus-peak primary mass with a conditional mass ratio power law.

The primary mass is a mixture of a truncated power law and a Gaussian peak, both tapered on the lower edge by a Planck window of width \(\delta\):

\[p(m_1) \propto \left[(1-\lambda)\,\frac{m_1^{-\alpha}}{Z_{\text{pl}}} + \lambda\,\mathcal{N}(m_1\mid\mu,\sigma^2)\right] S\!\left(\frac{m_1 - m_{\text{min}}}{\delta}\right)\]

\[p(q\mid m_1) \propto q^{\beta} S\!\left(\frac{m_1 q - m_{\text{min}}}{\delta}\right)\]

Normalisation is computed once at construction: the primary mass constant splits into a numerical integral over the taper region \([m_{\min}, m_{\min}+\delta]\) and the analytic integral above it, while the conditional mass ratio constant is tabulated over \(m_1\) and interpolated in log_prob().

Parameters:

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

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

• delta (ArrayLike) – Width of the smoothing window above \(m_{\text{min}}\).

• lambda_peak (ArrayLike) – Mixing fraction \(\lambda\) of the Gaussian peak, in \([0, 1]\).

• loc (ArrayLike) – Location \(\mu\) of the Gaussian peak.

• mmax (ArrayLike) – Maximum mass.

• mmin (ArrayLike) – Minimum mass.

• scale (ArrayLike) – Scale \(\sigma\) of the Gaussian peak.

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

Notes

Sampling is not implemented; this model evaluates densities only.

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

Log probability density at value.

Parameters:

value (ArrayLike) – Values whose last axis is \((m_1, q)\).

Returns:

The normalised log density, of shape batch_shape. The conditional normalisation of \(q\) is interpolated from the grid built at construction.

Return type:

ArrayLike

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, q)\) coordinates.

Return type:

constraints.Constraint

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

Bases: numpyro.distributions.Distribution

Double side truncated power law mass model.

As described in Equation 7 of Reconstructing phenomenological distributions of compact binaries via gravitational wave observations: the primary mass follows a truncated power law and the secondary is uniform between \(m_{\text{min}}\) and \(m_1\).

\[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.

• validate_args (Optional[bool], optional) – Whether to validate distribution parameters and inputs. Defaults to None.

log_prob(value)

Log probability density at value.

Parameters:

value (Array) – Values whose last axis is \((m_1, m_2)\).

Returns:

The log density, of shape batch_shape.

Return type:

Array

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

Draw samples from the model.

\(m_1\) is drawn from the truncated power law by inverse transform sampling, then \(m_2\) uniformly on \([m_{\min}, m_1]\).

Parameters:

• key (PRNGKeyArray) – JAX random key.

• sample_shape (tuple[int, ...]) – Shape of the sample batch to draw. Defaults to ().

Returns:

Samples of shape sample_shape + (2,), holding \((m_1, m_2)\).

Return type:

Array

support() → numpyro.distributions.constraints.Constraint

The support of the distribution.

Returns:

The mass sandwich \(m_{\min} \leq m_2 \leq m_1 \leq m_{\max}\), expressed in \((m_1, m_2)\) coordinates.

Return type:

constraints.Constraint