jejjohnson / jejjohnson/pyrox

gp: TransformFilterPrior — model surface and predictive pushforward for warped filters

Open
#206 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1
Forks
0
Avg merge
19h 2m
Merged PRs (30d)
16

Description

## Problem / Request

`TransformFilter` gives a filter. It does not give a **model**: no NumPyro sites on
the dynamics or noise parameters, no `factor`, no path into
`pyrox.inference.EnsembleMAP` / `EnsembleVI`, and — most importantly — no correct
`predict`.

Add `TransformFilterPrior` in `pyrox-gp`, mirroring `MarkovGPPrior`, and make it
the place where the two pushforward traps are handled once so users never meet
them.

## User Story

> As a user fitting a state-space model with a strictly positive state, I want to
> put priors on the dynamics and noise parameters and get calibrated predictions in
> physical units, so that I am not hand-writing an Optax loop and then
> exponentiating the posterior mean by mistake.

## Motivation

`pyrox-gp` already owns this pattern: `MarkovGPPrior` (`_markov.py:146`) wraps a
kernel, exposes `log_marginal`, and `markov_gp_factor` (`:453`) registers it with
NumPyro. That is what gives hyperpriors, SVI/MCMC, and seed ensembling.

Two reasons this issue is more than boilerplate:

**1. `predict` is where the traps live.** Both were measured, and both produce
plausible-looking wrong answers rather than errors:

- `Γ(E[z])` is the pushforward **median**, not the mean. Comparing it against a
physical-space mean made UKF look 15.2% *worse* than it is.
- Symmetric intervals on a warped predictive over-cover at 1.000 and **still put
15/80 intervals below zero** — after all the work of warping. Transformed
quantiles give 0.975, are narrower (2.1820 vs 2.2428), and cross zero zero times.

A user calling `TransformFilter` directly will hit both. A user calling
`prior.predict(...)` should not.

**2. Ensembling matters more here than usual.** The composed map `Γ⁻¹ ∘ f ∘ Γ` can
be more nonlinear than `f`, so moment-matched filters are more sensitive to
initialisation. Evidence: MCKF is the worst filter in physical coordinates
(RMSE 2.91) and the best in latent coordinates (0.416) — a 7× swing. `EnsembleMAP`
over seeds is the cheap defence and comes free from this surface.

### Scoping note, worth stating in the module docstring

`pyrox-gp`'s existing strategies — `LaplaceInference`, `GaussNewtonInference`,
`PosteriorLinearization`, `ExpectationPropagation`, and their Markov variants —
all call `_check_scalar_latent` and handle a non-Gaussian **likelihood** over a
scalar latent. This issue handles non-Gaussian **state support**. The axes are
orthogonal and compose; neither subsumes the other, and no existing strategy
should be modified.

## Proposed API

```python
class TransformFilterPrior(eqx.Module):
r"""State-space prior filtered in warped coordinates.

Wraps a `gauss_flows.TransformFilter` into the `pyrox_gp` model
surface so the dynamics, noise, and warp parameters can carry NumPyro
priors and be fitted with any `pyrox.inference` driver.

Prediction is returned in **physical** units with the pushforward
done correctly -- see `predict` and `predict_interval`.

Attributes:
filt: The configured `gauss_flows.TransformFilter`.
process_noise: $Q$, in the frame given by ``filt.noise_frame``.
obs_noise: $R$.
init_mean: Latent initial mean.
init_cov: Latent initial covariance.
"""

def log_marginal(self, y, mask=None) -> Float[Array, ""]:
r"""Moment-matched marginal log-likelihood.

Warning:
Exact only when the composed maps are affine. Otherwise this
is a surrogate: the innovation covariance is moment-matched,
not exact. Hyperparameters fitted by maximising it are
maximising the surrogate.
"""

def predict(self, y, mask=None, *, order: int = 32):
r"""Smoothed predictive moments in **physical** units.

Returns $(\mathbb{E}[x], \mathrm{Var}[x])$ by Gauss-Hermite
pushforward -- **not** $\Gamma(\mathbb{E}[z])$, which is the
pushforward median.
"""

def predict_interval(self, y, mask=None, *, level: float = 0.95):
r"""Physical-space credible interval by transformed quantiles.

Symmetric moment intervals are not offered: on a warped
predictive they over-cover *and* place mass outside the support.
"""

def transform_filter_factor(name, prior, y, mask=None) -> None:
"""Register the marginal log-likelihood with NumPyro.

Mirrors `pyrox_gp.markov_gp_factor`.
"""
```

## Design Snapshot

Intended usage, matching the `MarkovGPPrior` docstring style:

```python
import numpyro, numpyro.distributions as dist
from gauss_flows import TransformFilter
from gaussx import UnscentedIntegrator, nonlinear_kalman_filter
from pyrox_gp import TransformFilterPrior, transform_filter_factor

def model(y, mask=None):
log_q = numpyro.sample("log_q", dist.Normal(0.0, 1.0).expand([N]))
log_r = numpyro.sample("log_r", dist.Normal(0.0, 1.0).expand([M]))
filt = TransformFilter(
warp_state=log_warp,
inner=partial(nonlinear_kalman_filter, integrator=UnscentedIntegrator()),
)
prior = TransformFilterPrior(
filt, jnp.diag(jnp.exp(log_q)), jnp.diag(jnp.exp(log_r)), m0, P0,
)
transform_filter_factor("tf", prior, y, mask)
```

Optional-dependency guard, following the `optax` precedent at
`_inference_nongauss.py:790` — import inside the call, not at module scope, and
reuse the `pyrox-gp[flows]` extra introduced by the `NormalizingKalmanPrior`
issue rather than adding a second one.

## Mathematical Notes

```text
Latent: z_t = f̃(z_{t-1}) + q_t, f̃ = Γ_x⁻¹ ∘ f ∘ Γ_x
Physical: x_t = Γ_x(z_t)

PREDICTION. Smoothing gives z_t ~ N(m_t, S_t) in latent coordinates. Then

E[x_t] = ∫ Γ(z) N(z; m_t, S_t) dz <- quadrature, NOT Γ(m_t)
Var[x_t] = ∫ Γ(z)² N(...) dz − E[x_t]²

Γ(m_t) is the MEDIAN of the pushforward when Γ is monotone.

INTERVALS. For monotone Γ, quantiles transform exactly:

P(x < Γ(m + z_α s)) = P(z < m + z_α s) = α

so Γ(m ± z_α s) is exact at the nominal level, while E[x] ± z_α sd[x] is not:

method coverage (nominal 0.95) width crosses zero
E[x] ± 1.96 sd[x] 1.000 2.2428 15/80
Γ(m ± 1.96 s) 0.975 2.1820 0

The symmetric interval is WIDER, over-covers, and still puts mass below zero.

GAUSS-HERMITE CAVEAT, carried over from the TGP work: convergence plateaus around
3e-3 for spline warps and can get WORSE at higher order, because GH converges
spectrally only for analytic integrands. Order 32 is the sweet spot. Do not
document "raise the order until it converges" as a diagnostic.

LOG-MARGINAL is a moment-matched surrogate unless the composed maps are affine.
Say so in the docstring rather than letting users assume exactness by analogy
with markov_gp_factor.
```

## References & Existing Code

- Design doc: `transform_filter_design.md` §3.2, §3.3, §4
- Evidence: `ntf_fair.py` (mean-vs-median), `ntf_final.py` (interval construction)
- **The surface to mirror**: `packages/pyrox-gp/src/pyrox_gp/_markov.py:146` (`MarkovGPPrior`), `:376` (`ConditionedMarkovGP`), `:453` (`markov_gp_factor`), `:469` (`markov_gp_sample`)
- Optional-import precedent: `packages/pyrox-gp/src/pyrox_gp/_inference_nongauss.py:790`
- Orthogonal, do not modify: `_inference_nongauss.py:361-758`, `_inference_nongauss_markov.py:189-427`
- Ensemble drivers unlocked: `packages/pyrox/src/pyrox/inference/` (`EnsembleMAP`, `EnsembleVI`, `ensemble_predict`)
- Upstream: `gauss_flows.TransformFilter`, `gaussx.nonlinear_kalman_filter`, `gaussx.GaussHermiteIntegrator`

## Implementation Steps

- [ ] Add `packages/pyrox-gp/src/pyrox_gp/_transform_filter.py` with `TransformFilterPrior` and `transform_filter_factor`
- [ ] Reuse the existing `pyrox-gp[flows]` extra; do not add a second one
- [ ] Guard the `gauss_flows` import inside the call, with an install hint in the message
- [ ] Implement `predict` via Gauss-Hermite pushforward, default `order=32`
- [ ] Implement `predict_interval` with transformed quantiles; do not expose a symmetric option
- [ ] Document `log_marginal` as a surrogate for non-affine composed maps
- [ ] Export from `packages/pyrox-gp/src/pyrox_gp/__init__.py` behind the guard
- [ ] Add both to the `members:` list in `docs/api/gp.md`
- [ ] Skip the test module with `pytest.importorskip("gauss_flows")` so the base install stays green
- [ ] Add a `diagnostics()` helper reporting the sigma-point spread or condition number of the composed map, so users can see when the warp made the problem harder

## Definition of Done

- [ ] Code lands at the intended path
- [ ] Public API exported via `packages/pyrox-gp/src/pyrox_gp/__init__.py`
- [ ] Tests pass: `make test` both with and without the `flows` extra installed
- [ ] Lint + typecheck pass: `make lint && make typecheck`
- [ ] Docstrings (Google-style, `$...$` math, no Sphinx/RST) on all public symbols

## Testing

New file `packages/pyrox-gp/tests/gp/test_transform_filter.py`.

- [ ] **Reduction test:** an identity warp makes `log_marginal` equal a direct `gaussx.nonlinear_kalman_filter` call — exactly. Catches a mis-wired warp default
- [ ] **Affine-exactness test:** with an affine warp and linear dynamics, `log_marginal` equals `gaussx.kalman_filter`'s exact marginal to ≤ 1e-10, confirming the surrogate is exact where it should be
- [ ] **Trap 1 test:** `predict` differs from `Γ(smoothed mean)` for a non-affine warp, and matches a 400k-sample Monte-Carlo estimate of `E[x]` within MC tolerance
- [ ] **Trap 2 test:** `predict_interval` achieves coverage 0.975 with mean width 2.1820 and zero support violations; assert the symmetric construction would give 1.000 / 2.2428 / 15 violations, so the test documents why the option is absent
- [ ] Support test: no predictive interval crosses zero on the positive-state problem, for any inner integrator
- [ ] NumPyro integration: `numpyro.handlers.trace` shows the expected sample sites plus one `factor`; SVI runs 50 steps and the loss decreases
- [ ] `EnsembleMAP` over 4 seeds runs and returns finite parameters
- [ ] Import test: with `gauss_flows` absent, importing `pyrox_gp` succeeds and constructing `TransformFilterPrior` raises `ImportError` with the install hint
- [ ] Non-interference test: an existing `PosteriorLinearizationMarkov` test is unchanged — this issue must not touch the non-Gaussian-likelihood path
- [ ] Docstring examples execute (the repo's doctest gate)

## Documentation

- [ ] `docs/api/gp.md` — new section stating up front that this addresses non-Gaussian **state support**, distinct from the existing strategies' non-Gaussian **likelihood**, and that it needs `pyrox-gp[flows]`
- [ ] Notebook: positive-state model fitted with NumPyro priors, showing physical-space intervals from both interval constructions side by side
- [ ] Docstrings (covered by Definition of Done)

## Relationships

- Parent (theme epic): #
- Blocked by: # (gaussx: `nonlinear_kalman_filter`), # (gauss_flows: `TransformFilter`)
- Blocks: #
- Related: # (pyrox: `NormalizingKalmanPrior` — shares the `[flows]` extra and the quadrature pushforward)

Contributor guide

Open the contributing guide

Research direction

Start with packages/pyrox-gp/src/pyrox_gp/_markov.py, especially MarkovGPPrior and markov_gp_factor, then review transform_filter_design.md §§3.2–4 and the gauss_flows TransformFilter API. Add the model and factor at _transform_filter.py, export them, document them in docs/api/gp.md, and add packages/pyrox-gp/tests/gp/test_transform_filter.py. Done means the listed reduction, pushforward, interval, optional-import, NumPyro, ensemble, lint, typecheck, and documentation checks pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend-api-design, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.