jejjohnson / jejjohnson/pyrox

api(estimator): add sklearn-style immutable Estimator facade with GPEstimator validation

Open
#71 0 comments 0 reactions 0 assignees View on GitHub
area:core area:gp type:feature wave:4-structured
Dominant language
Python
Stars
1
Forks
0
Avg merge
19h 2m
Merged PRs (30d)
16

Description

## Problem / Request

`pyrox.api` already ships `EstimatorBase`, `FittedEstimator`, and `BNFEstimator` (with `BNFEstimatorMLE`, `BNFEstimatorVI`, and `FittedBNF`). The sklearn-style facade is real and in use by the BNF flagship demo (#73).

What's missing is the **GP analogue**: `GPEstimator` (and a fitted variant) that wraps the existing `pyrox.gp` model classes (`MaternKernel + ExactGP`, sparse / Markov variants) behind the same `cfg.fit(df).predict(df, quantiles=...)` ergonomic. Today users have to hand-wire a `numpyro` model, a guide, an SVI loop, and `Predictive` — exactly the friction `EstimatorBase` was designed to remove.

## User Story

> As a data-scientist user with a `pandas.DataFrame` of feature columns + a numeric target, I want
>
> ```python
> est = pyrox.api.GPEstimator(
> kernel="matern52",
> feature_cols=["lat", "lon", "elevation"],
> target_col="ppm",
> )
> fitted = est.fit(df_train, seed=0)
> df_pred = fitted.predict(df_test, quantiles=(0.025, 0.5, 0.975))
> ```
>
> so I get a calibrated GP regression model with one call, matching the `BNFEstimator` ergonomic.

## Motivation

Concrete reasons this lands now:

- **Closes the BNF-vs-GP API symmetry.** `BNFEstimator` exists and is documented (#73); a user who wants a GP fallback instead has to drop down two layers.
- **Removes the manual SVI loop from tutorials.** Most pyrox.gp tutorials currently spend 30+ lines wiring a model + guide + SVI step + `Predictive`. `GPEstimator` collapses this to two lines.
- **Validation surface** — the issue title mentions "validation" because the `Estimator` facade is the natural place for `feature_cols` schema checks (dtypes, NaN handling, column order). Reusing the BNF facade's existing validation hooks.

The lower-level `pyrox.gp` API stays unchanged — `GPEstimator` is sugar over it, not a replacement.

## Mathematics

The estimator is a thin wrapper around the standard GP regression posterior

$$
p(y_* \mid X_*, D) \;=\; \mathcal{N}\!\bigl(\mu_*, \Sigma_*\bigr),
$$

with

$$
\mu_* = K_{*X}(K_{XX} + \sigma^2 I)^{-1} y,
\qquad
\Sigma_* = K_{**} - K_{*X}(K_{XX} + \sigma^2 I)^{-1} K_{X*}.
$$

The estimator's `.predict(df, quantiles=...)` reduces the posterior to the requested quantiles via `numpyro.distributions.Normal.icdf` on the marginal $\mathcal{N}(\mu_*, \mathrm{diag}(\Sigma_*))$ at each test point.

No new math — the value-add is the immutable `Fitted*` pattern, dataframe schema validation, and the sklearn-style ergonomic.

## Proposed API

```python
class GPEstimator(EstimatorBase):
"""sklearn-style facade over pyrox.gp.

Args:
feature_cols: Required input columns (validated against df).
target_col: Required target column.
kernel: Kernel spec — "rbf", "matern52", "matern32", or a Kernel instance.
mean_fn: Optional mean function.
inducing: Optional inducing-point spec ("kmeans", "uniform", or array).
max_iter: SVI iterations.
"""

feature_cols: tuple[str, ...]
target_col: str
kernel: str | pgp.Kernel
inducing: str | int | None = None
max_iter: int = 1_000
learning_rate: float = 0.01

def fit(self, df: pd.DataFrame, *, seed: int) -> "FittedGP":
"""Fit and return an immutable FittedGP carrying the learned params."""

class FittedGP(FittedEstimator):
"""Immutable container with the trained GP and predict() method."""

posterior_params: dict
feature_cols: tuple[str, ...]
target_col: str

def predict(
self,
df: pd.DataFrame,
*,
quantiles: tuple[float, ...] | None = None,
return_std: bool = False,
) -> pd.DataFrame:
"""Return predictions; columns include the requested quantiles or mean+std."""
```

## Example usage

```python
import pandas as pd, pyrox.api as pyx

df_train = pd.DataFrame({"lat": [...], "lon": [...], "elev": [...], "co2": [...]})
df_test = pd.DataFrame({"lat": [...], "lon": [...], "elev": [...]})

est = pyx.GPEstimator(
feature_cols=["lat", "lon", "elev"],
target_col="co2",
kernel="matern52",
inducing=200, # sparse approx with 200 inducing points
max_iter=2_000,
)
fitted = est.fit(df_train, seed=0)
predictions = fitted.predict(df_test, quantiles=(0.025, 0.5, 0.975))
# predictions columns: lat, lon, elev, co2_q025, co2_q500, co2_q975
```

## Tasks / sub-tasks

- [ ] **Spike** — design the kernel-spec parsing so `"matern52"` resolves to `pyrox.gp.kernels.Matern52()` without circular imports between `pyrox.api` and `pyrox.gp`.
- [ ] **`GPEstimator`** — `src/pyrox/api/_gp.py`. Subclass `EstimatorBase`; reuse the dataframe-validation helpers from `_bnf.py`.
- [ ] **`FittedGP`** — same module; subclass `FittedEstimator`. Carries the posterior params and the kernel.
- [ ] **Sparse / inducing path** — `inducing=N` triggers a sparse-variational fit using the existing `pyrox.gp` sparse machinery.
- [ ] **Public export** — `pyrox.api.__init__`.
- [ ] **Tests** — (i) on a 50-point synthetic regression dataset, `GPEstimator.fit(...).predict(...)` matches a hand-wired `pyrox.gp` exact-GP fit to `1e-4`; (ii) inducing=N triggers the sparse path; (iii) schema validation rejects missing columns / wrong dtypes; (iv) quantiles agree with `scipy.stats.norm.ppf` on a Gaussian posterior; (v) `jit` / `vmap` compatibility through `fit`.
- [ ] **Docs** — `pyrox.api` reference entry; cross-link from the BNF demo (#73) as "GP fallback".

## References

- Existing facade: [`src/pyrox/api/_estimator.py`](https://github.com/jejjohnson/pyrox/blob/main/src/pyrox/api/_estimator.py) (`EstimatorBase`, `FittedEstimator`).
- Existing BNF facade: [`src/pyrox/api/_bnf.py`](https://github.com/jejjohnson/pyrox/blob/main/src/pyrox/api/_bnf.py) (`BNFEstimator`, `BNFEstimatorVI`, `FittedBNF`).
- Existing GP machinery: [`src/pyrox/gp/`](https://github.com/jejjohnson/pyrox/blob/main/src/pyrox/gp).
- Flagship demo waiting on this: [`docs/notebooks/bayesian_neural_fields.ipynb`](#73).
- sklearn estimator API: .

## Relationships

- Parent (wave epic): #33
- Blocked by: #
- Blocks: #
- Related: #73

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.