GaussianProcess

class pmrf.discrepancy_models.GaussianProcess(kernel: Callable[[Array, Array], Array], jitter: float = 1e-10, *, name: str | None = None)

Bases: DiscrepancyModel

Gaussian process discrepancy model with a covariance kernel.

Maps model predictions to a Gaussian Process distribution over frequency.

The kernel is responsible for returning the correlation between two input points. Given an input y of shape (*batch_shape, event_dims), the kernel must accept two inputs (x1, x2) of the same shape (scalar or vector), and return an array that is broadcastable to *batch_shape.

This easily allows for kernel batching. For example, to create multiple RBF kernels that model the last batch dimension D with independent kernels, simply create a kernel with parameters of shape (D,).

See pmrf.DiscrepancyModel for more information on general discrepancy models.

__call__(y_event: Array, x: Array) → AbstractDistribution

Evaluate the Gaussian process distribution over the given inputs.

Parameters:

• y_event (jnp.ndarray) – The model prediction in event space, with shape (…, N).

• x (jnp.ndarray) – The frequency points, with shape (N,).

Returns:

A multivariate normal distribution parameterized by the mean y_event and the covariance matrix generated by the kernel.

Return type:

dist.AbstractDistribution

jitter: float = 1e-10

A small scalar added to the diagonal of the covariance matrix for numerical stability.

kernel: Callable[[Array, Array], Array]

The covariance function that computes the correlation between two input arrays.