MarginalLogLikelihood

class pmrf.objectives.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.stats.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.stats.discrepancy_models for common discrepancy models. A GaussianProcess with a GaussianLikelihood is evaluated in closed form (log_prob()), which factorizes each distinct covariance matrix once. The noise may be constant or vary along the event axis.
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.objectives.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

linearize(model: PyTree, frequency: Frequency, space: Literal['raw', 'declared', 'physical'] = 'declared', **kwargs) → Linearization

Linearise the fit at model.

The residual $r = \tilde h - h(\theta)$ is formed in event space as in predict_discrepancy(), and differentiated with respect to the model’s free parameters, giving $J = -\partial r / \partial \theta$. With a conditional event transform the observation’s event depends on the model too, and $J$ includes that dependence. Only the model is differentiated: the kernel’s hyperparameters and the noise are held at their values in this evaluator.

$\Sigma_D = K + \Sigma_n$ is factorized per event block, with the noise variance read from the GaussianLikelihood at model’s prediction. It may vary along frequency. Without a discrepancy, $\Sigma_D = \Sigma_n$.

Parameters:

model (PyTree) – The fitted model. Must still be wrapped.
frequency (Frequency) – The frequency the evaluator was fitted over, matching observed.
space ({'declared', 'physical', 'raw'}, default='declared') – The space of the parameter values differentiated with respect to, as for pmrf.utils.derivative().
**kwargs – Passed to the predictor.

Returns:

The Jacobian, residual, factor of $\Sigma_D$ and Fisher matrix. Pass it to pmrf.stats.linearization.posterior_covariance().

Return type:

pmrf.stats.linearization.Linearization

Raises:

TypeError – If the discrepancy is neither None nor a GaussianProcess, or the likelihood is not a GaussianLikelihood.
ValueError – If orthogonal discrepancy is enabled.

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

Predict the discrepancy at new frequencies, conditioned on the fit.

The residual $r = \tilde h - h(\theta)$ of model at the fit frequency is formed in event space, with the resolved event transform applied to both the prediction and the observation, as in __call__(). The noise variance is read from the GaussianLikelihood as in the closed-form log-likelihood, and may vary along frequency. Both are passed to predict().

Kernel length scales are in the fit frequency’s unit, so new_frequency is converted to that unit first.

Parameters:

model (PyTree) – The fitted model.
frequency (Frequency) – The frequency the evaluator was fitted over, matching observed.
new_frequency (Frequency) – The frequencies to predict the discrepancy at, in any unit.
**kwargs – Passed to the predictor.

Returns:

The distribution of the discrepancy $\delta$ at new_frequency in event space, batched over the event blocks, with event shape (len(new_frequency),). It excludes measurement noise.

Return type:

dist.AbstractDistribution

Raises:

TypeError – If the discrepancy is not a GaussianProcess, or the likelihood is not a GaussianLikelihood.
ValueError – If orthogonal discrepancy is enabled.

predict_joint(model: PyTree, frequency: Frequency, new_frequency: Frequency, covariance: Array | None = None, **kwargs) → tuple[tuple[str, ...], AbstractDistribution]

Predict the free parameters and discrepancy together as a joint prior.

Mathematical Formulation

At the fitted model, let $A = K_{*h}\Sigma_D^{-1}$ and let $\Sigma_\text{post}$ be the linearised parameter covariance. Then

$$hatdelta = A r, qquad Sigma_delta = K_{**} - A K_{*h}^top

A JSigma_text{post}J^top A^top$$

\[\Sigma_{\delta\theta} = -A J\Sigma_\text{post}\]

The parameter term couples event blocks, giving one dense Gaussian over $[\theta;\delta]$. Kernel hyperparameters are held fixed, as in linearize(). New frequencies are converted to the fit’s unit.

If the Gaussian support leaves a selected parameter’s validity, attach it with prf.prior(..., truncate='unnormalised'). This retains the supplied density inside the valid region and gives values outside that region zero density, without adding a truncation normaliser. Such a prior is suitable for MAP fits, scoring, differentiation and posterior covariance; every sampling path rejects it before invoking a sampler.

Parameters:

model (PyTree) – The fitted model at its MAP, still carrying its parameters and prior.
frequency (Frequency) – The fit frequencies, matching observed.
new_frequency (Frequency) – The frequencies to predict the discrepancy at, in any unit.
covariance (jnp.ndarray or None, default=None) – Parameter covariance in declared space and the linearisation’s name order. If omitted, computed from this fit and the model’s prior with pmrf.stats.linearization.posterior_covariance(). Supply a covariance from combined linearisations when several fits inform the parameters.
**kwargs – Passed to the predictor.

Returns:

names (tuple of str) – The model’s free parameter names in linearisation order.
distribution (dist.AbstractDistribution) – One multivariate normal over the declared parameter values, flattened per name, followed by discrepancy values of shape (*batch, N_*) flattened in C order. Attach it with prf.prior(transfer, [*names, 'cable.discrepancy.values'], distribution) when its support lies inside each validity, or add truncate='unnormalised' when that support crosses a validity.

Raises:

TypeError – If the discrepancy is not a Gaussian process or the likelihood is not Gaussian.
ValueError – If orthogonal discrepancy is enabled or covariance has the wrong shape for the free parameters.

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.