gwkokab.analysis.utils.priors

Turning prior_cfg.json into priors, constants, aliases and lazy priors.

The prior configuration is keyed by regex rather than by parameter name; once match_all() has resolved those keys against the actual hyper-parameter names, get_processed_priors() interprets each value. There are four kinds:

• a dict with a "dist" field becomes an instantiated NumPyro or GWKokab prior;

• a number becomes a constant, frozen rather than sampled;

• a string naming another parameter becomes an alias, tying the two to one sampled dimension;

• a dict whose field values are strings becomes a lazy prior: a jax.tree_util.Partial plus a dependency map, so the prior’s own hyper-parameters are themselves sampled.

This module also supplies DirichletElement(), a prior that has no NumPyro equivalent: it expresses one coordinate of a simplex conditional on those before it, which is what lets a set of mixing fractions be sampled coordinate by coordinate as independent named sites.

Functions

DirichletElement(→ numpyro.distributions.Distribution)	Conditional Dirichlet Distribution.
get_processed_priors(→ dict)	Get the processed priors from a list of parameters.

Module Contents

gwkokab.analysis.utils.priors.DirichletElement(order: int, n_dimensions: int, validate_args: bool | None = None, **kwargs) → numpyro.distributions.Distribution

Conditional Dirichlet Distribution.

A DirichletElement is a distribution over the \(N\)-dimensional simplex, where \(N\) is the n_dimensions parameter. The \(n\)-th order Dirichlet distribution given by the order. The distribution is defined as,

\[\forall n \in \{0, N-2\}, p(x_n \mid x_0, \cdots, x_{n-1}) = \frac{(N-n-1)(1-\alpha_n-x_n)^{N-n-1}}{(1-\alpha_n)^{N-n-1}}, \qquad x_n \in [0, \alpha_n]\]

where \(\alpha_n = \sum_{i=0}^{n-1} \alpha_i\) and \(x_n\) is the order-th element of the simplex.

\[p(x_{N-1}\mid x_0, \cdots, x_{N-2}) = \frac{1}{1-\alpha_{N-1}}, \qquad x_{N-1} \in [0,\alpha_{N-1}]\]

Parameters:

• order (int) – order of the DirichletElement

• n_dimensions (int) – number of dimensions

• validate_args (Optional[bool], optional) – whether to validate the arguments, by default None

Returns:

DirichletElement distribution

Return type:

Distribution

Raises:

KeyError – Missing concentration parameters for DirichletElement of order {order}

gwkokab.analysis.utils.priors.get_processed_priors(params: List[str], priors: dict) → dict

Get the processed priors from a list of parameters.

A processed prior is either an instantiated prior or a tuple of (jax.tree_util.Partial, lazy_vars) where lazy_vars is a dictionary of lazy variables.

Parameters:

• params (List[str]) – list of parameters

• priors (dict) – dictionary of priors

Returns:

dictionary of processed priors

Return type:

dict

Raises:

ValueError – if the prior value is invalid