MarginalLogLikelihood

class pmrf.evaluators.MarginalLogLikelihood(predictor: Callable[[PyTree, Frequency], jnp.ndarray], observed: Any, likelihood: Callable[[jnp.ndarray | dist.AbstractDistribution], dist.AbstractDistribution], discrepancy: Callable[[jnp.ndarray, jnp.ndarray], dist.AbstractDistribution] | None = None, use_orthogonal_discrepancy: bool = False, orthogonal_basis: OrthogonalBasis | None = None, orthogonal_rcond: float | None = None, orthogonal_recompute: bool = False, event_transform: bij.AbstractBijector | Callable[[jnp.ndarray], bij.AbstractBijector] = None, event_ndims: int = 1, *, name: str | None = None, metadata: Any = None)

Bases: AbstractEvaluator

Computes the log of the probability of observed data by conditioning a likelihood function on a model prediction while marginalizing out a potential discrepancy model.

Includes a mapping from model “observation space” to “event space”. By default, this defines the frequency axis as the probabilistic event, by moving it to the last axis before passing it to the likelihood/discrepancy. However, an bijective transform can be applied to model probability in an arbitrary latent space.

Parameters:

predictor (Callable[[jaxtyping.PyTree, pmrf.frequency.Frequency], jax.jaxlib._jax.Array]) – The predictor (e.g. another Evaluator) that extracts model features. Can be a function or a PyTree with optional parameters.
observed (numpy.ndarray) – The observed data that the log probability will be computed of. Must have a shape that matches the shape of the predictor output.
likelihood (Callable[[jax.jaxlib._jax.Array | distreqx.distributions._distribution.AbstractDistribution], distreqx.distributions._distribution.AbstractDistribution]) – The likelihood function that takes the model prediction and returns the probability of observing some data. Can be a function or a PyTree with optional parameters. See pmrf.likelihoods for common likelihoods.
discrepancy (Callable[[jax.jaxlib._jax.Array, jax.jaxlib._jax.Array], distreqx.distributions._distribution.AbstractDistribution] | None) – An optional discrepancy model to cater for model misspecification. Can be a function or a PyTree with optional parameters. See pmrf.discrepancy_models for common discrepancy models.
use_orthogonal_discrepancy (bool) – Constrain a Gaussian-process discrepancy to the complement of the free-parameter tangent space. This retains the full-data likelihood; it is not REML.
orthogonal_basis (pmrf.evaluators.OrthogonalBasis | None) – A fixed tangent basis. Prefer creating it with with_orthogonal_reference() rather than constructing it directly.
orthogonal_rcond (float | None) – Required relative singular-value cutoff for the column-scaled Jacobian SVD. The cutoff is applied independently to each batch element.
orthogonal_recompute (bool) – If true, recompute the basis at every evaluation and differentiate through it. The default is false: a fixed reference is faster and defines a stable, normalized density. Configure it once with with_orthogonal_reference().
event_transform (distreqx.bijectors._bijector.AbstractBijector | Callable[[jax.jaxlib._jax.Array], distreqx.bijectors._bijector.AbstractBijector]) –
A bijective transform that maps from “observation space” (predicted features) to “event space” (probability). Can be:
- None, to use the default mapping (frequency as the event axis and independant real/imag);
- an distreqx.bijectors.AbstractBijector, applied as a fixed (“static”) transform to both the prediction and the observation;
- a callable taking the prediction in observation space and returning an distreqx.bijectors.AbstractBijector (a “conditional” transform). The returned bijector is resolved once per evaluation and applied to both the prediction and the observation, so the residual is expressed in the prediction’s own frame. This lets the residual basis depend on the prediction.
A conditional transform must be static in the PyTree sense: it is stored as a static field, so it should be a plain function or another hashable callable that holds no traced parameters. Any inferred parameters belong in the likelihood.

Normalization. A prediction-dependent change of variables introduces a Jacobian determinant. This class adds the log-determinant term properly: when the transform is conditional, jnp.sum(transform.forward_log_det_jacobian(observed)) of the resolved bijector is added to the returned log-likelihood. Conditional transforms are therefore not required to be volume preserving; for a volume-preserving transform (such as a unitary rotation) the term is exactly zero and contributes nothing.

The sum follows distreqx’s own convention, in that the bijector is trusted to report its log-determinant with one contribution per independent coordinate it transforms. A bijector parameterized by scalars may report a single scalar that distreqx does not broadcast over the input, in which case the summed term counts that contribution once rather than once per element.

The term is deliberately not added for a static (or default) transform, where it is a constant offset independent of the model. Omitting it keeps the default behaviour unchanged and shifts the log-likelihood only by a constant, which affects neither optimization nor posterior sampling.
event_ndims (int) – The number of trailing event dimensions in event space to use as the event shape. Defaults to 1.

__call__(model: PyTree, frequency: Frequency, **kwargs) → Array

Evaluate the model response over the specified frequency range.

Parameters:

model (PyTree) – The parameter PyTree to evaluate.
freq (Frequency) – The frequency object defining the evaluation points.
**kwargs (dict) – Additional keyword arguments for the evaluation process.

Returns:

The evaluated model response.

Return type:

jnp.ndarray

predictive_distribution(model: PyTree, frequency: Frequency, **kwargs) → AbstractDistribution

Returns the full predictive distribution of an observed event for a given model.

The returned distribution is in event space. To draw a sample from this distribution in observation space, see MarginalLogLikelihood.sample_observation().

resolve_event_transform(y_pred: Array) → AbstractBijector

Resolve event_transform to a concrete bijector for a given prediction.

A static transform is returned unchanged. A conditional transform is called with the prediction in observation space and must return an distreqx.bijectors.AbstractBijector.

Parameters:: y_pred (jnp.ndarray) – The model prediction in observation space.
Returns:: The resolved transform, to be applied to both prediction and observation.
Return type:: bij.AbstractBijector

sample_observation(key: Array, model: PyTree, frequency: Frequency, **kwargs) → Array: Returns a sample from the predictive distribution in observation space.

with_orthogonal_reference(model: PyTree, frequency: Frequency, **kwargs) → MarginalLogLikelihood: Return a copy with its tangent basis fixed at model and frequency.

discrepancy: Callable[[Array, Array], AbstractDistribution] | None = None: The optional discrepancy model.

event_ndims: int = 1: The number of trailing event dimensions.

event_transform: AbstractBijector | Callable[[Array], AbstractBijector] = None: The bijective event transform, or a callable resolving one from the prediction.

property has_conditional_event_transform: bool

Whether event_transform is prediction-dependent (“conditional”).

True when event_transform is a callable that resolves to a bijector from the prediction, rather than being a bijector itself.

likelihood: Callable[[Array | AbstractDistribution], AbstractDistribution]: The active likelihood function.

observed: ndarray: The observed data.

orthogonal_basis: OrthogonalBasis | None = None: Fixed tangent basis, normally populated by with_orthogonal_reference().

orthogonal_rcond: float | None = None: Relative singular-value cutoff used to determine tangent rank.

orthogonal_recompute: bool = False: Recompute the tangent basis at every call, including its exact derivative.

predictor: Callable[[PyTree, Frequency], Array]: The active predictor instance.

use_orthogonal_discrepancy: bool = False: Flag for orthogonal discrepancy.