GaussianProcess

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

Bases: AbstractDiscrepancyModel

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. See pmrf.stats.covariance_kernels for built-in covariance kernels.

Parameters:

kernel (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]) – The covariance kernel function that computes the correlation between two input arrays. Can be a function or a callable PyTree. See pmrf.stats.covariance_kernels for built-in covariance kernels.
jitter (float, default=1e-10) – A small scalar added to the diagonal of the covariance matrix for numerical stability.

__call__(y_event: Array, x: Array, orthogonal_projection: Array | None = None) → 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,).
orthogonal_projection (jnp.ndarray, optional) – An optional matrix P of shape (..., N, N) which the kernel matrix is projected onto using P @ K @ P^T. This can be used to specify the subspace which the kernel is allowed.

Returns:

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

Return type:

dist.AbstractDistribution

See also

pmrf.stats.covariance_kernels.gram: Builds the covariance matrix used here.

log_prob(y_event: Array, observed: Array, x: Array, noise_variance: Array) → Array

Evaluate the summed log density of observed under the GP plus Gaussian noise.

Each batch entry of observed is distributed as \(\mathcal{N}(y, K + \Sigma_n)\), with \(\Sigma_n = \mathrm{diag}(\sigma^2)\). Equal to the log probability of the distribution built by __call__() and pmrf.stats.likelihoods.GaussianLikelihood, summed over the batch.

M = K + Sigma_n is formed and factorized at the broadcast batch shape of the kernel’s Gram batch and noise_variance, rather than the full batch shape, and the residuals sharing each M are solved together.

Parameters:

y_event (jnp.ndarray) – The model prediction in event space, with shape (*batch_shape, N).
observed (jnp.ndarray) – The observation in event space, broadcastable to y_event.
x (jnp.ndarray) – The frequency points, with shape (N,).
noise_variance (jnp.ndarray) – The noise variance \(\sigma^2\), with the event as its last axis and broadcastable to (*batch_shape, N). A scalar, or a last axis of size one, is constant along the event axis.

Returns:

The scalar log density, summed over the batch.

Return type:

jnp.ndarray

orthogonal_log_prob(y_event: Array, observed: Array, x: Array, noise_variance: Array, basis) → Array

Evaluate the full-data P K P^T + sigma^2 I Gaussian density.

This uses its nonsingular block factorization; it is not REML. In particular, the tangent-space block is retained because it depends on the fitted mean and measurement noise.

Unlike log_prob(), noise_variance must be constant along the event axis, so it has no event axis: it broadcasts to the event batch shape.

predict(residual: Array, x: Array, x_new: Array, noise_variance: Array) → AbstractDistribution

Predict the discrepancy at new frequencies, given residuals at the fit frequencies.

Conditions the GP on the residuals \(r\) at the fit frequencies \(x_A\) with Gaussian noise covariance \(\Sigma_n = \mathrm{diag}(\sigma^2)\), and returns the distribution of the discrepancy \(\delta\) at the new frequencies \(x_B\) for every event block:

\[\mu = K_{BA} (K_{AA} + \Sigma_n)^{-1} r\]

\[\Sigma = K_{BB} - K_{BA} (K_{AA} + \Sigma_n)^{-1} K_{AB}\]

The prediction is of \(\delta\) alone: it excludes measurement noise. The GP’s jitter is added to \(K_{AA}\) and \(K_{BB}\), so \(\Sigma\) stays positive definite.

As in log_prob(), K_AA + Sigma_n is formed and factorized at the broadcast shape of the kernel’s Gram batch and noise_variance, and the residuals sharing each matrix are solved together.

Parameters:

residual (jnp.ndarray) – The residuals at the fit frequencies in event space, with shape (*batch_shape, N_A).
x (jnp.ndarray) – The fit frequency points, with shape (N_A,).
x_new (jnp.ndarray) – The frequency points to predict at, with shape (N_B,).
noise_variance (jnp.ndarray) – The noise variance \(\sigma^2\), with the event as its last axis and broadcastable to (*batch_shape, N_A). A scalar, or a last axis of size one, is constant along the event axis.

Returns:

A multivariate normal distribution over the discrepancy at x_new, batched over the event blocks, with event shape (N_B,).

Return type:

dist.AbstractDistribution

jitter: float = 1e-10: The added jitter.

kernel: Callable[[Array, Array], Array]: The covariance kernel.