gwkokab.inference

Provides essential classes and functions for the inference module.

Submodules

• gwkokab.inference.factory
• gwkokab.inference.flowMC_analytical_poisson_likelihood
• gwkokab.inference.flowMC_discrete_poisson_likelihood
• gwkokab.inference.numpyro_analytical_poisson_likelihood
• gwkokab.inference.numpyro_discrete_poisson_likelihood
• gwkokab.inference.poissonlikelihood_utils

Functions

analytical_poisson_likelihood_fn(→ jaxtyping.Array)
discrete_poisson_likelihood_fn(→ jaxtyping.Array)
flowMC_analytical_poisson_likelihood(...)
flowMC_discrete_poisson_likelihood(...)	This class is used to provide a likelihood function for the inhomogeneous Poisson
numpyro_analytical_poisson_likelihood(...)
numpyro_discrete_poisson_likelihood(...)

Package Contents

gwkokab.inference.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

gwkokab.inference.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

gwkokab.inference.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]

gwkokab.inference.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: List[collections.abc.Callable[Ellipsis, jax.Array]] | None, 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,

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

where, \(\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 \(\Lambda\), \(\mu(\Lambda)\) is the expected number of detected mergers for that population, and \(\ell_n(\lambda)\) is the likelihood for the \(n\)-th observed event’s parameters. Using Bayes’ theorem, we can obtain the posterior \(p(\Lambda\mid\text{data})\) by multiplying the likelihood by a prior \(\pi(\Lambda)\).

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

\[\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}}\]

gwkokab.inference.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]

gwkokab.inference.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: List[collections.abc.Callable[Ellipsis, jax.Array]] | None, 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]