gwkokab.inference

Provides essential classes and functions for the inference module.

Submodules

• gwkokab.inference.factory
• gwkokab.inference.flowMC_analytical_gwalk_poisson_likelihood
• gwkokab.inference.flowMC_discrete_poisson_likelihood
• gwkokab.inference.numpyro_analytical_gwalk_poisson_likelihood
• gwkokab.inference.numpyro_discrete_poisson_likelihood
• gwkokab.inference.poissonlikelihood_utils

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.
flowMC_analytical_gwalk_poisson_likelihood(...)	Build the flowMC log-posterior for the analytical GWalk data representation.
flowMC_discrete_poisson_likelihood(...)	Build the flowMC log-posterior for the discrete data representation.
numpyro_analytical_gwalk_poisson_likelihood(...)	Build the NumPyro model for the analytical GWalk data representation.
numpyro_discrete_poisson_likelihood(...)	Build the NumPyro model for the discrete data representation.

Package Contents

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

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

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

gwkokab.inference.flowMC_analytical_gwalk_poisson_likelihood(dist_fn: Callable[..., 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]

Build the flowMC log-posterior for the analytical GWalk 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 over points resampled from each event’s Gaussian summary.

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]) – Unused here; flowMC takes its priors from priors as a joint distribution rather than site by site.

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

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

• 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 log_posterior(x, data) callable. Non-finite values are mapped to \(-\infty\), so the sampler rejects such points rather than propagating NaN.

Return type:

Callable[[Array, Dict[str, Any]], Array]

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

Build the flowMC log-posterior 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.

Using Bayes’ theorem, the posterior \(p(\Lambda\mid\text{data})\) follows by multiplying by a prior \(\pi(\Lambda)\):

\[p(\Lambda\mid\text{data}) \propto \pi(\Lambda)\,\mathcal{L}(\Lambda)\]

The integral inside the likelihood is evaluated by Monte Carlo as

\[\int \ell_n(\lambda)\,\rho(\lambda\mid\Lambda)\,\mathrm{d}\lambda \propto \int \frac{p(\lambda \mid \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}}\]

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]) – Unused here; flowMC takes its priors from priors as a joint distribution rather than site by site.

• 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 JIT-compiled log_posterior(x, data). Non-finite values are mapped to \(-\infty\), so the sampler rejects such points rather than propagating NaN.

Return type:

Callable[[Array, Dict[str, Any]], Array]

gwkokab.inference.numpyro_analytical_gwalk_poisson_likelihood(dist_fn: collections.abc.Callable[..., 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]

Build the NumPyro model for the analytical GWalk 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 over points resampled from each event’s Gaussian summary.

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.

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

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

• 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 event data and registering the likelihood as a factor site.

Return type:

Callable

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

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