Report 10 — `pipekit-cycle`: time-stepping, DA, observation operators

   Domain libraries           geotoolz │ xr_toolz
                                  ▲
                                  │
   Infrastructure         pipekit-cycle (this report)
                                  ▲
                                  │ depends on pipekit.state
                                  │
   Framework                   pipekit ◄── pipekit-array
                                  ▲
                                  │
   Substrate              georeader │ xarray ecosystem

pipekit-cycle/
  __init__.py             # public re-exports
  _src/
    cycle.py              # Cycle, EnsembleCycle, WindowedCycle, Recurrence
    protocols.py          # ForwardModel, ObservationOperator, AnalysisStep
    da.py                 # DACycle, EnsembleDACycle, SmootherCycle
    obs.py                # IdentityObs, LinearObs, CallableObs, CompositeObs
    forward.py            # CompositeForward, CallableForward, NeuralForward (stub)

class Cycle(StatefulOperator):
    """Apply step_op for n_steps, threading state through.
    
    Equivalent shape to jax.lax.scan but works on any carrier.
    Outputs: (final_carrier, final_state, history) where history is
    an optional sequence of intermediate states.
    
    `step_op` may be any `StatefulOperator`, or a `ForwardModel` —
    a ForwardModel is wrapped on construction via the
    `ForwardModel → StatefulOperator` shim, so domain libraries can
    pass their `ForwardModel` adapters directly without importing
    pipekit-cycle internals.
    """
    step_op: StatefulOperator | ForwardModel
    n_steps: int
    save_history: bool = False
    history_stride: int = 1   # save every nth step

class EnsembleCycle(StatefulOperator):
    """Run Cycle over an ensemble of initial states in parallel.
    
    The carrier is a stacked ensemble (leading axis = ensemble member).
    Operators inside step_op see one member at a time — vectorisation
    is delegated to the user's operator (e.g., pipekit-array's BatchedMap,
    or pipekit-jax's filter_vmap).
    """
    step_op: StatefulOperator
    n_steps: int
    n_members: int

class WindowedCycle(StatefulOperator):
    """Sliding-window analysis. Apply step_op over rolling windows
    of an input stream, advancing by stride.
    
    Used for windowed reanalysis (4D-Var over fixed assimilation windows).
    """
    step_op: StatefulOperator
    window: int
    stride: int

class Recurrence(StatefulOperator):
    """Apply step_op repeatedly until convergence condition is met.
    
    For iterative inference: fixed-point methods, Newton iteration,
    EM, variational inference. The condition_op is itself an Operator
    that returns a bool from the current state.
    """
    step_op: StatefulOperator
    condition_op: Operator   # state → bool
    max_iters: int = 1000

@runtime_checkable
class ForwardModel(Protocol):
    """Advance state forward in time by dt.
    
    Implementations: plumax.PlumeForward, somax.OceanForward, neural
    emulators, hybrid physics+ML models.
    """
    def step(self, state: Any, dt: float) -> Any: ...
    @property
    def dt(self) -> float: ...
    @property  
    def state_signature(self) -> Signature | None: ...

@runtime_checkable
class ObservationOperator(Protocol):
    """Map model state to predicted observations (the H operator in DA).
    
    H(x) gives 'what would the observations look like if the state were x?'.
    Critical for innovation = obs - H(forecast).
    """
    def __call__(self, state: Any) -> Any: ...   # state → predicted_obs
    def linearize(self, state: Any) -> Any: ...   # tangent linear, optional

@runtime_checkable
class AnalysisStep(Protocol):
    """Combine forecast state with observations to produce analysis state.
    
    Implementations: filterx.EnKFAnalysis, vardax.FourDVarAnalysis,
    particle filter, ensemble Kalman smoother.
    """
    def __call__(self, forecast: Any, obs: Any, *, obs_op: ObservationOperator, 
                 obs_err_cov: Any) -> Any: ...

class DACycle(StatefulOperator):
    """Standard forecast-analysis cycle.
    
    Each call:
        1. Forecast: state → forward_model.step(state, dt) → forecast_state
        2. Read observations for this time window from obs_source
        3. Analyse: analysis_step(forecast_state, obs, obs_op, err_cov) → analysis_state
        4. Carry analysis_state forward as the new state
    """
    forward_model: ForwardModel
    obs_op: ObservationOperator
    analysis_step: AnalysisStep
    obs_source: Operator | None = None   # observation injection (typically a catalog reader)
    n_steps: int
    save_history: bool = True

class EnsembleDACycle(StatefulOperator):
    """Ensemble-based DA — EnKF, ETKF, LETKF, etc.
    
    State is a stack of ensemble members. Analysis step is the ensemble
    update. Concrete analysis_step implementations come from filterx.
    """
    forward_model: ForwardModel
    obs_op: ObservationOperator
    analysis_step: AnalysisStep   # an ensemble analysis (filterx.EnKFAnalysis)
    obs_source: Operator | None
    n_steps: int
    n_members: int

class SmootherCycle(StatefulOperator):
    """4D-Var-shaped: forward pass over window, backward adjoint pass.
    
    Useful for variational analysis where the full window is solved
    jointly. analysis_step here is a 4D-Var solver from vardax.
    """
    forward_model: ForwardModel
    obs_op: ObservationOperator
    analysis_step: AnalysisStep   # a smoother (vardax.FourDVarSolver)
    window: int
    stride: int

class IdentityObs(Operator):
    """Direct measurement: H = I. The state is observable as-is."""

class LinearObs(Operator):
    """H is a linear matrix. Useful when state and obs differ in dimension
    via a fixed linear map (e.g., subsampling, channel-mixing).
    """
    H: ndarray   # array-shaped, multi-backend via array_namespace

class CallableObs(Operator):
    """H is an arbitrary callable. Catch-all for nonlinear observation
    operators (radiative-transfer-style: state to radiances).
    
    The callable might itself be a pipekit pipeline (e.g., a forward
    radiative transfer model wrapped as an Operator).
    """
    fn: Callable
    forbid_in_yaml: ClassVar[bool] = True

class CompositeObs(Operator):
    """H = h2 ∘ h1. Compose observation operators when the obs path
    has multiple stages (e.g., chemistry transport → column integrator
    → satellite-instrument-response convolution).
    """
    components: tuple[ObservationOperator, ...]

class CallableForward(Operator):
    """Wrap any callable as a ForwardModel. The escape hatch."""
    fn: Callable
    dt: float

class CompositeForward(Operator):
    """Compose multiple forward models — e.g., advection ∘ chemistry ∘ deposition.
    
    Each component advances state by its own dt; CompositeForward enforces
    consistent total dt across components.
    """
    components: tuple[ForwardModel, ...]

class NeuralForward(Operator):
    """Stub for neural emulator forward models. Wraps a pipekit-array
    ModelOp or pipekit-jax JaxModelOp into the ForwardModel protocol.
    
    The actual neural model is trained via pipekit-train (Report 11)
    and stored in the model registry (Report 12).
    """
    model_op: Operator   # ModelOp or JaxModelOp
    dt: float

# In filterx — your Kalman / particle filter library
# Pure algorithm code; no pipekit dependency in the core.

# filterx.adapters.pipekit (new, ~150 LOC, extras-gated via filterx[pipekit])
class EnKFAnalysis:
    """Ensemble Kalman analysis step. Satisfies pipekit_cycle.AnalysisStep
    protocol via structural typing.
    """
    inflation: float = 1.0
    localization: Any = None
    
    def __call__(self, forecast_ensemble, obs, *, obs_op, obs_err_cov):
        # Wrap filterx's EnKF math in the AnalysisStep signature
        ...

class ParticleFilterAnalysis:
    """Particle filter / sequential Monte Carlo. Same protocol."""
    ...

# vardax.adapters.pipekit
class FourDVarAnalysis:
    """4D-Var solver as an AnalysisStep. Wraps vardax's variational
    solver in the cycle-friendly interface.
    """
    background_cov: Any
    inner_iters: int = 50
    outer_iters: int = 3
    
    def __call__(self, background_state, obs_window, *, obs_op, obs_err_cov):
        ...

# plumax.adapters.pipekit
class PlumeForward:
    """Plume dispersion forward model as a pipekit_cycle.ForwardModel.
    Wraps plumax's plume math in the step(state, dt) interface.
    """
    dispersion_params: Any
    
    def step(self, plume_state, dt):
        ...

class PlumeObservationOp:
    """Map plume state → predicted column-integrated concentrations.
    
    Specific to satellite-observable plume products. Domain-specific
    observation operator; satisfies pipekit_cycle.ObservationOperator.
    """
    instrument_response: Any
    
    def __call__(self, plume_state):
        ...

import pipekit as pk
import pipekit_cycle as pc
import geocatalog as gc
import geotoolz as gz
from filterx.adapters.pipekit import EnKFAnalysis
from plumax.adapters.pipekit import ChemistryForward, ColumnObs
from statecatalog import StateCatalog

# Pieces from the algorithm + domain libraries
forward = ChemistryForward(species=["CH4", "NH3"], dt=3600.0)
obs_op  = ColumnObs(instrument="TROPOMI")
analysis = EnKFAnalysis(inflation=1.05, localization="gaspari-cohn")

# Catalog source for observations at each cycle step
obs_catalog = gc.open_catalog("s3://tropomi/2024/*.parquet")
obs_source = pk.Sequential([
    gc.CatalogTimeQuery(window="hourly"),   # custom op: pick obs for current cycle time
    gz.tropomi.LoadL2(),
    gz.qa.AssertValidFraction(min_valid=0.3),
])

# The cycle itself
da = pc.EnsembleDACycle(
    forward_model=forward,
    obs_op=obs_op,
    analysis_step=analysis,
    obs_source=obs_source,
    n_steps=24,         # one day
    n_members=40,
    save_history=True,
)

# Run from an initial state, write outputs to the state catalog
state_cat = StateCatalog.open("s3://reanalysis/methane_2024/")
final_state, history = da(initial_state, initial_carry_state)
state_cat.write(history, model_config_hash=pk.config_hash(da))

# Just a forecast — no observations, no analysis.
forecast = pc.Cycle(
    step_op=forward,         # the ChemistryForward from above
    n_steps=72,              # 3-day forecast
    save_history=True,
    history_stride=1,        # save every hour
)

trajectory = forecast(analysis_state, initial_carry_state)
state_cat.write(trajectory, model_config_hash=pk.config_hash(forecast))

# Convergence-driven recurrence — not time-stepping.
retrieval = pc.Recurrence(
    step_op=LMStep(jacobian_fn=...),       # one LM update
    condition_op=pk.qc.AssertCallable(
        predicate=lambda state: state.residual_norm < 1e-6 or state.iter >= 100
    ),
    max_iters=100,
)

retrieved_state = retrieval(initial_guess, initial_iter_state)

# Same cycle structure, but the forward model is a learned emulator
# loaded from the model registry.

import pipekit_experiment as pe

registry = pe.S3ModelRegistry("s3://models/")
neural_forward = pc.NeuralForward(
    model_op=registry.load("methane_emulator_v3"),   # content-addressed
    dt=3600.0,
)

# Same DACycle structure — only the forward model changed
da_emulated = pc.EnsembleDACycle(
    forward_model=neural_forward,          # ← swapped
    obs_op=obs_op,                          # same
    analysis_step=analysis,                 # same
    obs_source=obs_source,                  # same
    n_steps=24,
    n_members=40,
)

[project]
name = "pipekit-cycle"
version = "0.1.0"
dependencies = [
    "pipekit>=0.1",
    "numpy>=2.0",
]

[project.optional-dependencies]
array  = ["pipekit-array>=0.1"]    # for NeuralForward / NeuralObservation via ModelOp
jax    = ["pipekit-jax>=0.1"]      # for ScanCycle, JaxDACycle (deferred)
catalog = ["statecatalog>=0.1"]    # for state-aware cycles that read/write state catalogs
all    = ["pipekit-cycle[array,catalog]"]

sub = pk.Sequential([
    StatelessOpA(),     # carrier → carrier
    StatefulOpB(),       # (carrier, state) → (carrier, state)
    StatelessOpC(),
    StatefulOpD(),
])