research: DataAssimBench — OSSE observation generation + benchmark recipe for pipekit-cycle / pipekit-evaluate

# Comparative Analysis: `DataAssimBench` vs `pipekit-cycle` / `pipekit-evaluate`

## Context

[StevePny/DataAssimBench](https://github.com/StevePny/DataAssimBench) (`dabench`) is a JAX-based benchmarking suite for data assimilation (DA): synthetic-truth generators (Lorenz63/96, SQG turbulence, QG variants, ERA5/NetCDF loaders), an `Observer` that manufactures synthetic observations from a truth trajectory, observation operators, and four DA "cyclers" (3D-Var, ETKF, 4D-Var, backprop-4D-Var). A companion repo, [DataAssimBench-Examples](https://github.com/StevePny/DataAssimBench-Examples), holds tutorial notebooks plus the complete benchmark protocol from Solvik, Penny & Hoyer (2024) on backprop-4DVar.

This survey asked what pipekit can borrow. Architecturally, `dabench` is the inverse of pipekit's design — an inheritance-heavy `DACycler` base class with `_in_4d` / `_uses_ensemble` flags, subclass-required `Model` wrappers, and custom `StateVector`/`ObsVector` containers that duplicate xarray — so its *architecture* should not be copied; pipekit-cycle's runtime-checkable Protocols are strictly better. What is worth taking is three *features*: (1) the **Observer** abstraction — an OSSE observation *generator*, which pipekit-cycle lacks entirely; (2) the **OSSE benchmark recipe** (nature run → observe → assimilate → score vs. truth, with a no-DA baseline) as the concrete first recipe for the `pipekit-evaluate` benchmark ladder; (3) **window time-alignment utilities** for collecting ragged observation streams into fixed-shape analysis windows.

> OSSE = Observing System Simulation Experiment: a model run with known initial conditions serves as "truth" (the *nature run*), synthetic observations are sampled from it with controlled noise, and a DA method assimilates them — so the analysis error is exactly measurable. It is the standard way to evaluate any DA method.

---

## 1. What `dabench` Contains (relevant subset)

```
dabench/
├── data/ # truth generators: Lorenz63/96, SQGTurb, PyQG(+JAX), QGS, ERA5/GCP, generic NetCDF
├── observer/
│ └── _observer.py # Observer — samples synthetic observations from a nature run
├── obsop/
│ └── _obsop.py # ObsOp — linear H matrix or nonlinear h(x)
├── dacycler/
│ ├── _dacycler.py # base cycle(): jax.lax.scan over analysis+forecast cycles
│ ├── _var3d.py / _etkf.py / _var4d.py / _var4d_backprop.py
│ └── _utils.py # time alignment: group + pad ragged obs per analysis window
├── metrics/ # JAX MSE / RMSE / MAE / Pearson-r
└── vector/ # StateVector / ObsVector containers (NOT worth copying)
```

### A. `Observer` (`dabench/observer/_observer.py`)

Constructor (abridged):

```python
Observer(
state_vec, # the nature run
random_time_density=1.0, # fraction of timesteps observed, OR
random_time_count=None, # absolute number of observed times
random_location_density=1.0, # fraction of state locations observed, OR
random_location_count=None, # absolute count (scalar or per-dim tuple)
times=None, locations=None, # ... or fully explicit indices
stationary_observers=True, # fixed network vs. new locations each time
error_bias=0.0, error_sd=0.0, # Gaussian noise; scalar or per-variable array
error_positive_only=False, # clip for positive-definite quantities
random_seed=99,
)
```

- Time and location sampling are independently specified by *density* (fraction) or *count* (absolute), or given explicitly.
- `stationary_observers=False` draws new locations per timestep (a moving network), padding ragged counts so downstream code sees fixed shapes.
- `.observe()` returns the observation values **plus** `system_index` (which state elements were observed) **plus** `errors` (the actual noise realisations drawn) — so an experiment can audit its own synthetic-obs process after the fact.

### B. Window time alignment (`dabench/dacycler/_dacycler.py`, `_utils.py`)

```python
cycler.cycle(
input_state, start_time, obs_vector,
n_cycles,
analysis_window=0.2, # window length in time units
analysis_time_in_window=None, # where analysis is valid; default = mid-window
return_forecast=False,
)
```

- Observations are grouped per analysis window and **padded to a uniform per-window count**, so the entire multi-cycle loop runs under a single `jax.lax.scan` (one compilation, scannable, differentiable end-to-end).
- `analysis_time_in_window` decouples "window of observations used" from "time at which the analysis is valid" (default: middle of window, not the edge).

### C. OSSE benchmark protocol (examples repo)

The Solvik, Penny & Hoyer (2024) notebooks (`examples/da_cycler/Solvik_Penny_Hoyer_2024-Backprop4DVar/`) run a published, citable benchmark:

- **Sweep axes**: system dimension ∈ {6, 20, 36, 72, 144, 256} (Lorenz96), number of observed variables ∈ {6 … 36}, observation error σ ∈ {0.1 … 2.0}.
- **Baseline**: a *no-DA free-running forecast* is always included as the floor every method must beat.
- **Metrics**: analysis RMSE vs. truth + wall-clock time, reported jointly.

---

## 2. Comparison with pipekit

### A. Already in pipekit (direct equivalents)

| dabench feature | pipekit equivalent | Path | Notes |
|---|---|---|---|
| `ObsOp` (H matrix / h(x)) | `IdentityObs`, `LinearObs`, `CallableObs`, `CompositeObs` | `packages/pipekit-cycle/src/pipekit_cycle/obs.py:25-109` | ours compose; theirs is a single class |
| `DACycler.cycle()` orchestration | `Cycle`, `DACycle`, `EnsembleDACycle`, `SmootherCycle` | `pipekit_cycle/cycle.py`, `pipekit_cycle/da.py` | ours are carrier-agnostic `StatefulOperator`s |
| `Model` wrapper (subclass-required) | `ForwardModel` Protocol | `pipekit_cycle/protocols.py:36` | protocol seam beats mandatory subclassing |
| TLM trajectories from `generate(return_tlm=True)` | `TangentLinearModel` Protocol | `pipekit_cycle/protocols.py:136` | we declare the seam; algorithm libs implement |
| `metrics/` (RMSE etc.) | pipekit-evaluate scorer taxonomy | planned | no action; covered by existing plan |

### B. Already in pipekit but missing enhancements from `dabench`

#### B1. Ragged-observation window alignment + analysis placement (MEDIUM PRIORITY)
- **`dabench`**: groups an irregular observation stream into per-window sets, pads each window to a uniform observation count (so the cycle loop is shape-stable / scannable), and lets the caller place the analysis time anywhere in the window (`analysis_time_in_window`, default mid-window).
- **pipekit**: `WindowedCycle` (`pipekit_cycle/cycle.py:195`) and `SmootherCycle` (`pipekit_cycle/da.py:227`) assume the window's observations arrive already grouped; there is no utility that builds fixed-shape windows from an irregular obs stream, and no notion of where in the window the analysis is valid.
- **What's needed**: a carrier-agnostic helper in `pipekit_cycle` — e.g. `align_observations(obs_times, t0, window, stride) → list of (per-window index groups, padding mask)` (pure-Python index bookkeeping, no arrays needed) — plus an optional `analysis_time_in_window` parameter on `SmootherCycle`.
- **Impact**: required for any real observation stream (irregular satellite times); the fixed-shape grouping is also exactly what makes a downstream JAX consumer's whole cycle loop jit/scan-able.

### C. Missing completely from pipekit

#### C1. Observation generator ("Observer") for OSSEs (HIGH PRIORITY)
- **What it is**: the §1.A component — given a truth stream, decide *which* times and locations are observed (density/count/explicit, stationary or moving network), apply a noise model, and emit observations plus a full audit record (sampled indices + drawn noise).
- **Why useful**: pipekit-cycle currently has observation *operators* (mapping state → obs space) but nothing that *manufactures* observations from truth. Without it, every OSSE user hand-rolls sampling and noise injection, and results aren't reproducible/auditable across experiments. This is also the missing input leg for the pipekit-evaluate benchmark ladder (C2).
- **Where in pipekit**: proposed `packages/pipekit-cycle/src/pipekit_cycle/observer.py`. Design constraint — pipekit-cycle depends on `pipekit` only (pure Python), so split:
- the **sampling plan** (which time indices, which location indices, stationary flag, seed) is pure-Python index bookkeeping → lives in pipekit-cycle, e.g. an `ObservationPlan` dataclass + `Observer` operator;
- the **noise injection** is array math → delegate to the existing `ObservationNoise` protocol (`pipekit_cycle/protocols.py:248`), implemented downstream (vardax, filterax, pipekit-array).
- **Sketch** (exact API for the follow-up feature issue):

```python
@dataclass
class ObservationPlan:
time_indices: list[int]
location_indices: list[list[int]] # one entry total if stationary, else per time
stationary: bool
seed: int

class Observer(Operator):
"""Sample observations from a truth stream per an ObservationPlan."""
def __init__(self, plan: ObservationPlan, obs_op, noise: ObservationNoise | None = None): ...
def __call__(self, truth_stream): ... # → (observations, plan record)
```

#### C2. OSSE benchmark recipe + no-DA reference for pipekit-evaluate (HIGH PRIORITY)
- **What it is**: the §1.C protocol as a first concrete recipe for the benchmark ladder: nature run → `Observer` → cycle → score-vs-truth, with (a) canonical sweep axes for the benchmark cube — system dimension, observation count/density, observation error σ — and (b) the **no-DA free-running forecast** as the reference rule every method must beat.
- **Why useful**: `pipekit-evaluate` has the scaffolding (`benchmark/cube.py`, `benchmark/reference.py`, `benchmark/run.py`) but no canonical recipe yet. This one is proven and published (Solvik, Penny & Hoyer 2024), and its baseline maps one-to-one onto the existing reference-rule concept.
- **Where in pipekit**: `packages/pipekit-evaluate/src/pipekit_evaluate/benchmark/` — OSSE axes for `cube.py`, no-DA baseline for `reference.py`, the loop in `run.py`. Depends on C1.

### D. pipekit has it, `dabench` doesn't (context only — no action)
- Protocol seams (`ForwardModel`, `ObservationOperator`, `AnalysisStep`, `ObservationNoise`, …) vs. their base-class flags (`_in_4d`, `_uses_ensemble`) — algorithm libraries plug in without subclassing.
- Carrier-agnostic core — `dabench` is welded to JAX + xarray (`StateVector`/`ObsVector` duplicate xarray awkwardly; the `xarray_jax` threading through `lax.scan` is visibly painful).
- `EnsembleCycle`, `Recurrence`, control/observe/qc operator families, registry/tracker protocols — no equivalents there.

---

## 3. Summary Table

| Feature | `dabench` | pipekit | Status |
|---|---|---|---|
| Observation operators (H / h(x)) | ✓ | ✓ | **Already have** |
| Forward-model seam | ✓ (subclass) | ✓ (Protocol) | We're ahead |
| Cycle orchestration | ✓ | ✓ | **Already have** |
| Ragged-obs window alignment (shape-stable padding) | ✓ | ✗ | **Enhancement needed** (B1) |
| `analysis_time_in_window` placement | ✓ | ✗ | **Enhancement needed** (B1) |
| OSSE observation generator + audit record | ✓ | ✗ | **Missing** (C1) |
| No-DA baseline as reference rule | ✓ | scaffold only | **Missing** (C2) |
| Canonical OSSE sweep axes for benchmark cube | ✓ | ✗ | **Missing** (C2) |
| Carrier-agnostic core, Protocol seams | ✗ | ✓ | We're ahead |

---

## 4. Recommended Integration Priority

### Phase 1: Observer (enables the benchmark work)
1. `ObservationPlan` + `Observer` in `pipekit_cycle` (C1) — pure-Python sampling plan, noise via the existing `ObservationNoise` seam.

### Phase 2: OSSE benchmark recipe
2. OSSE recipe in `pipekit-evaluate` with no-DA reference rule + canonical sweep axes (C2). Depends on Phase 1.

### Phase 3: Window ergonomics
3. `align_observations` helper + `analysis_time_in_window` on `SmootherCycle` (B1). Independent of Phases 1–2; needed before any irregular real-data stream.

---

## 5. Proposed Follow-up Issues

- [ ] `feat(cycle): ObservationPlan + Observer — OSSE observation generator` — covers Phase 1 / C1
- [ ] `feat(evaluate): OSSE benchmark recipe with no-DA reference rule and canonical sweep axes` — covers Phase 2 / C2
- [ ] `feat(cycle): ragged-observation window alignment + analysis_time_in_window` — covers Phase 3 / B1

---

## References
- DataAssimBench — https://github.com/StevePny/DataAssimBench
- DataAssimBench-Examples — https://github.com/StevePny/DataAssimBench-Examples
- Solvik, Penny & Hoyer (2024), backprop-4DVar benchmark — notebooks under `examples/da_cycler/Solvik_Penny_Hoyer_2024-Backprop4DVar/`
- Hunt et al. (2007) — ETKF formulation used by their `_etkf.py`

## Relationships
- Parent (theme epic, if any): —
- Blocked by: —
- Blocks (follow-up issues from §5 reference this as parent research): TBD
- Related: jejjohnson/vardax#53 (solver guards, TLM/Lyapunov utilities, benchmark protocol); jejjohnson/xrtoolz#244 (state-vector round-trip + observation-sampling masks)