gwkokab.inference.poissonlikelihood_utils¶

The inhomogeneous-Poisson log-likelihood.

This is where the actual likelihood arithmetic lives; the four modules alongside it only adapt these two functions to a sampler backend. The likelihood is

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

The two functions differ in how the per-event integral is evaluated: discrete_poisson_likelihood_fn() averages over that event’s posterior samples, reweighted by the PE prior they were drawn under, while analytical_gwalk_poisson_likelihood_fn() uses samples resampled from the event’s Gaussian summary and their log offsets.

Both optionally subtract a variance-tapering penalty, so that hyper-parameters where the Monte Carlo estimate is too noisy to trust are pushed away from rather than silently accepted.

Functions¶

analytical_gwalk_poisson_likelihood_fn(→ jaxtyping.Array)

Poisson log-likelihood from per-event Gaussian summaries.

discrete_poisson_likelihood_fn(→ jaxtyping.Array)

Poisson log-likelihood from per-event posterior samples.

Module Contents¶

gwkokab.inference.poissonlikelihood_utils.analytical_gwalk_poisson_likelihood_fn(model_instance: numpyro.distributions.distribution.Distribution, poisson_mean_estimator: collections.abc.Callable[..., 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[source]¶

Poisson log-likelihood from per-event Gaussian summaries.

Each event is summarised by a mean and covariance rather than by a sample cloud; the samples here have been resampled from that summary, and ln_offsets carries the accompanying log weight. Samples falling outside the model’s support are excluded from the reduction rather than contributing \(-\infty\).

Parameters:

  • model_instance (Distribution) – The population model at the current hyper-parameters.

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

  • samples_stack (Array) – Resampled event coordinates, of shape (n_events, n_samples, n_parameters).

  • ln_offsets (Array) – Log weight of each resampled point.

  • pmean_kwargs (Dict[str, Any]) – Extra arguments for the Poisson mean estimator; must include T_obs.

  • 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:

The scalar log-likelihood.

Return type:

Array

gwkokab.inference.poissonlikelihood_utils.discrete_poisson_likelihood_fn(model_instance: numpyro.distributions.distribution.Distribution, poisson_mean_estimator: collections.abc.Callable[..., Tuple[jaxtyping.Array, jaxtyping.Array]], data_group: Tuple[jaxtyping.Array, ...], log_ref_priors_group: Tuple[jaxtyping.Array, ...], masks_group: Tuple[jaxtyping.Array, ...], pmean_kwargs: Dict[str, Any], N_pes: Tuple[jaxtyping.Array, ...], variance_cut_threshold: float | None) → jaxtyping.Array[source]¶

Poisson log-likelihood from per-event posterior samples.

Each event’s integral is a Monte Carlo average over its posterior samples, divided by the PE prior they were drawn under:

\[\int \ell_n(\lambda)\,\rho(\lambda\mid\Lambda)\,\mathrm{d}\lambda \approx \frac{1}{M_n} \sum_{i=1}^{M_n} \frac{\rho(\lambda_{n,i}\mid\Lambda)}{\pi_{n,i}}\]

Events are processed in buckets of similar sample count – see gwkokab.analysis.utils.jenks – so each bucket is one padded, rectangular array. Padding is masked out twice over: the padded coordinates are replaced by a point known to be in the support before the model is evaluated, so no NaN is produced, and their contributions are then masked to \(-\infty\) so they add nothing to the sum.

Parameters:

  • model_instance (Distribution) – The population model at the current hyper-parameters.

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

  • data_group (Tuple[Array, ...]) – One padded array of posterior samples per bucket of events.

  • log_ref_priors_group (Tuple[Array, ...]) – Log PE prior \(\log\pi_{n,i}\) of each sample, bucketed to match.

  • masks_group (Tuple[Array, ...]) – Boolean masks marking the real samples in each bucket, as opposed to padding.

  • pmean_kwargs (Dict[str, Any]) – Extra arguments for the Poisson mean estimator; must include T_obs.

  • N_pes (Tuple[Array, ...]) – Number of real posterior samples \(M_n\) per event, bucketed to match.

  • 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:

The scalar log-likelihood.

Return type:

Array