gp(integrations): add GPflow, LogFalkon, SPDE-FEM, metrics, and advanced model-pattern docs
- 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
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