gwkokab.models ============== .. py:module:: gwkokab.models .. autoapi-nested-parse:: Provides implementation of various models, constraints, and transformations using `NumPyro `_. Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/gwkokab/models/constraints/index /autoapi/gwkokab/models/hybrids/index /autoapi/gwkokab/models/mass/index /autoapi/gwkokab/models/redshift/index /autoapi/gwkokab/models/spin/index /autoapi/gwkokab/models/sundry/index /autoapi/gwkokab/models/transformations/index /autoapi/gwkokab/models/utils/index /autoapi/gwkokab/models/wrappers/index Classes ------- .. autoapisummary:: gwkokab.models.BrokenPowerlaw gwkokab.models.BrokenPowerlawTwoPeak gwkokab.models.GWTC4EffectiveSpinSkewNormalModel gwkokab.models.MadauDickinsonRedshiftModel gwkokab.models.PowerlawPrimaryMassRatio gwkokab.models.PowerlawRedshiftModel gwkokab.models.SmoothedTwoComponentPrimaryMassRatio gwkokab.models.Wysocki2019MassModel Functions --------- .. autoapisummary:: gwkokab.models.BetaFromMeanVar gwkokab.models.GaussianSpinModel gwkokab.models.GenericTiltModel gwkokab.models.MultiSourceModel gwkokab.models.NDIsotropicAndTruncatedNormalMixture gwkokab.models.NPowerlawMGaussian gwkokab.models.SubPopulationModel gwkokab.models.TwoTruncatedNormalMixture Package Contents ---------------- .. py:class:: BrokenPowerlaw(alpha1: jaxtyping.ArrayLike, alpha2: jaxtyping.ArrayLike, mbreak: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, validate_args: Optional[bool] = None) Bases: :py:obj:`numpyro.distributions.Distribution` Base class for probability distributions in NumPyro. The design largely follows from :mod:`torch.distributions`. :param 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. :param 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`. :param validate_args: Whether to enable validation of distribution parameters and arguments to `.log_prob` method. As an example: .. doctest:: >>> import jax.numpy as jnp >>> import numpyro.distributions as dist >>> d = dist.Dirichlet(jnp.ones((2, 3, 4))) >>> d.batch_shape (2, 3) >>> d.event_shape (4,) .. py:method:: log_prob(value: jaxtyping.ArrayLike) -> jaxtyping.ArrayLike Evaluates the log probability density for a batch of samples given by `value`. :param value: A batch of samples from the distribution. :return: an array with shape `value.shape[:-self.event_shape]` :rtype: ArrayLike .. py:method:: support() -> numpyro.distributions.constraints.Constraint The support of this distribution. Subclasses can override this as a class attribute or as a property. .. py:class:: 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: Optional[bool] = None) Bases: :py:obj:`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: .. math:: 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} .. math:: 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) .. math:: 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) .. math:: 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 :math:`\mathcal{N}_{[m_{1,\mathrm{min}}, \infty)}(m \mid \mu, \sigma)` is a left truncated normal distribution with mean :math:`\mu` and standard deviation :math:`\sigma`, :math:`\lambda_0` and :math:`\lambda_1` are the mixing fractions of the broken powerlaw and first Gaussian component respectively, :math:`\delta_{m_1}` is the smoothing scale for the primary mass, and :math:`S(x)` is the exponential of :func:`~gwkokab.utils.kernel.log_planck_taper_window`. .. py:method:: log_prob(value: jaxtyping.ArrayLike) -> jaxtyping.ArrayLike Evaluates the log probability density for a batch of samples given by `value`. :param value: A batch of samples from the distribution. :return: an array with shape `value.shape[:-self.event_shape]` :rtype: ArrayLike .. py:method:: support() -> numpyro.distributions.constraints.Constraint The support of this distribution. Subclasses can override this as a class attribute or as a property. .. py:class:: GWTC4EffectiveSpinSkewNormalModel(loc: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, epsilon: jaxtyping.ArrayLike, *, validate_args: Optional[bool] = None) Bases: :py:obj:`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 `_. .. math:: 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} where :math:`\mathcal{N}_{[-1,1]}(x | \mu, \sigma)` is the truncated normal distribution with mean :math:`\mu` and standard deviation :math:`\sigma`, truncated to the interval :math:`[-1, 1]`. The normalization constant is expressed as: .. math:: \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, :math:`\Phi(x)` is the cumulative distribution function of the standard normal distribution. .. py:method:: log_prob(value: jaxtyping.ArrayLike) -> jaxtyping.ArrayLike Evaluates the log probability density for a batch of samples given by `value`. :param value: A batch of samples from the distribution. :return: an array with shape `value.shape[:-self.event_shape]` :rtype: ArrayLike .. py:attribute:: support The support of this distribution. Subclasses can override this as a class attribute or as a property. .. py:class:: MadauDickinsonRedshiftModel(z_max: jax.Array, kappa: jax.Array, gamma: jax.Array, z_peak: jax.Array, *, validate_args: Optional[bool] = None) Bases: :py:obj:`_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: .. math:: 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} 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. .. py:method:: log_psi_of_z(z: jax.Array) -> jax.Array Evaluate the psi function at a given redshift. .. math:: \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) :param z: Redshift(s) to evaluate. :type z: ArrayLike :returns: Values of the psi function. :rtype: ArrayLike .. py:class:: PowerlawPrimaryMassRatio(alpha: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, *, validate_args: Optional[bool] = None) Bases: :py:obj:`numpyro.distributions.Distribution` Power law model for two-dimensional mass distribution, modelling primary mass and conditional mass ratio distribution. .. math:: p(m_1,q\mid\alpha,\beta) = p(m_1\mid\alpha)p(q \mid m_1, \beta) .. math:: \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*} :param alpha: Power law index for primary mass :type alpha: ArrayLike :param beta: Power law index for mass ratio :type beta: ArrayLike :param mmin: Minimum mass :type mmin: ArrayLike :param mmax: Maximum mass :type mmax: ArrayLike .. py:method:: log_prob(value) Evaluates the log probability density for a batch of samples given by `value`. :param value: A batch of samples from the distribution. :return: an array with shape `value.shape[:-self.event_shape]` :rtype: ArrayLike .. py:method:: 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. :param jax.random.key key: the rng_key key to be used for the distribution. :param tuple sample_shape: the sample shape for the distribution. :return: an array of shape `sample_shape + batch_shape + event_shape` :rtype: numpy.ndarray .. py:method:: support() -> numpyro.distributions.constraints.Constraint The support of this distribution. Subclasses can override this as a class attribute or as a property. .. py:class:: PowerlawRedshiftModel(z_max: jax.Array, kappa: jax.Array, *, validate_args: Optional[bool] = None) Bases: :py:obj:`_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: .. math:: 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. :param kappa: The power-law exponent :math:`\kappa`. :type kappa: float :param z_max: The maximum redshift (upper limit of the support). :type z_max: float .. py:method:: log_psi_of_z(z: jax.Array) -> jax.Array Evaluate the psi function at a given redshift. .. math:: \ln\psi(z) = \kappa \ln(1 + z) :param z: Redshift(s) to evaluate. :type z: ArrayLike :returns: Values of the psi function. :rtype: ArrayLike .. py:class:: 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: :py:obj:`numpyro.distributions.Distribution` Base class for probability distributions in NumPyro. The design largely follows from :mod:`torch.distributions`. :param 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. :param 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`. :param validate_args: Whether to enable validation of distribution parameters and arguments to `.log_prob` method. As an example: .. doctest:: >>> import jax.numpy as jnp >>> import numpyro.distributions as dist >>> d = dist.Dirichlet(jnp.ones((2, 3, 4))) >>> d.batch_shape (2, 3) >>> d.event_shape (4,) .. py:method:: log_prob(value: jaxtyping.ArrayLike) -> jaxtyping.ArrayLike Evaluates the log probability density for a batch of samples given by `value`. :param value: A batch of samples from the distribution. :return: an array with shape `value.shape[:-self.event_shape]` :rtype: ArrayLike .. py:method:: support() -> numpyro.distributions.constraints.Constraint The support of this distribution. Subclasses can override this as a class attribute or as a property. .. py:class:: Wysocki2019MassModel(alpha_m: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, mmax: jaxtyping.ArrayLike, *, validate_args: Optional[bool] = None) Bases: :py:obj:`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 `_. .. math:: 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}}} :param alpha_m: index of the power law distribution :type alpha_m: ArrayLike :param mmin: lower mass limit :type mmin: ArrayLike :param mmax: upper mass limit :type mmax: ArrayLike .. py:method:: log_prob(value) Evaluates the log probability density for a batch of samples given by `value`. :param value: A batch of samples from the distribution. :return: an array with shape `value.shape[:-self.event_shape]` :rtype: ArrayLike .. py:method:: 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. :param jax.random.key key: the rng_key key to be used for the distribution. :param tuple sample_shape: the sample shape for the distribution. :return: an array of shape `sample_shape + batch_shape + event_shape` :rtype: numpy.ndarray .. py:method:: support() -> numpyro.distributions.constraints.Constraint The support of this distribution. Subclasses can override this as a class attribute or as a property. .. py:function:: BetaFromMeanVar(mean: jaxtyping.ArrayLike, variance: jaxtyping.ArrayLike, *, validate_args: Optional[bool] = None) -> numpyro.distributions.Beta Beta distribution parameterized by the expected value and variance. :param mean: Expected value of the beta distribution. :type mean: ArrayLike :param variance: Variance of the beta distribution. :type variance: ArrayLike :param loc: lower bound of the beta distribution, defaults to 0.0 :type loc: ArrayLike :param scale: width of the beta distribution, defaults to 1.0 :type scale: ArrayLike :returns: Beta distribution with the specified mean and variance. :rtype: Beta .. py:function:: GaussianSpinModel(mu_eff: jaxtyping.ArrayLike, sigma_eff: jaxtyping.ArrayLike, mu_p: jaxtyping.ArrayLike, sigma_p: jaxtyping.ArrayLike, rho: jaxtyping.ArrayLike, *, validate_args: Optional[bool] = 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 `_. .. math:: \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) where :math:`\chi_{\text{eff}}` is the effective spin and :math:`\chi_{\text{eff}}\in[-1,1]` and :math:`\chi_{p}` is the precessing spin and :math:`\chi_{p}\in[0,1]`. :param mu_eff: mean of the effective spin :type mu_eff: ArrayLike :param sigma_eff: standard deviation of the effective spin :type sigma_eff: ArrayLike :param mu_p: mean of the precessing spin :type mu_p: ArrayLike :param sigma_p: standard deviation of the precessing spin :type sigma_p: ArrayLike :param rho: correlation coefficient between the effective and precessing spins :type rho: ArrayLike :returns: Multivariate normal distribution for the effective and precessing spins :rtype: MultivariateNormal .. py:function:: 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: Optional[bool] = 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. :param zeta: Weight of the Gaussian component. :type zeta: ArrayLike :param loc1: Location parameter of the first Gaussian component. :type loc1: ArrayLike :param loc2: Location parameter of the second Gaussian component. :type loc2: ArrayLike :param scale1: Scale parameter of the first Gaussian component. :type scale1: ArrayLike :param scale2: Scale parameter of the second Gaussian component. :type scale2: ArrayLike :param low1: Minimum cosine tilt angle of the first component, by default -1.0 :type low1: ArrayLike, optional :param low2: Minimum cosine tilt angle of the second component, by default -1.0 :type low2: ArrayLike, optional :param high1: Maximum cosine tilt angle of the first component, by default 1.0 :type high1: ArrayLike, optional :param high2: Maximum cosine tilt angle of the second component, by default 1.0 :type high2: ArrayLike, optional :param validate_args: Whether to validate the arguments, by default None :type validate_args: Optional[bool], optional :returns: Mixture model of spin orientations with minimum and maximum tilt constraints for each spin. :rtype: MixtureGeneral .. py:function:: 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 .. py:function:: NDIsotropicAndTruncatedNormalMixture(zeta: jaxtyping.ArrayLike, loc: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, isotropic_low: jaxtyping.ArrayLike, isotropic_high: jaxtyping.ArrayLike, gaussian_low: Optional[jaxtyping.ArrayLike], gaussian_high: Optional[jaxtyping.ArrayLike], *, batch_dim: int = 1, validate_args: Optional[bool] = None) -> numpyro.distributions.MixtureGeneral General N-dimensional mixture model of an isotropic uniform distribution and a truncated normal distribution. .. math:: 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 :math:`\mathcal{U}(\cdot)` is the isotropic uniform distribution between :math:`\boldsymbol{a}=\left< a_1, a_2, \ldots, a_N \right>` and :math:`\boldsymbol{b}=\left< b_1, b_2, \ldots, b_N \right>`, and :math:`\mathcal{N}_{[\boldsymbol{L}, \boldsymbol{U}]}(\cdot)` is the truncated normal distribution with mean :math:`\boldsymbol{\mu}=\left< \mu_1, \mu_2, \ldots, \mu_N \right>`, standard deviation :math:`\boldsymbol{\sigma}=\left< \sigma_1, \sigma_2, \ldots, \sigma_N \right>`, lower bound :math:`\boldsymbol{L}=\left< L_1, L_2, \ldots, L_N \right>`, and upper bound :math:`\boldsymbol{U}=\left< U_1, U_2, \ldots, U_N \right>`. :param zeta: The mixing probability of the second component. :type zeta: ArrayLike :param loc: The mean of the truncated normal distribution. :type loc: ArrayLike :param scale: The standard deviation of the truncated normal distribution. :type scale: ArrayLike :param isotropic_low: The lower bound of the isotropic uniform distribution. :type isotropic_low: ArrayLike :param isotropic_high: The upper bound of the isotropic uniform distribution. :type isotropic_high: ArrayLike :param gaussian_low: The lower bound of the truncated normal distribution. :type gaussian_low: Optional[ArrayLike] :param gaussian_high: The upper bound of the truncated normal distribution. :type gaussian_high: Optional[ArrayLike] :param batch_dim: The batch dimension of the distributions, by default 1 :type batch_dim: int, optional :param validate_args: Whether to validate the parameters of the distributions, by default None :type validate_args: Optional[bool], optional :returns: N-dimensional mixture model of an isotropic uniform distribution and a truncated normal distribution. :rtype: MixtureGeneral .. py:function:: 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 :code:`N_pl` components are power-law and the next :code:`N_g` components are Gaussian. * Log rates are named as :code:`log_rate_{i}` where :code:`i` is the index of the component. * First :code:`N_pl` log rates are for power-law components and the next :code:`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 :class:`PowerlawPrimaryMassRatio` and for Gaussian we have :class:`~numpyro.distributions.truncated.TruncatedNormal` distribution. .. math:: (m_1, m_2) \sim \text{PowerlawPrimaryMassRatio}(\alpha, \beta, m_{\text{min}}, m_{\text{max}}) .. math:: (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 :math:`\mu` and :math:`\sigma^2` where :math:`\mu` is the mean and :math:`\sigma^2` is the variance of the beta distribution. .. math:: \chi_i \sim \mathrm{Beta}(\mu, \sigma^2) .. warning:: Not every choice of :math:`\mu` and :math:`\sigma^2` will result in a valid beta distribution. The beta distribution is only valid when :math:`\mu \in (0, 1)`, :math:`\sigma^2 \in (0, 0.25)`, and :math:`\mu(1-\mu) > \sigma^2`. Refer to the `link `_ for more information. **Tilt distribution**: Tilt is taken from a :class:`~numpyro.distributions.truncated.TruncatedNormal` distribution, with fixed mean :math:`\mu=1` and fixed bounds, :math:`a=-1` and :math:`b=1`. .. math:: \cos(\theta_i) \sim \mathcal{N}_{[-1, 1]}(\sigma^2\mid\mu=1) **Eccentricity distribution**: Eccentricity is taken from :class:`~numpyro.distributions.truncated.TruncatedNormal` distribution. .. math:: \varepsilon_i \sim \mathcal{N}_{[a,b]}(\mu, \sigma^2) .. attention:: Interestingly, in :class:`~numpyro.distributions.truncated.TruncatedNormal` distribution, if any of the bounds are not provided, it will be set to :math:`\pm\infty`. For example, if we set :math:`\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: .. code:: bash ___ with an exception for the powerlaw mass distribution where the :code:`` is ignored. For example, spin is taken from a beta distribution whose parameters are :code:`mean` and :code:`variance`. The naming convention for the spin parameters would be: .. code:: text chi[1-2]_mean_(pl|g)_[0-N_pl+N_g] chi[1-2]_variance_(pl|g)_[0-N_pl+N_g] :param N_pl: Number of power-law components :type N_pl: int :param N_g: Number of Gaussian components :type N_g: int :param use_spin: whether to include spin, defaults to False :type use_spin: bool :param use_tilt: whether to include tilt, defaults to False :type use_tilt: bool :param use_eccentricity_mixture: whether to include eccentricity, defaults to False :type use_eccentricity_mixture: bool :param use_mean_anomaly: whether to include mean_anomaly, defaults to False :type use_mean_anomaly: bool :param use_powerlaw_redshift: whether to include redshift, defaults to False :type use_powerlaw_redshift: bool :param use_cos_iota: whether to include cos_iota, defaults to False :type use_cos_iota: bool :param use_polarization_angle: whether to include polarization_angle, defaults to False :type use_polarization_angle: bool :param use_right_ascension: whether to include right_ascension, defaults to False :type use_right_ascension: bool :param use_sin_declination: whether to include sin_declination, defaults to False :type use_sin_declination: bool :param use_detection_time: whether to include detection_time, defaults to False :type use_detection_time: bool :param use_phi_1: whether to include phi_1, defaults to False :type use_phi_1: bool :param use_phi_2: whether to include phi_2, defaults to False :type use_phi_2: bool :param use_phi_12: whether to include phi_12, defaults to False :type use_phi_12: bool :param use_phi_orb: whether to include phi_orb, defaults to False :type use_phi_orb: bool :param validate_args: whether to validate arguments, defaults to None :type validate_args: Optional[bool], optional :returns: scaled mixture of distributions :rtype: ScaledMixture .. py:function:: 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 .. py:function:: TwoTruncatedNormalMixture(comp1_high: Optional[jaxtyping.ArrayLike], comp1_loc: jaxtyping.ArrayLike, comp1_low: Optional[jaxtyping.ArrayLike], comp1_scale: jaxtyping.ArrayLike, comp2_high: Optional[jaxtyping.ArrayLike], comp2_loc: jaxtyping.ArrayLike, comp2_low: Optional[jaxtyping.ArrayLike], comp2_scale: jaxtyping.ArrayLike, zeta: jaxtyping.ArrayLike, *, validate_args: Optional[bool] = None) -> numpyro.distributions.MixtureGeneral Create a mixture model of two truncated normal distributions. :param comp1_high: Upper truncation for the first component. :type comp1_high: Optional[ArrayLike] :param comp1_loc: Location parameter for the first component. :type comp1_loc: ArrayLike :param comp1_low: Lower truncation for the first component. :type comp1_low: Optional[ArrayLike] :param comp1_scale: Scale parameter for the first component. :type comp1_scale: ArrayLike :param comp2_high: Upper truncation for the second component. :type comp2_high: Optional[ArrayLike] :param comp2_loc: Location parameter for the second component. :type comp2_loc: ArrayLike :param comp2_low: Lower truncation for the second component. :type comp2_low: Optional[ArrayLike] :param comp2_scale: Scale parameter for the second component. :type comp2_scale: ArrayLike :param zeta: Mixing proportion for the second component. :type zeta: ArrayLike :param validate_args: Whether to validate the arguments, by default None :type validate_args: Optional[bool], optional :returns: A mixture of two truncated normal distributions. :rtype: MixtureGeneral