gwkokab.models

Provides implementation of various models, constraints, and transformations using NumPyro.

Submodules

• gwkokab.models.constraints
• gwkokab.models.hybrids
• gwkokab.models.mass
• gwkokab.models.redshift
• gwkokab.models.spin
• gwkokab.models.sundry
• gwkokab.models.transformations
• gwkokab.models.utils
• gwkokab.models.wrappers

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
GWTC4EffectiveSpinSkewNormalModel	GWTC-4 effective spin skew normal model.
MadauDickinsonRedshiftModel	Redshift distribution for compact binary mergers modeled after the Madau-
PowerlawPrimaryMassRatio	Power law model for two-dimensional mass distribution, modelling primary mass and
PowerlawRedshiftModel	Redshift distribution for compact binary mergers modeled as a power law modulated
SmoothedTwoComponentPrimaryMassRatio	Base class for probability distributions in NumPyro. The design largely
Wysocki2019MassModel	It is a double side truncated power law distribution, as described in

Functions

BetaFromMeanVar(→ numpyro.distributions.Beta)	Beta distribution parameterized by the expected value and variance.
GaussianSpinModel(...)	Bivariate normal distribution for the effective and precessing spins.
GenericTiltModel(→ numpyro.distributions.MixtureGeneral)	A mixture model of spin orientations with isotropic and normally distributed
MultiSourceModel(→ gwkokab.models.utils.ScaledMixture)
NDIsotropicAndTruncatedNormalMixture(...)	General N-dimensional mixture model of an isotropic uniform distribution and a
NPowerlawMGaussian(→ gwkokab.models.utils.ScaledMixture)	Create a mixture of power-law and Gaussian components.
SubPopulationModel(→ gwkokab.models.utils.ScaledMixture)
TwoTruncatedNormalMixture(...)	Create a mixture model of two truncated normal distributions.

Package Contents

class gwkokab.models.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.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.GWTC4EffectiveSpinSkewNormalModel(loc: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, epsilon: jaxtyping.ArrayLike, *, validate_args: bool | None = None)

Bases: numpyro.distributions.Distribution

GWTC-4 effective spin skew normal model.

This class implements effective spin skew normal model introduced in equation (B37) GWTC-4.0: Population Properties of Merging Compact Binaries.

\[\begin{split}p(\chi_\mathrm{eff} | \mu, \sigma, \epsilon) \propto \begin{cases} (1 + \epsilon) \mathcal{N}_{[-1,1]}(\chi_\mathrm{eff} | \mu, \sigma (1 + \epsilon)), & \chi_\mathrm{eff} \leq \mu \\ (1 - \epsilon) \mathcal{N}_{[-1,1]}(\chi_\mathrm{eff} | \mu, \sigma (1 - \epsilon)), & \chi_\mathrm{eff} > \mu \end{cases}\end{split}\]

where \(\mathcal{N}_{[-1,1]}(x | \mu, \sigma)\) is the truncated normal distribution with mean \(\mu\) and standard deviation \(\sigma\), truncated to the interval \([-1, 1]\).

The normalization constant is expressed as:

\[\mathcal{Z} = \frac{1 - \epsilon}{2} \left[\frac{\mathrm{erf}\left( \displaystyle -\frac{1 + \mu}{\sqrt{2}\sigma (1 - \epsilon)} \right)}{\Phi\left( \displaystyle\frac{1 - \mu}{\sigma (1 - \epsilon)} \right) - \Phi\left( \displaystyle\frac{-1 - \mu}{\sigma (1 - \epsilon)} \right)} \right] -\frac{1 + \epsilon}{2} \left[ \frac{\mathrm{erf}\left( \displaystyle -\frac{1 + \mu}{\sqrt{2}\sigma (1 + \epsilon)} \right)}{\Phi\left( \displaystyle\frac{1 - \mu}{\sigma (1 + \epsilon)} \right) - \Phi\left( \displaystyle\frac{-1 - \mu}{\sigma (1 + \epsilon)} \right)} \right]\]

where, \(\Phi(x)\) is the cumulative distribution function of the standard normal distribution.

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

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

class gwkokab.models.MadauDickinsonRedshiftModel(z_max: jax.Array, kappa: jax.Array, gamma: jax.Array, z_peak: jax.Array, *, validate_args: bool | None = None)

Bases: _RedshiftModel

Redshift distribution for compact binary mergers modeled after the Madau- Dickinson star formation rate, modulated by the cosmological volume element.

The probability density function is defined as:

\[\begin{split}p(z) \propto \frac{dV_c/dz(z) \cdot (1 + z)^{\kappa - 1}}{1 + \left(\\frac{1 + z}{1 + z_{peak}}\right)^{\gamma}}, \ \qquad 0 \leq z \leq z_{max}\end{split}\]

where:

• dV_c/dz is the differential comoving volume element,

• gamma is the high-redshift slope,

• kappa is the low-redshift slope,

• z_max is the upper redshift cutoff.

• z_peak is the redshift at which the merger rate peaks,

This distribution is normalized numerically on a fixed redshift grid.

log_psi_of_z(z: jax.Array) → jax.Array

Evaluate the psi function at a given redshift.

\[\ln\psi(z) = \kappa \ln(1 + z) + \ln\left(1 + (1 + z_{peak})^{\gamma}\right) - \ln\left((1 + z_{peak})^{\gamma} + (1 + z)^{\gamma}\right)\]

Parameters:

z (ArrayLike) – Redshift(s) to evaluate.

Returns:

Values of the psi function.

Return type:

ArrayLike

class gwkokab.models.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.PowerlawRedshiftModel(z_max: jax.Array, kappa: jax.Array, *, validate_args: bool | None = None)

Bases: _RedshiftModel

Redshift distribution for compact binary mergers modeled as a power law modulated by the cosmological volume element.

The probability density function is defined as:

\[p(z) \propto \frac{dV_c/dz(z) \cdot (1 + z)^{\kappa - 1}}}, \qquad 0 \leq z \leq z_{max}\]

where:

• dV_c/dz is the differential comoving volume element,

• is the redshift evolution power-law index,

• z_max is the upper redshift cutoff.

This distribution is normalized numerically on a fixed redshift grid.

Parameters:

• kappa (float) – The power-law exponent \(\kappa\).

• z_max (float) – The maximum redshift (upper limit of the support).

log_psi_of_z(z: jax.Array) → jax.Array

Evaluate the psi function at a given redshift.

\[\ln\psi(z) = \kappa \ln(1 + z)\]

Parameters:

z (ArrayLike) – Redshift(s) to evaluate.

Returns:

Values of the psi function.

Return type:

ArrayLike

class gwkokab.models.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.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.

gwkokab.models.BetaFromMeanVar(mean: jaxtyping.ArrayLike, variance: jaxtyping.ArrayLike, *, validate_args: bool | None = None) → numpyro.distributions.Beta

Beta distribution parameterized by the expected value and variance.

Parameters:

• mean (ArrayLike) – Expected value of the beta distribution.

• variance (ArrayLike) – Variance of the beta distribution.

• loc (ArrayLike) – lower bound of the beta distribution, defaults to 0.0

• scale (ArrayLike) – width of the beta distribution, defaults to 1.0

Returns:

Beta distribution with the specified mean and variance.

Return type:

Beta

gwkokab.models.GaussianSpinModel(mu_eff: jaxtyping.ArrayLike, sigma_eff: jaxtyping.ArrayLike, mu_p: jaxtyping.ArrayLike, sigma_p: jaxtyping.ArrayLike, rho: jaxtyping.ArrayLike, *, validate_args: bool | None = None) → numpyro.distributions.MultivariateNormal

Bivariate normal distribution for the effective and precessing spins. See Eq. (D3) and (D4) in Population Properties of Compact Objects from the Second LIGO-Virgo Gravitational-Wave Transient Catalog.

\[\begin{split}\left(\chi_{\text{eff}}, \chi_{p}\right) \sim \mathcal{N}\left( \begin{bmatrix} \mu_{\text{eff}} \\ \mu_{p} \end{bmatrix}, \begin{bmatrix} \sigma_{\text{eff}}^2 & \rho \sigma_{\text{eff}} \sigma_{p} \\ \rho \sigma_{\text{eff}} \sigma_{p} & \sigma_{p}^2 \end{bmatrix} \right)\end{split}\]

where \(\chi_{\text{eff}}\) is the effective spin and \(\chi_{\text{eff}}\in[-1,1]\) and \(\chi_{p}\) is the precessing spin and \(\chi_{p}\in[0,1]\).

Parameters:

• mu_eff (ArrayLike) – mean of the effective spin

• sigma_eff (ArrayLike) – standard deviation of the effective spin

• mu_p (ArrayLike) – mean of the precessing spin

• sigma_p (ArrayLike) – standard deviation of the precessing spin

• rho (ArrayLike) – correlation coefficient between the effective and precessing spins

Returns:

Multivariate normal distribution for the effective and precessing spins

Return type:

MultivariateNormal

gwkokab.models.GenericTiltModel(zeta: jaxtyping.ArrayLike, loc1: jaxtyping.ArrayLike, loc2: jaxtyping.ArrayLike, scale1: jaxtyping.ArrayLike, scale2: jaxtyping.ArrayLike, low1: jaxtyping.ArrayLike = -1.0, low2: jaxtyping.ArrayLike = -1.0, high1: jaxtyping.ArrayLike = 1.0, high2: jaxtyping.ArrayLike = 1.0, *, validate_args: bool | None = None) → numpyro.distributions.MixtureGeneral

A mixture model of spin orientations with isotropic and normally distributed components, with a minimum and maximum tilt constraint for each spin.

Parameters:

• zeta (ArrayLike) – Weight of the Gaussian component.

• loc1 (ArrayLike) – Location parameter of the first Gaussian component.

• loc2 (ArrayLike) – Location parameter of the second Gaussian component.

• scale1 (ArrayLike) – Scale parameter of the first Gaussian component.

• scale2 (ArrayLike) – Scale parameter of the second Gaussian component.

• low1 (ArrayLike, optional) – Minimum cosine tilt angle of the first component, by default -1.0

• low2 (ArrayLike, optional) – Minimum cosine tilt angle of the second component, by default -1.0

• high1 (ArrayLike, optional) – Maximum cosine tilt angle of the first component, by default 1.0

• high2 (ArrayLike, optional) – Maximum cosine tilt angle of the second component, by default 1.0

• validate_args (Optional[bool], optional) – Whether to validate the arguments, by default None

Returns:

Mixture model of spin orientations with minimum and maximum tilt constraints for each spin.

Return type:

MixtureGeneral

gwkokab.models.MultiSourceModel(N_spl: int, N_bpl: int, N_gpl: int, N_gg: int, use_beta_spin_magnitude: bool = False, use_spin_magnitude_mixture: bool = False, use_truncated_normal_spin_x: bool = False, use_truncated_normal_spin_y: bool = False, use_truncated_normal_spin_z: bool = False, use_chi_eff_mixture: bool = False, use_skew_normal_chi_eff: bool = False, use_truncated_normal_chi_p: bool = False, use_tilt: bool = False, use_eccentricity_mixture: bool = False, use_eccentricity_powerlaw: bool = False, use_mean_anomaly: bool = False, use_powerlaw_redshift: bool = False, use_madau_dickinson_redshift: bool = False, *, validate_args=None, **params) → gwkokab.models.utils.ScaledMixture

gwkokab.models.NDIsotropicAndTruncatedNormalMixture(zeta: jaxtyping.ArrayLike, loc: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, isotropic_low: jaxtyping.ArrayLike, isotropic_high: jaxtyping.ArrayLike, gaussian_low: jaxtyping.ArrayLike | None, gaussian_high: jaxtyping.ArrayLike | None, *, batch_dim: int = 1, validate_args: bool | None = None) → numpyro.distributions.MixtureGeneral

General N-dimensional mixture model of an isotropic uniform distribution and a truncated normal distribution.

\[p(\mathbf{x}\mid\zeta,\boldsymbol{\mu},\boldsymbol{\sigma}) = (1-\zeta)\mathcal{U}(\mathbf{x}\mid \boldsymbol{a},\boldsymbol{b}) + \zeta\mathcal{N}_{[\boldsymbol{L}, \boldsymbol{U}]}(\mathbf{x}\mid \boldsymbol{\mu},\boldsymbol{\sigma})\]

where \(\mathcal{U}(\cdot)\) is the isotropic uniform distribution between \(\boldsymbol{a}=\left< a_1, a_2, \ldots, a_N \right>\) and \(\boldsymbol{b}=\left< b_1, b_2, \ldots, b_N \right>\), and \(\mathcal{N}_{[\boldsymbol{L}, \boldsymbol{U}]}(\cdot)\) is the truncated normal distribution with mean \(\boldsymbol{\mu}=\left< \mu_1, \mu_2, \ldots, \mu_N \right>\), standard deviation \(\boldsymbol{\sigma}=\left< \sigma_1, \sigma_2, \ldots, \sigma_N \right>\), lower bound \(\boldsymbol{L}=\left< L_1, L_2, \ldots, L_N \right>\), and upper bound \(\boldsymbol{U}=\left< U_1, U_2, \ldots, U_N \right>\).

Parameters:

• zeta (ArrayLike) – The mixing probability of the second component.

• loc (ArrayLike) – The mean of the truncated normal distribution.

• scale (ArrayLike) – The standard deviation of the truncated normal distribution.

• isotropic_low (ArrayLike) – The lower bound of the isotropic uniform distribution.

• isotropic_high (ArrayLike) – The upper bound of the isotropic uniform distribution.

• gaussian_low (Optional[ArrayLike]) – The lower bound of the truncated normal distribution.

• gaussian_high (Optional[ArrayLike]) – The upper bound of the truncated normal distribution.

• batch_dim (int, optional) – The batch dimension of the distributions, by default 1

• validate_args (Optional[bool], optional) – Whether to validate the parameters of the distributions, by default None

Returns:

N-dimensional mixture model of an isotropic uniform distribution and a truncated normal distribution.

Return type:

MixtureGeneral

gwkokab.models.NPowerlawMGaussian(N_pl: int, N_g: int, use_beta_spin_magnitude: bool = False, use_spin_magnitude_mixture: bool = False, use_truncated_normal_spin_x: bool = False, use_truncated_normal_spin_y: bool = False, use_truncated_normal_spin_z: bool = False, use_chi_eff_mixture: bool = False, use_skew_normal_chi_eff: bool = False, use_truncated_normal_chi_p: bool = False, use_tilt: bool = False, use_eccentricity_mixture: bool = False, use_eccentricity_powerlaw: bool = False, use_powerlaw_redshift: bool = False, use_madau_dickinson_redshift: bool = False, use_cos_iota: bool = False, use_phi_12: bool = False, use_polarization_angle: bool = False, use_right_ascension: bool = False, use_sin_declination: bool = False, use_detection_time: bool = False, use_phi_1: bool = False, use_phi_2: bool = False, use_phi_orb: bool = False, use_mean_anomaly: bool = False, *, validate_args=None, **params) → gwkokab.models.utils.ScaledMixture

Create a mixture of power-law and Gaussian components.

This model has a lot of parameters, we can not list all of them here. Therefore, we are providing the general description of the each sub model and their parameters.

Important

Important information about this models are as follows:

• The first N_pl components are power-law and the next N_g components are Gaussian.

• Log rates are named as log_rate_{i} where i is the index of the component.

• First N_pl log rates are for power-law components and the next N_g log rates are for Gaussian components.

• All log rates are in terms of natural logarithm.

Note

Mass distribution: For powerlaw mass distribution is PowerlawPrimaryMassRatio and for Gaussian we have TruncatedNormal distribution.

\[(m_1, m_2) \sim \text{PowerlawPrimaryMassRatio}(\alpha, \beta, m_{\text{min}}, m_{\text{max}})\]

\[(m_1, m_2) \sim \mathcal{N}_{[a,b]}(\mu, \sigma^2)\]

Spin distribution: Spin is taken from a beta distribution. The beta distribution is parameterized by \(\mu\) and \(\sigma^2\) where \(\mu\) is the mean and \(\sigma^2\) is the variance of the beta distribution.

\[\chi_i \sim \mathrm{Beta}(\mu, \sigma^2)\]

Warning

Not every choice of \(\mu\) and \(\sigma^2\) will result in a valid beta distribution. The beta distribution is only valid when \(\mu \in (0, 1)\), \(\sigma^2 \in (0, 0.25)\), and \(\mu(1-\mu) > \sigma^2\). Refer to the link for more information.

Tilt distribution: Tilt is taken from a TruncatedNormal distribution, with fixed mean \(\mu=1\) and fixed bounds, \(a=-1\) and \(b=1\).

\[\cos(\theta_i) \sim \mathcal{N}_{[-1, 1]}(\sigma^2\mid\mu=1)\]

Eccentricity distribution: Eccentricity is taken from TruncatedNormal distribution.

\[\varepsilon_i \sim \mathcal{N}_{[a,b]}(\mu, \sigma^2)\]

Attention

Interestingly, in TruncatedNormal distribution, if any of the bounds are not provided, it will be set to \(\pm\infty\). For example, if we set \(\mu=0\) and do not provide the upper bound, then the resulting distribution would be a half normal distribution.

The naming of the parameters follows the following convention:

<parameter name>_<model parameter>_<component type>_<component number>

with an exception for the powerlaw mass distribution where the <parameter name> is ignored. For example, spin is taken from a beta distribution whose parameters are mean and variance. The naming convention for the spin parameters would be:

chi[1-2]_mean_(pl|g)_[0-N_pl+N_g]
chi[1-2]_variance_(pl|g)_[0-N_pl+N_g]

Parameters:

• N_pl (int) – Number of power-law components

• N_g (int) – Number of Gaussian components

• use_spin (bool) – whether to include spin, defaults to False

• use_tilt (bool) – whether to include tilt, defaults to False

• use_eccentricity_mixture (bool) – whether to include eccentricity, defaults to False

• use_mean_anomaly (bool) – whether to include mean_anomaly, defaults to False

• use_powerlaw_redshift (bool) – whether to include redshift, defaults to False

• use_cos_iota (bool) – whether to include cos_iota, defaults to False

• use_polarization_angle (bool) – whether to include polarization_angle, defaults to False

• use_right_ascension (bool) – whether to include right_ascension, defaults to False

• use_sin_declination (bool) – whether to include sin_declination, defaults to False

• use_detection_time (bool) – whether to include detection_time, defaults to False

• use_phi_1 (bool) – whether to include phi_1, defaults to False

• use_phi_2 (bool) – whether to include phi_2, defaults to False

• use_phi_12 (bool) – whether to include phi_12, defaults to False

• use_phi_orb (bool) – whether to include phi_orb, defaults to False

• validate_args (Optional[bool], optional) – whether to validate arguments, defaults to None

Returns:

scaled mixture of distributions

Return type:

ScaledMixture

gwkokab.models.SubPopulationModel(N_spl: int, N_bpl: int, N_gpl: int, use_beta_spin_magnitude: bool = False, use_spin_magnitude_mixture: bool = False, use_truncated_normal_spin_x: bool = False, use_truncated_normal_spin_y: bool = False, use_truncated_normal_spin_z: bool = False, use_chi_eff_mixture: bool = False, use_skew_normal_chi_eff: bool = False, use_truncated_normal_chi_p: bool = False, use_tilt: bool = False, use_eccentricity_mixture: bool = False, use_eccentricity_powerlaw: bool = False, use_mean_anomaly: bool = False, use_powerlaw_redshift: bool = False, use_madau_dickinson_redshift: bool = False, *, validate_args=None, **params) → gwkokab.models.utils.ScaledMixture

gwkokab.models.TwoTruncatedNormalMixture(comp1_high: jaxtyping.ArrayLike | None, comp1_loc: jaxtyping.ArrayLike, comp1_low: jaxtyping.ArrayLike | None, comp1_scale: jaxtyping.ArrayLike, comp2_high: jaxtyping.ArrayLike | None, comp2_loc: jaxtyping.ArrayLike, comp2_low: jaxtyping.ArrayLike | None, comp2_scale: jaxtyping.ArrayLike, zeta: jaxtyping.ArrayLike, *, validate_args: bool | None = None) → numpyro.distributions.MixtureGeneral

Create a mixture model of two truncated normal distributions.

Parameters:

• comp1_high (Optional[ArrayLike]) – Upper truncation for the first component.

• comp1_loc (ArrayLike) – Location parameter for the first component.

• comp1_low (Optional[ArrayLike]) – Lower truncation for the first component.

• comp1_scale (ArrayLike) – Scale parameter for the first component.

• comp2_high (Optional[ArrayLike]) – Upper truncation for the second component.

• comp2_loc (ArrayLike) – Location parameter for the second component.

• comp2_low (Optional[ArrayLike]) – Lower truncation for the second component.

• comp2_scale (ArrayLike) – Scale parameter for the second component.

• zeta (ArrayLike) – Mixing proportion for the second component.

• validate_args (Optional[bool], optional) – Whether to validate the arguments, by default None

Returns:

A mixture of two truncated normal distributions.

Return type:

MixtureGeneral