pipekit-cycle: protocols & primitives for variational (4DVar) and nudging (BFN) data assimilation

## TL;DR

`pipekit-cycle` cleanly models the **filtering / ensemble** family of data assimilation (`DACycle`, `EnsembleDACycle`, `SmootherCycle`). It has **no vocabulary for the two other major DA families**: **variational** (3D/4D-Var) and **nudging** (back-and-forth nudging, BFN).

An audit of the MASSH sea-surface-height mapping codebase ([`leguillf/MASSH@VarDyn`](https://github.com/leguillf/MASSH/tree/VarDyn)) — which ships a full **4DVar-SW** and **BFN-QG** stack — surfaced a coherent set of missing **protocols** and **composition primitives**. This issue proposes adding them while keeping pipekit-cycle **algorithm-agnostic**: the numerics (minimizer internals, filter math) stay in downstream algorithm libraries (filterx, vardax, …); pipekit ships only the *contracts* and thin *composition operators* that those libraries satisfy and compose.

Scope of this issue: **protocols + primitives/operators only.** The concrete 4DVar/BFN/EnKF *algorithms* are explicitly out of scope (a follow-up round).

---

## Background & motivation

Classical DA decomposes into three families, and pipekit-cycle currently only spans one:

| Family | Update mechanism | pipekit-cycle today |
|---|---|---|
| **Sequential / ensemble** | `analysis = forecast + K·(obs − H·forecast)`; Kalman gain from covariances | ✅ `DACycle`, `EnsembleDACycle`, `SmootherCycle`, `ObservationOperator`, `AnalysisStep` |
| **Variational (3D/4D-Var)** | minimize `J(x) = ½‖x−x_b‖²_B⁻¹ + ½‖H(x)−y‖²_R⁻¹` over a window, using the **adjoint** of the forward model and obs operator to get `∇J` | ❌ nothing |
| **Nudging (BFN)** | integrate the model forward and backward, relaxing the state toward observations with a gain term; iterate to convergence | ❌ nothing |

The current protocols (`packages/pipekit-cycle/src/pipekit_cycle/protocols.py`) encode only the sequential contract:

- `ForwardModel` exposes `step(state, dt)` / `dt` / `state_signature` — **no tangent-linear, no adjoint** (`protocols.py:35-41`).
- `ObservationOperator` exposes `__call__` + `linearize` — **no adjoint/transpose application** (`protocols.py:60-62`).
- `AnalysisStep` is `(forecast, obs, *, obs_op, obs_err_cov) → analysis` — **no notion of a differentiable cost `J` and its gradient** (`protocols.py:79-86`).
- `obs_err_cov` is threaded as an opaque `Any` (`DAState.obs_err_cov`) — **no covariance contract** with the `√` / `inv-√` operations preconditioning needs.

MASSH demonstrates exactly what's missing. Verbatim from `MASSH@VarDyn`:

- **Model adjoint/TLM are first-class**: every dynamical model carries `def step(...)`, `def step_tgl(...)`, `def step_adj(...)` — `Qgm` (`models/model_qg1l/jqgm.py`), `SW` (`models/model_qgsw/sw.py`), `CSWm` (`models/model_sw1l/jswm.py`).
- **Obs operator carries its adjoint**: every `Obsop_*` has `def misfit(self, t, State)` (forward innovation) and `def adj(self, t, adState, R)` (`src/obsop.py`).
- **Cost function is explicit**: `class Variational` → `cost`, `grad`, `cost_and_grad`, with `J = np.float64(0.5 * (Jo + Jb))`, `Jo += _m.dot(self.R.inv(misfit))`, `Jb = X0.dot(X0)` / `np.dot(X0, self.B.inv(X0))` (`src/tools_4Dvar.py`).
- **Covariance with √ / inv-√**: `class Cov` → `inv` (`1/sigma**2 * X`), `sqr` (`sigma * X`), `invsqr` (`1/sigma * X`) (`src/tools_4Dvar.py`).
- **Reduced control-vector transform**: 8 `Basis_*` classes, each exposing `operg` (control→state) and `operg_transpose` (its adjoint), driven from `Variational` via `self.basis.operg(...)` / `operg_transpose(...)` (`src/basis.py`).
- **Minimizer driver**: `Inv_4Dvar` → `opt.minimize(wrapper, ..., method=config.INV.opt_method, jac=wrapper.jac, callback=callback)` wrapped by `class Wrapper`, with `ConvergenceReached` / `CrazyGradient` retry control (`src/inv.py`).
- **BFN nudging**: `class bfn_qg1l` → `compute_nudging_term`, `convergence(path_forth, path_back)`, `update_parameter`; plus `bfn_nudge_smoothing`, `bfn_select_obs_temporal_window` (`src/tools_bfn.py`).
- **Localization kernel**: `def gaspari_cohn(r, c)` (`models/model_qg1l/jqgm.py`), reused throughout the basis and nudging tapers.

---

## User stories

- **As a variational-DA researcher**, I want to express a 4D-Var problem as a pipekit graph — forward model + obs operator + their adjoints + B/R covariances + a reduced control vector — so I can plug in any minimizer and any forward model that satisfies the protocols, without rewriting the outer loop.
- **As an algorithm-library author (vardax)**, I want runtime-checkable protocols for `CostFunction`, `Covariance`, and `ControlVectorTransform` so my classes interop with pipekit cycles **without importing pipekit-cycle**, exactly as filter classes do today.
- **As an SSH-mapping practitioner**, I want a ready-made `ForwardBackwardCycle` + `NudgingStep` so I can build a BFN-QG mapper by supplying only the QG model and a nudging gain.
- **As a hybrid-ML modeler**, I want to wrap a differentiable (JAX/PyTorch) emulator as a `ForwardModel` and have its `tangent_linear` / `adjoint` come "for free" from autodiff, then drop it straight into a 4D-Var driver.
- **As a maintainer**, I want these additions to stay **pure-Python, dependency-free, and algorithm-agnostic**, consistent with the existing protocol-driven design.

---

## API proposal

All additions live in `pipekit-cycle`. Protocols extend `pipekit_cycle.protocols`; operators get new modules. Nothing here imports a numerical backend.

### 1. Adjoint / tangent-linear contracts

```python
# pipekit_cycle/protocols.py (extend)
from typing import Any, Protocol, runtime_checkable

@runtime_checkable
class TangentLinearModel(Protocol):
"""A ForwardModel that also exposes its tangent-linear and adjoint."""
def step(self, state: Any, dt: float) -> Any: ...
def tangent_linear(self, dstate: Any, state: Any, dt: float) -> Any: ...
"""Propagate a perturbation `dstate` linearised about `state`."""
def adjoint(self, adstate: Any, state: Any, dt: float) -> Any: ...
"""Apply the adjoint (transpose) of the tangent-linear step."""

@runtime_checkable
class AdjointObservationOperator(Protocol):
"""An ObservationOperator that can apply H and Hᵀ."""
def __call__(self, state: Any) -> Any: ... # H(x)
def linearize(self, state: Any) -> Any: ... # tangent-linear H
def adjoint(self, dobs: Any, state: Any) -> Any: ... # Hᵀ · dobs
```

`ForwardModel` and `ObservationOperator` stay as-is (backward compatible); the new protocols are the differentiable refinements. A `@runtime_checkable` `isinstance(model, TangentLinearModel)` lets a driver detect adjoint support and fall back to autodiff otherwise.

### 2. Cost function (variational objective)

```python
# pipekit_cycle/protocols.py
@runtime_checkable
class CostFunction(Protocol):
"""A differentiable objective J(x) over a control/state vector."""
def value(self, x: Any) -> float: ...
def grad(self, x: Any) -> Any: ...
def value_and_grad(self, x: Any) -> tuple[float, Any]: ...
```

### 3. Covariance (B and R) with sqrt / inverse machinery

```python
# pipekit_cycle/protocols.py
@runtime_checkable
class Covariance(Protocol):
"""Symmetric positive-definite operator with the ops DA needs.

Mirrors MASSH `Cov`: apply, inverse-apply, and the square-root /
inverse-square-root used for preconditioning the control vector.
"""
def apply(self, x: Any) -> Any: ... # C · x
def inv(self, x: Any) -> Any: ... # C⁻¹ · x
def sqrt(self, x: Any) -> Any: ... # C^{1/2} · x
def inv_sqrt(self, x: Any) -> Any: ... # C^{-1/2} · x
```

### 4. Control-vector transform (change of variable / B^{1/2})

```python
# pipekit_cycle/protocols.py
@runtime_checkable
class ControlVectorTransform(Protocol):
"""Map a reduced control vector to state space and back (its adjoint).

Mirrors MASSH `Basis.operg` / `Basis.operg_transpose`.
"""
def to_state(self, x_ctrl: Any, t: float | None = None) -> Any: ...
def to_control_adjoint(self, adstate: Any, t: float | None = None) -> Any: ...
```

### 5. Minimizer contract

```python
# pipekit_cycle/protocols.py
@runtime_checkable
class Minimizer(Protocol):
"""Drives a CostFunction to a (local) minimum. Adapter over scipy/optax/etc."""
def minimize(self, cost: "CostFunction", x0: Any) -> Any: ...
```

### 6. New composition operators

```python
# pipekit_cycle/innovation.py
class Innovation(Operator):
"""Compute the (optionally R-weighted) innovation d = obs − H(forecast)."""
def __init__(self, obs_op: ObservationOperator,
obs_err_cov: Covariance | None = None): ...

# pipekit_cycle/window.py
class ObservationWindow(Operator):
"""Select observations active within [t0, t1] (temporal binning),
mirroring MASSH `bfn_select_obs_temporal_window`."""
def __init__(self, obs_source: Operator, window: float, stride: float): ...

# pipekit_cycle/nudging.py
class NudgingStep(Operator):
"""One relaxation step: x ← x + K · taper · (obs − H(x)).

`gain` (the nudging coefficient K) and `taper` (e.g. Gaspari-Cohn in
space/time) are supplied by the caller. Mirrors `bfn_qg1l.compute_nudging_term`.
"""
def __init__(self, obs_op, gain, taper=None): ...

# pipekit_cycle/cycle.py (new wrapper)
class ForwardBackwardCycle(StatefulOperator):
"""Back-and-forth nudging: integrate forward with NudgingStep, then
backward, iterate until `converged(forth, back) < tol` or `max_iters`.
Mirrors MASSH BFN outer loop (`bfn_qg1l.convergence`)."""
def __init__(self, forward_model, nudging_step, *,
max_iters: int = 10, tol: float = 1e-3): ...

# pipekit_cycle/variational.py (driver; algorithm-agnostic)
class FourDVar(StatefulOperator):
"""Assemble a windowed 4D-Var cost from (model, obs_op, B, R, control)
and hand it to a Minimizer. Computes ∇J via the model+obs adjoints when
available, else via autodiff. Mirrors MASSH `Inv_4Dvar`."""
def __init__(self, forward_model, obs_op, *, background, B, R,
control: ControlVectorTransform | None = None,
minimizer: Minimizer, n_steps: int, precondition: bool = True): ...
```

### 7. Correlation-kernel utility

```python
# pipekit_cycle/kernels.py (pure functions; no numpy dependency — duck-typed xp)
def gaspari_cohn(r, c, *, xp=math): ... # compact-support 5th-order piecewise poly
def gaussian_kernel(r, length_scale, *, xp=math): ...
def matern_kernel(r, length_scale, nu=1.5, *, xp=math): ...
```

> **Design note.** Following the existing pattern (`forward.py`, `obs.py`), each concrete operator is a `pipekit.Operator` that *also* satisfies the relevant protocol structurally, so libraries can supply their own implementations without subclassing.

---

## API examples

**4D-Var (sketch).** Forward model + obs operator with adjoints, B/R covariances, a wavelet control vector, and a scipy L-BFGS minimizer:

```python
from pipekit_cycle import FourDVar, Innovation

dvar = FourDVar(
forward_model=qg_model, # satisfies TangentLinearModel (has tangent_linear/adjoint)
obs_op=altimetry_obs, # satisfies AdjointObservationOperator
background=x_b,
B=BackgroundCov(sigma=0.1), # satisfies Covariance (apply/inv/sqrt/inv_sqrt)
R=DiagObsCov(sigma=0.03),
control=WaveletControl(...), # satisfies ControlVectorTransform (to_state/to_control_adjoint)
minimizer=ScipyMinimizer(method="L-BFGS-B", maxiter=200),
n_steps=window_steps,
precondition=True, # optimize in B^{1/2} space
)
analysis = dvar(forecast0, DAState(t=0.0, obs_err_cov=R))
```

**BFN-QG (sketch).** Back-and-forth nudging into a QG model with a Gaspari-Cohn taper:

```python
from pipekit_cycle import ForwardBackwardCycle, NudgingStep
from pipekit_cycle.kernels import gaspari_cohn

nudge = NudgingStep(
obs_op=ssh_obs,
gain=0.5,
taper=lambda r: gaspari_cohn(r, c=3.0),
)
bfn = ForwardBackwardCycle(forward_model=qg_model, nudging_step=nudge,
max_iters=10, tol=1e-3)
mapped = bfn(first_guess, DAState(t=t0))
```

**Autodiff fallback.** A JAX emulator that doesn't hand-code an adjoint still works:

```python
# NeuralForward gains tangent_linear/adjoint via jax.jvp / jax.vjp in an adapter,
# so isinstance(model, TangentLinearModel) is True and FourDVar uses it directly.
```

---

## TODO

- [ ] Add `TangentLinearModel`, `AdjointObservationOperator`, `CostFunction`, `Covariance`, `ControlVectorTransform`, `Minimizer` protocols to `protocols.py` (+ `__all__`, runtime-checkable).
- [ ] `kernels.py`: `gaspari_cohn`, `gaussian_kernel`, `matern_kernel` (duck-typed `xp`, pure-Python default).
- [ ] `innovation.py`: `Innovation` operator (+ optional R-weighting).
- [ ] `window.py`: `ObservationWindow` selector.
- [ ] `nudging.py`: `NudgingStep` operator.
- [ ] `cycle.py`: `ForwardBackwardCycle` wrapper (+ convergence callback hook).
- [ ] `variational.py`: `FourDVar` driver (adjoint-from-protocol with autodiff fallback; preconditioned + unpreconditioned control; checkpointing hook for the adjoint sweep).
- [ ] Adjoint **checkpointing primitive** for the trajectory (stride + recompute), so the backward sweep doesn't need the whole forward trajectory in memory (MASSH uses `jax.checkpoint`).
- [ ] Re-export all new symbols from `pipekit_cycle/__init__.py`.
- [ ] Tests: `gaspari_cohn` compact-support + partition-of-unity; adjoint/tangent dot-product test helper (` == `, mirroring MASSH `adjoint_test`); `FourDVar` gradient check via finite differences (mirroring MASSH `grad_test`); BFN convergence on a linear toy model.
- [ ] Docs notebook: 4D-Var and BFN on a 1-D toy model (Burgers or linear advection) end-to-end.
- [ ] Update `CLAUDE.md` workspace map + master-plan Report 10 cross-reference.

## Open questions

- Should `tangent_linear`/`adjoint` be **separate protocols** (as proposed) or **optional methods** on `ForwardModel`/`ObservationOperator`? Separate protocols keep the base contract minimal and let drivers detect capability via `isinstance`.
- Does the adjoint-checkpointing primitive belong in `pipekit-cycle` or in core `pipekit.state`? It's reusable beyond DA.
- `Minimizer` adapters (scipy, optax, jaxopt) — ship thin adapters here, or leave entirely to downstream libs and only define the protocol?

## References

- MASSH (SSH mapping; 4DVar-SW + BFN-QG): https://github.com/leguillf/MASSH/tree/VarDyn
- `src/tools_4Dvar.py` — `Cov`, `Variational` (cost/grad/cost_and_grad, Jb/Jo, preconditioning)
- `src/inv.py` — `Inv_4Dvar`, `Wrapper`, `scipy.optimize.minimize` driver
- `src/tools_bfn.py` — `bfn_qg1l` (nudging term, forward/backward convergence)
- `src/basis.py` — `Basis_*` reduced control vectors (`operg` / `operg_transpose`)
- `src/obsop.py` — `Obsop_*` (`misfit` / `adj`)
- `models/model_qg1l/jqgm.py` — `step_tgl` / `step_adj`, `gaspari_cohn`
- Current pipekit-cycle contracts: `packages/pipekit-cycle/src/pipekit_cycle/protocols.py`, `forward.py`, `obs.py`, `da.py`
- Master plan Report 10 (pipekit-cycle), §2.3 (protocols), §2.5–2.6 (obs/forward)
- Background:
- Le Dimet & Talagrand (1986), *Variational algorithms for analysis and assimilation of meteorological observations.* Tellus A.
- Auroux & Blum (2008), *A nudging-based data assimilation method: the Back and Forth Nudging algorithm.* Nonlin. Processes Geophys.
- Gaspari & Cohn (1999), *Construction of correlation functions in two and three dimensions.* QJRMS.
- Le Guillou et al. (2021), *Mapping altimetry in the forthcoming SWOT era by back-and-forth nudging a one-layer QG model.* JTECH.