gwkokab.inference
=================

.. py:module:: gwkokab.inference

.. autoapi-nested-parse::

   Provides essential classes and functions for the inference module.



Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/gwkokab/inference/factory/index
   /autoapi/gwkokab/inference/flowMC_analytical_poisson_likelihood/index
   /autoapi/gwkokab/inference/flowMC_discrete_poisson_likelihood/index
   /autoapi/gwkokab/inference/numpyro_analytical_poisson_likelihood/index
   /autoapi/gwkokab/inference/numpyro_discrete_poisson_likelihood/index
   /autoapi/gwkokab/inference/poissonlikelihood_utils/index


Functions
---------

.. autoapisummary::

   gwkokab.inference.analytical_poisson_likelihood_fn
   gwkokab.inference.discrete_poisson_likelihood_fn
   gwkokab.inference.flowMC_analytical_poisson_likelihood
   gwkokab.inference.flowMC_discrete_poisson_likelihood
   gwkokab.inference.numpyro_analytical_poisson_likelihood
   gwkokab.inference.numpyro_discrete_poisson_likelihood


Package Contents
----------------

.. py:function:: analytical_poisson_likelihood_fn(model_instance: numpyro.distributions.distribution.Distribution, poisson_mean_estimator: collections.abc.Callable[Ellipsis, tuple[jaxtyping.Array, jaxtyping.Array]], samples_stack: jaxtyping.Array, ln_offsets: jaxtyping.Array, pmean_kwargs: Dict[str, Any], variance_cut_threshold: float | None) -> jaxtyping.Array

.. py:function:: discrete_poisson_likelihood_fn(model_instance: numpyro.distributions.distribution.Distribution, poisson_mean_estimator: collections.abc.Callable[Ellipsis, Tuple[jaxtyping.Array, jaxtyping.Array]], data_group: Tuple[jaxtyping.Array, Ellipsis], log_ref_priors_group: Tuple[jaxtyping.Array, Ellipsis], masks_group: Tuple[jaxtyping.Array, Ellipsis], pmean_kwargs: Dict[str, Any], N_pes: Tuple[jaxtyping.Array, Ellipsis], variance_cut_threshold: float | None) -> jaxtyping.Array

.. py:function:: flowMC_analytical_poisson_likelihood(dist_fn: Callable[Ellipsis, numpyro.distributions.Distribution], priors: gwkokab.models.utils.JointDistribution, variables: Dict[str, numpyro.distributions.Distribution], constant_params: Dict[str, Any], variables_index: Dict[str, int], poisson_mean_estimator: Callable[[gwkokab.models.utils.ScaledMixture], tuple[jaxtyping.Array, jaxtyping.Array]], variance_cut_threshold: float | None) -> Callable[[jaxtyping.Array, Dict[str, Any]], jaxtyping.Array]

.. py:function:: flowMC_discrete_poisson_likelihood(dist_fn: collections.abc.Callable[Ellipsis, numpyro.distributions.distribution.Distribution], priors: gwkokab.models.utils.JointDistribution, variables: Dict[str, numpyro.distributions.distribution.Distribution], variables_index: Dict[str, int], poisson_mean_estimator: collections.abc.Callable[[gwkokab.models.utils.ScaledMixture], jax.Array], where_fns: Optional[List[collections.abc.Callable[Ellipsis, jax.Array]]], constants: Dict[str, jax.Array], variance_cut_threshold: float | None) -> collections.abc.Callable[[jax.Array, Dict[str, Any]], jax.Array]

   This class is used to provide a likelihood function for the inhomogeneous Poisson
   process. The likelihood is given by,

   .. math::
       \log\mathcal{L}(\Lambda) \propto -\mu(\Lambda)
       +\log\sum_{n=1}^N \int \ell_n(\lambda) \rho(\lambda\mid\Lambda)
       \mathrm{d}\lambda


   where, :math:`\displaystyle\rho(\lambda\mid\Lambda) =
   \frac{\mathrm{d}N}{\mathrm{d}V\mathrm{d}t \mathrm{d}\lambda}` is the merger
   rate density for a population parameterized by :math:`\Lambda`, :math:`\mu(\Lambda)` is
   the expected number of detected mergers for that population, and
   :math:`\ell_n(\lambda)` is the likelihood for the :math:`n`-th observed event's
   parameters. Using Bayes' theorem, we can obtain the posterior
   :math:`p(\Lambda\mid\text{data})` by multiplying the likelihood by a prior
   :math:`\pi(\Lambda)`.

   .. math::
       p(\Lambda\mid\text{data}) \propto \pi(\Lambda) \mathcal{L}(\Lambda)

   The integral inside the main likelihood expression is then evaluated via
   Monte Carlo as

   .. math::
       \int \ell_n(\lambda) \rho(\lambda\mid\Lambda) \mathrm{d}\lambda \propto
       \int \frac{p(\lambda | \mathrm{data}_n)}{\pi_n(\lambda)}
       \rho(\lambda\mid\Lambda) \mathrm{d}\lambda \approx
       \frac{1}{N_{\mathrm{samples}}}
       \sum_{i=1}^{N_{\mathrm{samples}}}
       \frac{\rho(\lambda_{n,i}\mid\Lambda)}{\pi_{n,i}}


.. py:function:: numpyro_analytical_poisson_likelihood(dist_fn: collections.abc.Callable[Ellipsis, numpyro.distributions.Distribution], priors: gwkokab.models.utils.JointDistribution, variables: Dict[str, numpyro.distributions.Distribution], constant_params: Dict[str, Any], variables_index: Dict[str, int], poisson_mean_estimator: collections.abc.Callable[[gwkokab.models.utils.ScaledMixture], tuple[jaxtyping.Array, jaxtyping.Array]], variance_cut_threshold: float | None) -> collections.abc.Callable[[jaxtyping.Array, jaxtyping.Array, Dict[str, Any]], None]

.. py:function:: numpyro_discrete_poisson_likelihood(dist_fn: collections.abc.Callable[Ellipsis, numpyro.distributions.Distribution], priors: gwkokab.models.utils.JointDistribution, variables: Dict[str, numpyro.distributions.Distribution], variables_index: Dict[str, int], poisson_mean_estimator: collections.abc.Callable[[gwkokab.models.utils.ScaledMixture], jax.Array], where_fns: Optional[List[collections.abc.Callable[Ellipsis, jax.Array]]], constants: Dict[str, jax.Array], variance_cut_threshold: float | None) -> collections.abc.Callable[[Tuple[jax.Array, Ellipsis], Tuple[jax.Array, Ellipsis], Tuple[jax.Array, Ellipsis]], jax.Array]