gwkokab.models.mass =================== .. py:module:: gwkokab.models.mass Classes ------- .. autoapisummary:: gwkokab.models.mass.BrokenPowerlaw gwkokab.models.mass.BrokenPowerlawTwoPeak gwkokab.models.mass.GaussianPrimaryMassRatio gwkokab.models.mass.GenericSmoothedPowerlawMassRatio gwkokab.models.mass.PowerlawPrimaryMassRatio gwkokab.models.mass.SmoothedBrokenPowerlawMassRatioPowerlaw gwkokab.models.mass.SmoothedGaussianPrimaryMassRatio gwkokab.models.mass.SmoothedPowerlawPrimaryMassRatio gwkokab.models.mass.SmoothedTwoComponentPrimaryMassRatio gwkokab.models.mass.Wysocki2019MassModel 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:: GaussianPrimaryMassRatio(loc: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, mmin: jaxtyping.ArrayLike, mmax: 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,) :param loc: Location parameter for primary mass :type loc: ArrayLike :param scale: Scale parameter for primary mass :type scale: 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:: support() -> numpyro.distributions.constraints.Constraint The support of this distribution. Subclasses can override this as a class attribute or as a property. .. py:class:: GenericSmoothedPowerlawMassRatio(primary_mass_distribution: numpyro.distributions.Distribution, beta: jaxtyping.ArrayLike, delta_m1: jaxtyping.ArrayLike, delta_m2: jaxtyping.ArrayLike, m2min: 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:attribute:: args_constraints .. py:attribute:: primary_mass_distribution .. 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:: 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: 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:: SmoothedGaussianPrimaryMassRatio(loc, scale, beta, m1min, m2min, mmax, delta_m1, delta_m2, *, validate_args=None) Bases: :py:obj:`numpyro.distributions.Distribution` :class:`~numpyro.distributions.continuous.Normal` with smoothing kernel on the lower edge. .. math:: 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) .. math:: 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 :func:`~gwkokab.utils.kernel.log_planck_taper_window`. .. attention:: If :code:`low` or :code:`high` are not provided to the `TruncatedNormal`, they default to :math:`-\infty` or :math:`+\infty`, respectively. This class relies on this behavior to produce the desired distribution when bounds are unspecified. .. 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:: 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: :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:: 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.