gwkokab.inference.numpyro_discrete_poisson_likelihood¶

NumPyro model over per-event posterior samples.

Where the flowMC modules return a plain function of a flat vector, NumPyro needs a model: each hyper-parameter is drawn at its own named sample site and the likelihood enters as a factor. The site names are the hyper-parameter names sorted alphabetically, which is the same ordering variables_index records.

When priors is a LazyJointDistribution, the variables are sampled in two passes: the independent ones first, then the lazy ones in topological order, each built from the values it depends on. NumPyro needs this because every variable must be an individually named sample site.

Functions¶

numpyro_discrete_poisson_likelihood(...)

Build the NumPyro model for the discrete data representation.

Module Contents¶

gwkokab.inference.numpyro_discrete_poisson_likelihood.numpyro_discrete_poisson_likelihood(dist_fn: collections.abc.Callable[..., 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[..., jax.Array]] | None, constants: Dict[str, jax.Array], variance_cut_threshold: float | None) → collections.abc.Callable[[Tuple[jax.Array, ...], Tuple[jax.Array, ...], Tuple[jax.Array, ...]], jax.Array][source]¶

Build the NumPyro model for the discrete data representation.

The likelihood is that of the inhomogeneous Poisson process,

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

where \(\rho(\lambda\mid\Lambda) = \mathrm{d}N/\mathrm{d}V\,\mathrm{d}t\,\mathrm{d}\lambda\) is the merger rate density of a population parameterised by \(\Lambda\), \(\mu(\Lambda)\) is the expected number of detections for that population, and \(\ell_n(\lambda)\) is the likelihood for the \(n\)-th observed event’s parameters.

with the per-event integral evaluated as a Monte Carlo average over that event’s posterior samples, reweighted by the PE prior they were drawn under.

Parameters:

  • dist_fn (Callable[..., Distribution]) – Builds the population model from the sampled variables and the constants.

  • priors (JointDistribution) – Joint prior over the sampled variables, ordered as variables_index says.

  • variables (Dict[str, Distribution]) – The sampled variables and their priors.

  • variables_index (Dict[str, int]) – Maps each variable name to its position in the flat parameter vector. This is the ordering recorded as variables_index in the output HDF5.

  • poisson_mean_estimator (Callable) – Estimator of \(\mu(\Lambda)\), returning (mean, variance). See gwkokab.poisson_mean.

  • where_fns (Optional[List[Callable[..., Array]]]) – Extra validity predicates on the hyper-parameters, beyond the prior support. Points failing any of them get \(-\infty\). None skips the check entirely.

  • constants (Dict[str, Array]) – Hyper-parameters frozen by the prior configuration rather than sampled.

  • variance_cut_threshold (float | None) – Threshold above which the Monte Carlo variance of the estimate is penalised by variance_tapering_fn(). None disables the penalty.

Returns:

A NumPyro model taking the bucketed event data and registering the likelihood as a factor site.

Return type:

Callable