feat(interpolate): Gaussian Process / kriging (gp_to_grid, fillnan_gp via sklearn)

# Design: Gaussian Process / kriging interpolation

**Status:** design draft (MVP)
**Author:** @jejjohnson
**Date:** 2026-05-07
**Backend:** `sklearn.gaussian_process.GaussianProcessRegressor` (new optional dep `xr_toolz[ml-interp]`)
**Module home:** new file [`src/xr_toolz/interpolate/_src/gp.py`](src/xr_toolz/interpolate/_src/gp.py) + operators in [`interpolate/operators.py`](src/xr_toolz/interpolate/operators.py)
**Related:** [#159](https://github.com/jejjohnson/xr_toolz/issues/159) (grid_to_points — different use case, no uncertainty), kNN/IDW issue (the cheap counterpart), [#158](https://github.com/jejjohnson/xr_toolz/issues/158) (dask)

---

## 1. User story

> As a user mapping **sparse observations** (along-track altimetry, sparse buoys, irregular profiles) onto a regular grid, I want **optimal interpolation** that gives me both a **smooth field** and a **per-pixel uncertainty estimate** — the same machinery the altimetry community calls "DUACS-style mapping" or "objective analysis."

For the dense-point case, kNN/IDW is fine. For sparse points where the structure of the field matters and uncertainty matters (downstream uses: assimilation, error budgets, "where can I trust this map"), GP / kriging is the right tool. No primitive in `xr_toolz` covers this today.

## 2. Why sklearn (and not GPy / GPflow / pykrige)

| Option | Verdict |
|---|---|
| **`sklearn.gaussian_process.GaussianProcessRegressor`** | **Pick.** Mature, sklearn-style API, broad kernel library (RBF, Matern, RationalQuadratic, etc.), built-in marginal-likelihood optimization, returns std-dev with `return_std=True`. |
| `pykrige` | Geostatistics-specific, exposes variogram models the geostats community expects (spherical, exponential, Gaussian variograms). Heavier dep. |
| `GPy`, `GPflow` | More flexible / scalable. Heavyweight (TF/PyTorch), excessive for the standard use case. |
| `dask-ml` GP | Doesn't exist as a first-class primitive yet. |
| Hand-rolled | The math is small but stable hyperparameter optimization isn't. |

Pick **sklearn**. Fits behind an optional `xr_toolz[ml-interp]` extra so users who don't need it don't carry it.

## 3. Scope (MVP)

In:
- **`gp_to_grid`** — scattered → regular grid with `mean` and optional `std` outputs.
- **`gp_to_points`** — scattered → arbitrary target points.
- **`fillnan_gp`** — gridded field gap-fill via GP from finite neighbours (slow but uncertainty-aware).
- **`GPToGrid`**, **`GPToPoints`**, **`FillNaNGP`** operators.
- **Kernel presets** keyed by string: `"matern"` (default ν=3/2), `"rbf"`, `"rational_quadratic"`. Plus pass-through for any sklearn kernel object.

Out (file as follow-ups):
- **Anisotropic kernels** (separate length-scales per axis) — supported via passing a custom sklearn kernel; presets ship isotropic.
- **Gaspari–Cohn / compactly-supported covariances** — important for assimilation, requires custom kernel; punt.
- **Sparse / inducing-point GPs** — for N > 10³ points where dense GP is intractable. Major project; separate issue.
- **Uncertainty quantification beyond posterior std** — credible intervals, sample paths.

## 4. Mathematics

For source $\{(\mathbf{x}_i, y_i)\}$ with kernel $k(\cdot, \cdot)$ and noise variance $\sigma_n^2$, the posterior at target $\mathbf{x}^*$ is

$$
\hat y(\mathbf{x}^*) = \mathbf{k}_*^\top (K + \sigma_n^2 I)^{-1} \mathbf{y},
\qquad
\sigma^2(\mathbf{x}^*) = k(\mathbf{x}^*, \mathbf{x}^*) - \mathbf{k}_*^\top (K + \sigma_n^2 I)^{-1} \mathbf{k}_*,
$$

with $K_{ij} = k(\mathbf{x}_i, \mathbf{x}_j)$ and $\mathbf{k}_* = (k(\mathbf{x}^*, \mathbf{x}_i))_i$. Hyperparameters (length-scale $\ell$, signal variance $\sigma_f^2$, noise $\sigma_n^2$) are fit by maximizing the log-marginal likelihood — sklearn does this automatically with L-BFGS-B restarts.

**Cost:** $O(N^3)$ for the fit, $O(N^2 M)$ for $M$ predictions. Practical ceiling $N \sim 10^3$–$10^4$ points without sparse approximations.

## 5. API

```python
def gp_to_grid(
lons, lats, values, grid,
*,
kernel: Literal["matern","rbf","rational_quadratic"] | object = "matern",
length_scale: float | tuple[float, float] = 1.0,
nu: float = 1.5, # for matern only
noise_level: float = 1e-3,
n_restarts: int = 3,
return_std: bool = False,
) -> xr.Dataset:
"""Returns a Dataset with `mean` (and `std` if return_std=True)."""

def gp_to_points(...): ...

def fillnan_gp(
da, *, lon="lon", lat="lat",
kernel="matern", length_scale=1.0, nu=1.5,
noise_level=1e-3, max_points: int = 2000,
return_std: bool = False,
) -> xr.DataArray | xr.Dataset:
"""Fill NaNs by GP from a subsample of `max_points` finite neighbours
(subsampled to keep cost bounded)."""
```

Operator wrappers follow the patterns from #159 / kNN issue.

## 6. Implementation sketch

```python
# _src/gp.py
import numpy as np
import xarray as xr

try:
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import (
Matern, RBF, RationalQuadratic, WhiteKernel, ConstantKernel,
)
_HAVE_SKLEARN_GP = True
except ImportError:
_HAVE_SKLEARN_GP = False

_KERNELS = {
"matern": lambda l, **kw: Matern(length_scale=l, nu=kw.get("nu", 1.5)),
"rbf": lambda l, **kw: RBF(length_scale=l),
"rational_quadratic": lambda l, **kw: RationalQuadratic(length_scale=l),
}

def _build_kernel(kernel, length_scale, nu, noise_level):
if not isinstance(kernel, str):
return kernel
base = _KERNELS[kernel](length_scale, nu=nu)
return ConstantKernel(1.0) * base + WhiteKernel(noise_level=noise_level)

def gp_to_grid(lons, lats, values, grid, *, kernel="matern",
length_scale=1.0, nu=1.5, noise_level=1e-3,
n_restarts=3, return_std=False):
if not _HAVE_SKLEARN_GP:
raise ImportError(
"GP interpolation requires scikit-learn. "
"Install with: pip install 'xr_toolz[ml-interp]'"
)
X = np.column_stack([np.ravel(lons), np.ravel(lats)])
y = np.ravel(values).astype(np.float64)
finite = np.isfinite(y)
k = _build_kernel(kernel, length_scale, nu, noise_level)
gp = GaussianProcessRegressor(
kernel=k, n_restarts_optimizer=n_restarts,
normalize_y=True,
)
gp.fit(X[finite], y[finite])
lon_g, lat_g = np.meshgrid(grid.lon, grid.lat, indexing="xy")
Q = np.column_stack([lon_g.ravel(), lat_g.ravel()])
if return_std:
mean, std = gp.predict(Q, return_std=True)
return xr.Dataset({
"mean": (("lat","lon"), mean.reshape(len(grid.lat), len(grid.lon))),
"std": (("lat","lon"), std.reshape(len(grid.lat), len(grid.lon))),
}, coords={"lat": grid.lat, "lon": grid.lon})
mean = gp.predict(Q)
return xr.DataArray(
mean.reshape(len(grid.lat), len(grid.lon)),
dims=("lat","lon"), coords={"lat": grid.lat, "lon": grid.lon},
)
```

## 7. Usage examples

```python
from xr_toolz.interpolate import gp_to_grid

# Sparse Argo SST onto a regular grid, with uncertainty
out = gp_to_grid(lons, lats, sst, grid,
kernel="matern", length_scale=2.0, nu=1.5,
return_std=True)
out["mean"].plot()
out["std"].plot() # uncertainty map
```

## 8. Tests (DoD)

1. **Skip if sklearn unavailable.** `pytest.importorskip("sklearn")`.
2. **Identity-at-source.** With `noise_level → 0`, prediction at a source point recovers the source value exactly.
3. **Smooth recovery.** $f(x,y)=\sin(x)\cos(y)$ sampled at 100 random points; max error on a 32×32 grid < `1e-2` with default Matern.
4. **Std drops near sources, grows far from them.** Standard GP behaviour — assert.
5. **Kernel preset dispatch.** Each of `"matern"`, `"rbf"`, `"rational_quadratic"` runs without error and returns finite values.
6. **Custom kernel pass-through.** Pass `kernel=Matern(length_scale=0.5)` directly; predictions match an sklearn-only reference.
7. **`return_std=True` returns Dataset; `False` returns DataArray.** Type contract.
8. **`fillnan_gp` subsamples large source sets.** With `max_points=100` and 10⁴ finite neighbours, the fit completes in reasonable time.
9. **Operator round-trip.** Custom kernel object stored as `` in config (mirrors `` convention).

## 9. Subtasks

1. Add `xr_toolz[ml-interp]` extra to `pyproject.toml`.
2. Implement `_src/gp.py` per §6.
3. Operators.
4. Public exports + namespace decision (export from `gap_fill` too).
5. Tests per §8.
6. Docs notebook: sparse altimetry-track SSH → GP map with uncertainty side-by-side; comparison with kNN/IDW (cheaper, no uncertainty) and biharmonic (no uncertainty, dense fields only).

## 10. Open questions

- **Default kernel.** Matern ν=3/2 is the right default for ocean / atmos data (less smooth than RBF, more realistic for turbulent fields). RBF is too smooth for SST/SSH. **Recommendation:** Matern.
- **Default `length_scale`.** No principled default — depends on the data. Sklearn will optimize from the initial value via marginal likelihood. **Recommendation:** `length_scale=1.0` initial, document that the user should pass a sensible scale (in coord units — degrees for lat/lon).
- **Hyperparameter optimization on or off?** sklearn's `n_restarts_optimizer=0` would freeze hyperparams (faster, but possibly worse). **Recommendation:** `n_restarts=3` default — costs a few seconds, big robustness improvement.
- **Memory ceiling.** GP is $O(N^3)$. `max_points=2000` for `fillnan_gp` is a defensive default. Document the limitation prominently. Sparse GPs are a future issue.
- **`normalize_y=True`.** Always normalize internally (matches sklearn convention for stable optimization). Document.
- **Should we expose `dask-ml`?** dask-ml lacks a first-class GP. **Recommendation:** punt until they have one, or until a user explicitly asks.
- **Spherical / haversine distance.** Same issue as IDW — Euclidean in (lon, lat) is biased near poles. **Recommendation:** flag in docs; require pre-projected coords for high-lat work; track as a follow-up.
- **Should `fillnan_gp` live in `gap_fill.py`?** Same call as kNN — yes, also export from there.

Start with the proposed new file src/xr_toolz/interpolate/_src/gp.py and compare operator patterns from interpolate/operators.py, issue #159, and the kNN issue. Review the pyproject.toml dependency setup and implement the listed MVP entry points before adding public exports and tests. Done means the §8 behavior, optional dependency handling, operator round-trip, and documentation notebook are covered.

Written by the indexing model from the issue text.