api(estimator): add sklearn-style immutable Estimator facade with GPEstimator validation
- 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
Assessment
This issue has not been assessed yet.