jejjohnson / jejjohnson/pyrox

gp(integrations): add GPflow, LogFalkon, SPDE-FEM, metrics, and advanced model-pattern docs

Open
#50 0 comments 0 reactions 0 assignees View on GitHub
area:docs area:gp area:integration type:feature wave:5-advanced
Dominant language
Python
Stars
1
Forks
0
Avg merge
19h 2m
Merged PRs (30d)
16

Description

## Problem / Request

The `pyrox.gp` surface today fits well-defined model patterns (dense / sparse / Markov / multi-output) but downstream evaluation and cross-ecosystem comparison are open-coded by each user. There's no `pyrox.gp.metrics` module for the standard calibration / scoring metrics (NLL, CRPS, log-marginal vs. test-NLL gap, reliability diagrams), no GPflow-compatibility shim for users migrating models, no LogFalkon backend, and no SPDE-FEM precision-matrix builder for spatial GMRFs.

This issue tracks the **integrations + metrics** scope as a single thread because the metrics module is the shared dependency for the GPflow / Falkon / SPDE bring-ups (each integration's correctness has to be verified against the metrics module).

## User Story

> As a user porting a GPflow regression model to `pyrox.gp` (or comparing a Falkon-style large-scale KRR to `pyrox.gp.markov` on the same dataset), I want one `pyrox.gp.metrics` module + a small set of integration shims so I can score, calibrate, and benchmark without rewriting the test harness for each framework.

## Motivation

Concrete use cases that need this thread:

- **GPflow → pyrox migration**: users with existing GPflow models need a one-line conversion path for kernels and likelihoods + the same evaluation suite to compare results.
- **Large-scale KRR benchmarking**: pyrox.gp.markov and Falkon target different regimes; reviewers need an apples-to-apples NLL / CRPS / RMSE table.
- **Spatial GMRF prior building**: SPDE-FEM is the standard way to build a Matérn precision matrix from a triangulated mesh (Lindgren–Rue 2011). Currently users have to build the precision matrix by hand and pass it to `MultivariateNormalPrecision` from gaussx.
- **Reproducible benchmarks** in the eventual paper / docs require all of the above on a shared test split.

Scope is wide; keep each integration narrow with explicit "supported / not supported" tables in the docstring.

## Mathematics

**Metrics (probabilistic regression).** For a held-out point with true value $y$ and predicted $\mathcal{N}(\mu, \sigma^2)$,

$$
\text{NLL}(y; \mu, \sigma) \;=\; \tfrac{1}{2}\log(2\pi\sigma^2) + \tfrac{(y - \mu)^2}{2 \sigma^2},
$$

$$
\text{CRPS}(y; \mu, \sigma) \;=\; \sigma\,\Big[\phi\!\big(\tfrac{y - \mu}{\sigma}\big) + \tfrac{y - \mu}{\sigma}\big(2\Phi\!\big(\tfrac{y - \mu}{\sigma}\big) - 1\big) - \tfrac{1}{\sqrt\pi}\Big].
$$

Reliability diagram: empirical coverage vs. nominal coverage on a grid of quantiles.

**SPDE-FEM Matérn precision.** For $u \sim \mathcal{N}(0, K_{\nu, \kappa})$ on a domain $\Omega$, the SPDE form $(\kappa^2 - \Delta)^{\alpha/2} u = \mathcal{W}$ with $\alpha = \nu + d/2$ has a sparse FEM precision

$$
Q \;=\; \tau^2 \, \bigl(\kappa^4 C + 2 \kappa^2 G + G C^{-1} G\bigr) \quad \text{for } \alpha = 2,
$$

where $C$ is the mass matrix and $G$ the stiffness matrix from a triangulated mesh. The precision is sparse with the same sparsity pattern as $G$.

**Falkon KRR (out-of-scope but referenced).** Inducing-point KRR with $O(N M^2)$ training cost; pyrox already plans gaussx#49 (Nyström preconditioner). This issue is the **comparison shim**, not a re-implementation.

## Proposed API

```python
# pyrox.gp.metrics — scoring
def regression_nll(y_true, q_pred: dist.Normal | dist.StudentT) -> Float[Array, ""]: ...
def regression_crps(y_true, q_pred) -> Float[Array, ""]: ...
def classification_nll(y_true, p_pred) -> Float[Array, ""]: ...
def reliability_curve(y_true, q_pred, n_quantiles=20) -> tuple[Array, Array]: ...

# pyrox.gp.integrations — narrow shims
def from_gpflow_kernel(gpflow_kernel) -> pyrox.gp.Kernel: ...
def from_gpflow_likelihood(gpflow_lik) -> pyrox.gp.Likelihood: ...
def spde_fem_matern_precision(
mesh_vertices: Float[Array, "V d"],
mesh_faces: Int[Array, "F 3"],
*,
kappa: float,
nu: int, # alpha = nu + d/2, alpha integer
) -> lx.AbstractLinearOperator: ... # sparse block-tridiag-like Q

def falkon_backend(model, *, n_inducing: int, **falkon_kwargs):
"""Forward to a Falkon (LogFalkon) backend; experimental."""
```

Each shim ships with a one-line "Supported / Not supported" table in its docstring.

## Example usage

```python
import jax.numpy as jnp, pyrox.gp as pgp

# Metrics — same shape regardless of model class
y_true = ... # (N,)
q_pred = model.predict(X_test) # numpyro.distributions.Normal
nll = pgp.metrics.regression_nll(y_true, q_pred)
crps = pgp.metrics.regression_crps(y_true, q_pred)
nominal, coverage = pgp.metrics.reliability_curve(y_true, q_pred)

# SPDE-FEM Matern prior on a triangulated mesh
mesh_v, mesh_f = build_2d_mesh(...)
Q = pgp.integrations.spde_fem_matern_precision(
mesh_v, mesh_f, kappa=1.0, nu=1, # alpha = 2 (Matern-3/2 in 2D)
)
# Q is a sparse precision; hand it to MultivariateNormalPrecision (gaussx)
```

## Tasks / sub-tasks

- [ ] **Spike** — pick the sparse-matrix backend for SPDE precision (likely `lineax.MatrixLinearOperator` over a `jax.experimental.sparse` BCOO matrix, or a dense fallback for small meshes).
- [ ] **`pyrox.gp.metrics`** — `src/pyrox/gp/_metrics.py`. NLL, CRPS, classification NLL, reliability curve. Pure JAX — no per-backend special cases.
- [ ] **`pyrox.gp.integrations.gpflow`** — narrow kernel + likelihood converters. Document exactly which GPflow classes are supported in the docstring.
- [ ] **`pyrox.gp.integrations.spde`** — `spde_fem_matern_precision` for $\alpha \in \{1, 2\}$ in 2D. Document the assumption that the mesh comes from an external triangulator (`scipy.spatial.Delaunay`, `meshzoo`, etc.).
- [ ] **`pyrox.gp.integrations.falkon`** — optional thin wrapper on `FalkonML/falkon`; behind a `pip install pyrox[falkon]` extra. Document as experimental.
- [ ] **Tests** — (i) NLL / CRPS match `scipy.stats.norm.logpdf` / closed-form CRPS on small problems; (ii) SPDE-FEM Matérn precision: empirical samples match dense Matérn covariance at small mesh sizes; (iii) GPflow kernel converter: `pyrox` and GPflow `K(X, X)` agree to `1e-6` on a 50-point toy.
- [ ] **Benchmark** — Falkon vs. `pyrox.gp.markov` on a $10^5$-point regression: training time, test NLL, test CRPS.
- [ ] **Docs** — `pyrox.gp.metrics` reference + `pyrox.gp.integrations` reference; one "model-pattern doc" page comparing standard pyrox GP models on a shared benchmark.

## References

- Lindgren, F., Rue, H. & Lindström, J. (2011). *An explicit link between Gaussian fields and Gaussian Markov random fields: the SPDE approach*. JRSS-B (SPDE-FEM construction).
- Gneiting, T. & Raftery, A. E. (2007). *Strictly proper scoring rules, prediction, and estimation*. JASA (CRPS / NLL / reliability).
- Meanti, G. *et al.* (2020). *Kernel methods through the roof: handling billions of points efficiently*. NeurIPS (Falkon).
- GPflow user guide: .

## Relationships

- Parent (wave epic): #43
- Blocked by: #
- Blocks: #
- Related: gaussx#49 (Falkon recipe), #48

Contributor guide

Open the contributing guide

Research direction

Start with src/pyrox/gp/_metrics.py and the proposed pyrox.gp.integrations.gpflow, spde, and falkon modules. Review the supported and unsupported API scope, then run the specified metric, SPDE-FEM, and GPflow comparison tests. Done means the requested shims, metrics, tests, benchmark, and reference documentation are covered.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
documentation, machine-learning, testing-qa
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.