pipekit-cycle¶

Time-stepping, data assimilation, observation operators. Built on pipekit.state (the StatefulOperator + CarryState primitives in pipekit core). Master plan reference: Report 10.

Cycle wrappers¶

`pipekit_cycle.cycle` ¶

Generic time-stepping and iteration wrappers.

These are the cycle abstractions — Cycle, EnsembleCycle, WindowedCycle, Recurrence — that lift a per-step StatefulOperator (or a ForwardModel) into a multi-step operator. Each is itself a StatefulOperator, so cycles compose into longer cycles and into Sequential pipelines.

Design choice: each cycle returns (carrier, state) to preserve the StatefulOperator contract. When save_history=True the intermediate carriers are accumulated on self.history — a plain list attribute available after the call. Treat the attribute as write-once-per-call: re-running the operator overwrites it.

See master plan Report 10, section 2.2.

`Cycle` ¶

Bases: StatefulOperator

Apply step_op for n_steps, threading state through each step.

Equivalent in shape to jax.lax.scan but carrier-agnostic.

Parameters:

Name	Type	Description	Default
`step_op`	`Any`	A `StatefulOperator` or a `ForwardModel`. The latter is wrapped automatically so domain libraries can pass their forward-model adapters directly.	required
`n_steps`	`int`	Number of steps to take. Must be ≥ 0; `0` is a no-op that returns the inputs unchanged.	required
`save_history`	`bool`	When `True`, intermediate `(carrier, state)` pairs are appended to `self.history` every `history_stride` steps.	`False`
`history_stride`	`int`	Save every `stride`-th step (default 1).	`1`

Returns from _apply(carrier, state): (final_carrier, final_state).

Raises:

Type	Description
`TypeError`	If `step_op` is neither a `StatefulOperator` nor a `ForwardModel`.
`ValueError`	If `n_steps < 0` or `history_stride < 1`.

Example::

cycle = Cycle(step_op=Advect(dt=3600.0), n_steps=24, save_history=True)
carrier, state = cycle(initial_carrier, initial_state)
# cycle.history is the per-step trajectory

Source code in packages/pipekit-cycle/src/pipekit_cycle/cycle.py

class Cycle(StatefulOperator):
    """Apply ``step_op`` for ``n_steps``, threading state through each step.

    Equivalent in shape to ``jax.lax.scan`` but carrier-agnostic.

    Args:
        step_op: A `StatefulOperator` or a `ForwardModel`. The latter
            is wrapped automatically so domain libraries can pass their
            forward-model adapters directly.
        n_steps: Number of steps to take. Must be ≥ 0; ``0`` is a
            no-op that returns the inputs unchanged.
        save_history: When ``True``, intermediate ``(carrier, state)``
            pairs are appended to ``self.history`` every
            ``history_stride`` steps.
        history_stride: Save every ``stride``-th step (default 1).

    Returns from ``_apply(carrier, state)``: ``(final_carrier, final_state)``.

    Raises:
        TypeError: If ``step_op`` is neither a `StatefulOperator` nor a
            `ForwardModel`.
        ValueError: If ``n_steps < 0`` or ``history_stride < 1``.

    Example::

        cycle = Cycle(step_op=Advect(dt=3600.0), n_steps=24, save_history=True)
        carrier, state = cycle(initial_carrier, initial_state)
        # cycle.history is the per-step trajectory
    """

    __config_mixin_auto__ = False

    def __init__(
        self,
        step_op: Any,
        n_steps: int,
        save_history: bool = False,
        history_stride: int = 1,
    ) -> None:
        if n_steps < 0:
            raise ValueError(f"Cycle.n_steps must be >= 0, got {n_steps}.")
        if history_stride < 1:
            raise ValueError(
                f"Cycle.history_stride must be >= 1, got {history_stride}."
            )
        self.step_op = _coerce_step(step_op)
        self.n_steps = n_steps
        self.save_history = save_history
        self.history_stride = history_stride
        self.history: list[tuple[Any, Any]] = []

    def _apply(self, carrier: Any, state: Any) -> tuple[Any, Any]:
        self.history = []
        for i in range(self.n_steps):
            carrier, state = self.step_op(carrier, state)
            if self.save_history and (i % self.history_stride == 0):
                self.history.append((carrier, state))
        return carrier, state

    def get_config(self) -> dict[str, Any]:
        return {
            "step_op": nested_config(self.step_op),
            "n_steps": self.n_steps,
            "save_history": self.save_history,
            "history_stride": self.history_stride,
        }

`EnsembleCycle` ¶

Bases: StatefulOperator

Run Cycle over a list of ensemble-member carriers.

The carrier is expected to be a sequence (list / tuple) of n_members individual carriers; the same step_op is applied to each member in turn, sharing the carry-state across members (the typical EnKF arrangement, where the state holds covariances or inflation factors that all members read).

Vectorisation across members is delegated to the user — for array-shaped carriers, wrap step_op in a batched / vmapped operator before passing it in.

Parameters:

Name	Type	Description	Default
`step_op`	`Any`	Per-member step. `StatefulOperator` or `ForwardModel`.	required
`n_steps`	`int`	Steps per member.	required
`n_members`	`int`	Expected ensemble size; the input carrier is length-checked at call time.	required

Returns from _apply(members, state): (list_of_final_members, state).

Raises:

Type	Description
`TypeError`	If `step_op` is neither a `StatefulOperator` nor a `ForwardModel`.
`ValueError`	If `n_steps < 0` or `n_members < 1` (from the constructor), or if the input ensemble's length differs from `n_members` (at call time).

Source code in packages/pipekit-cycle/src/pipekit_cycle/cycle.py

class EnsembleCycle(StatefulOperator):
    """Run `Cycle` over a list of ensemble-member carriers.

    The carrier is expected to be a sequence (list / tuple) of
    ``n_members`` individual carriers; the same ``step_op`` is applied
    to each member in turn, sharing the carry-state across members
    (the typical EnKF arrangement, where the state holds covariances
    or inflation factors that all members read).

    Vectorisation across members is delegated to the user — for
    array-shaped carriers, wrap ``step_op`` in a batched / vmapped
    operator before passing it in.

    Args:
        step_op: Per-member step. `StatefulOperator` or `ForwardModel`.
        n_steps: Steps per member.
        n_members: Expected ensemble size; the input carrier is
            length-checked at call time.

    Returns from ``_apply(members, state)``: ``(list_of_final_members, state)``.

    Raises:
        TypeError: If ``step_op`` is neither a `StatefulOperator` nor a
            `ForwardModel`.
        ValueError: If ``n_steps < 0`` or ``n_members < 1`` (from the
            constructor), or if the input ensemble's length differs
            from ``n_members`` (at call time).
    """

    __config_mixin_auto__ = False

    def __init__(self, step_op: Any, n_steps: int, n_members: int) -> None:
        if n_steps < 0:
            raise ValueError(f"EnsembleCycle.n_steps must be >= 0, got {n_steps}.")
        if n_members < 1:
            raise ValueError(f"EnsembleCycle.n_members must be >= 1, got {n_members}.")
        self.step_op = _coerce_step(step_op)
        self.n_steps = n_steps
        self.n_members = n_members

    def _apply(  # ty: ignore[invalid-method-override]
        self,
        members: list[Any],
        state: Any,
    ) -> tuple[list[Any], Any]:
        if len(members) != self.n_members:
            raise ValueError(
                f"EnsembleCycle expected {self.n_members} members, got {len(members)}."
            )
        out_members = list(members)
        for _ in range(self.n_steps):
            for j, carrier in enumerate(out_members):
                out_members[j], state = self.step_op(carrier, state)
        return out_members, state

    def get_config(self) -> dict[str, Any]:
        return {
            "step_op": nested_config(self.step_op),
            "n_steps": self.n_steps,
            "n_members": self.n_members,
        }

`Recurrence` ¶

Bases: StatefulOperator

Apply step_op repeatedly until condition_op returns truthy.

For iterative inference: fixed-point methods, Levenberg-Marquardt, EM, variational solvers. condition_op is an Operator that inspects the post-step state and returns a bool-ish value; the loop exits when it returns truthy or when max_iters is reached.

Parameters:

Name	Type	Description	Default
`step_op`	`Any`	Per-iteration step.	required
`condition_op`	`Operator`	`Operator` invoked as `condition_op(state)`. Truthy = stop.	required
`max_iters`	`int`	Safety cap on iterations (default 1000).	`1000`

Returns from _apply(carrier, state): (carrier, state) after either convergence or hitting max_iters.

Raises:

Type	Description
`TypeError`	If `step_op` is neither a `StatefulOperator` nor a `ForwardModel`, or if `condition_op` is not an `Operator`.
`ValueError`	If `max_iters < 1`.

Source code in packages/pipekit-cycle/src/pipekit_cycle/cycle.py

class Recurrence(StatefulOperator):
    """Apply ``step_op`` repeatedly until ``condition_op`` returns truthy.

    For iterative inference: fixed-point methods, Levenberg-Marquardt,
    EM, variational solvers. ``condition_op`` is an `Operator` that
    inspects the post-step state and returns a bool-ish value; the
    loop exits when it returns truthy or when ``max_iters`` is reached.

    Args:
        step_op: Per-iteration step.
        condition_op: `Operator` invoked as ``condition_op(state)``.
            Truthy = stop.
        max_iters: Safety cap on iterations (default 1000).

    Returns from ``_apply(carrier, state)``: ``(carrier, state)``
    after either convergence or hitting ``max_iters``.

    Raises:
        TypeError: If ``step_op`` is neither a `StatefulOperator` nor a
            `ForwardModel`, or if ``condition_op`` is not an `Operator`.
        ValueError: If ``max_iters < 1``.
    """

    __config_mixin_auto__ = False

    def __init__(
        self,
        step_op: Any,
        condition_op: Operator,
        max_iters: int = 1000,
    ) -> None:
        if not isinstance(condition_op, Operator):
            raise TypeError(
                "Recurrence.condition_op must be an Operator, "
                f"got {type(condition_op).__name__}."
            )
        if max_iters < 1:
            raise ValueError(f"Recurrence.max_iters must be >= 1, got {max_iters}.")
        self.step_op = _coerce_step(step_op)
        self.condition_op = condition_op
        self.max_iters = max_iters
        self.iters: int = 0
        self.converged: bool = False

    def _apply(self, carrier: Any, state: Any) -> tuple[Any, Any]:
        self.iters = 0
        self.converged = False
        for _ in range(self.max_iters):
            carrier, state = self.step_op(carrier, state)
            self.iters += 1
            if self.condition_op(state):
                self.converged = True
                break
        return carrier, state

    def get_config(self) -> dict[str, Any]:
        return {
            "step_op": nested_config(self.step_op),
            "condition_op": nested_config(self.condition_op),
            "max_iters": self.max_iters,
        }

`WindowedCycle` ¶

Bases: StatefulOperator

Sliding-window time-stepping over a stream of initial carriers.

The input is a list of carriers (one per stream position). For each window — windows start at positions 0, stride, 2*stride, … up to len(stream) - window — the cycle seeds a carrier from the first stream element of that window and threads it through window applications of step_op. State threads across windows.

Parameters:

Name	Type	Description	Default
`step_op`	`Any`	Per-step operator. `StatefulOperator` or `ForwardModel`.	required
`window`	`int`	Number of `step_op` applications per window.	required
`stride`	`int`	Step between successive window starts (in stream units).	`1`

Returns from _apply(stream, state): (list_of_window_final_carriers, state).

Raises:

Type	Description
`ValueError`	From the constructor for invalid `window` / `stride`, and from `_apply` when the stream is shorter than one window (which would otherwise silently produce zero windows).

Source code in packages/pipekit-cycle/src/pipekit_cycle/cycle.py

class WindowedCycle(StatefulOperator):
    """Sliding-window time-stepping over a stream of initial carriers.

    The input is a list of carriers (one per stream position). For
    each window — windows start at positions ``0, stride, 2*stride, …``
    up to ``len(stream) - window`` — the cycle seeds a carrier from
    the first stream element of that window and threads it through
    ``window`` applications of ``step_op``. State threads across
    windows.

    Args:
        step_op: Per-step operator. `StatefulOperator` or `ForwardModel`.
        window: Number of ``step_op`` applications per window.
        stride: Step between successive window starts (in stream units).

    Returns from ``_apply(stream, state)``:
        ``(list_of_window_final_carriers, state)``.

    Raises:
        ValueError: From the constructor for invalid ``window`` /
            ``stride``, and from ``_apply`` when the stream is shorter
            than one window (which would otherwise silently produce
            zero windows).
    """

    __config_mixin_auto__ = False

    def __init__(self, step_op: Any, window: int, stride: int = 1) -> None:
        if window < 1:
            raise ValueError(f"WindowedCycle.window must be >= 1, got {window}.")
        if stride < 1:
            raise ValueError(f"WindowedCycle.stride must be >= 1, got {stride}.")
        self.step_op = _coerce_step(step_op)
        self.window = window
        self.stride = stride

    def _apply(  # ty: ignore[invalid-method-override]
        self,
        stream: list[Any],
        state: Any,
    ) -> tuple[list[Any], Any]:
        outputs: list[Any] = []
        n = len(stream)
        if n < self.window:
            raise ValueError(
                f"WindowedCycle: stream of length {n} is shorter than one "
                f"window ({self.window}); no windows would be produced."
            )
        for start in range(0, n - self.window + 1, self.stride):
            carrier = stream[start]
            for _ in range(self.window):
                carrier, state = self.step_op(carrier, state)
            outputs.append(carrier)
        return outputs, state

    def get_config(self) -> dict[str, Any]:
        return {
            "step_op": nested_config(self.step_op),
            "window": self.window,
            "stride": self.stride,
        }

Protocols¶

`pipekit_cycle.protocols` ¶

Protocols decomposing data assimilation into its standard form.

The core triple — ForwardModel, ObservationOperator, AnalysisStep (predict, compare, update) — plus the reduced-order seams for variational and reduced-rank methods: ReducedBasis (a reduced control basis with a prior), TangentLinearModel (M' / M* of the dynamics), ErrorSubspace (a propagating low-rank covariance factor), and ReducedOrderModel (a Galerkin ROM — encode/decode + latent dynamics).

Ensemble-method seams: ObservationNoise (a noise model that can both report its covariance and draw samples), CovarianceLocalizer (distance-based covariance tapering), EnsembleInflator (post-analysis ensemble inflation), and LatentMap (an encode/decode pair for latent-space assimilation — the stateless subset of ReducedOrderModel).

Iterative-inversion seams: IterativeProcess (ensemble Kalman inversion / sampling — EKI, EKS, UKI) and StepScheduler (its step-size policy).

Each algorithm library (filterax, vardax, plumax, …) ships adapter classes that satisfy these protocols without importing pipekit-cycle. The protocols are runtime-checkable so isinstance(obj, ForwardModel) succeeds on any structurally compatible class.

See master plan Report 10, section 2.3.

`AnalysisStep` ¶

Bases: Protocol

Combine forecast state with observations to produce analysis state.

Implementations: ensemble Kalman analyses (EnKF, ETKF, LETKF), variational solvers (3D/4D-Var), particle filters, smoothers. Algorithm libraries supply concrete classes.

Members

call(forecast, obs, *, obs_op, obs_err_cov): Return the analysis state given the forecast, the observations, the observation operator, and the observation-error covariance.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class AnalysisStep(Protocol):
    """Combine forecast state with observations to produce analysis state.

    Implementations: ensemble Kalman analyses (EnKF, ETKF, LETKF),
    variational solvers (3D/4D-Var), particle filters, smoothers.
    Algorithm libraries supply concrete classes.

    Members:
        __call__(forecast, obs, *, obs_op, obs_err_cov): Return the
            analysis state given the forecast, the observations, the
            observation operator, and the observation-error covariance.
    """

    def __call__(
        self,
        forecast: Any,
        obs: Any,
        *,
        obs_op: ObservationOperator,
        obs_err_cov: Any,
    ) -> Any: ...

`CovarianceLocalizer` ¶

Bases: Protocol

Taper spurious long-range covariances in ensemble methods.

Localisation suppresses sampling noise in low-rank ensemble covariances by damping entries between distant locations (Gaspari-Cohn, Gaussian taper, hard cutoff, …). Cycles thread the localizer through to the AnalysisStep; supplying one is optional.

Members

call(covariance, coords): Return the localized covariance given per-element coordinates used to compute distances.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class CovarianceLocalizer(Protocol):
    """Taper spurious long-range covariances in ensemble methods.

    Localisation suppresses sampling noise in low-rank ensemble
    covariances by damping entries between distant locations
    (Gaspari-Cohn, Gaussian taper, hard cutoff, …). Cycles thread the
    localizer through to the `AnalysisStep`; supplying one is optional.

    Members:
        __call__(covariance, coords): Return the localized covariance
            given per-element coordinates used to compute distances.
    """

    def __call__(self, covariance: Any, coords: Any) -> Any: ...

`EnsembleInflator` ¶

Bases: Protocol

Inflate an analysis ensemble to counter variance underestimation.

Ensemble filters systematically underestimate spread; inflation (multiplicative, additive, relaxation-to-prior) corrects for it after each analysis. Relaxation methods blend in the forecast ensemble, hence the optional second argument.

Members

call(particles, forecast_particles=None, **kwargs): Return the inflated ensemble. forecast_particles is the prior ensemble for relaxation-to-prior schemes (RTPS / RTPP); pure multiplicative or additive schemes ignore it.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class EnsembleInflator(Protocol):
    """Inflate an analysis ensemble to counter variance underestimation.

    Ensemble filters systematically underestimate spread; inflation
    (multiplicative, additive, relaxation-to-prior) corrects for it
    after each analysis. Relaxation methods blend in the forecast
    ensemble, hence the optional second argument.

    Members:
        __call__(particles, forecast_particles=None, **kwargs): Return
            the inflated ensemble. ``forecast_particles`` is the prior
            ensemble for relaxation-to-prior schemes (RTPS / RTPP);
            pure multiplicative or additive schemes ignore it.
    """

    def __call__(
        self,
        particles: Any,
        forecast_particles: Any = None,
        **kwargs: Any,
    ) -> Any: ...

`ErrorSubspace` ¶

Bases: Protocol

A propagating low-rank error-covariance factor L (B ≈ L Lᵀ).

The reduced-rank representation behind SEEK / SEIK / RRSQRT and other reduced-order Kalman methods: instead of a full covariance, the error is carried as a small set of modes that are propagated through the linearised dynamics each forecast and updated by the analysis. Thread it alongside the model state (e.g. on DAState.extras) so a reduced-rank AnalysisStep can read and refresh it.

Members

propagate(model, state, dt): Advance the modes by dt about the linearisation point state — typically M' L via a TangentLinearModel, or finite differences of a ForwardModel. Returns the forecast error subspace. modes: The low-rank factor (its columns are the error modes). rank: Number of modes retained.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class ErrorSubspace(Protocol):
    """A propagating low-rank error-covariance factor ``L`` (``B ≈ L Lᵀ``).

    The reduced-rank representation behind SEEK / SEIK / RRSQRT and other
    reduced-order Kalman methods: instead of a full covariance, the error is
    carried as a small set of modes that are propagated through the
    linearised dynamics each forecast and updated by the analysis. Thread it
    alongside the model state (e.g. on ``DAState.extras``) so a reduced-rank
    `AnalysisStep` can read and refresh it.

    Members:
        propagate(model, state, dt): Advance the modes by ``dt`` about the
            linearisation point ``state`` — typically ``M' L`` via a
            `TangentLinearModel`, or finite differences of a `ForwardModel`.
            Returns the forecast error subspace.
        modes: The low-rank factor (its columns are the error modes).
        rank: Number of modes retained.
    """

    def propagate(self, model: Any, state: Any, dt: float) -> ErrorSubspace: ...

    @property
    def modes(self) -> Any: ...

    @property
    def rank(self) -> int: ...

`ForwardModel` ¶

Bases: Protocol

Advance a model state forward in time by dt.

Implementations: domain forward models (chemistry transport, ocean state, plume dispersion), neural emulators, hybrid physics + ML hybrids. Algorithm libraries provide adapters that satisfy this protocol structurally.

Members

step(state, dt): Return the state advanced by dt. dt: Default integration step. state_signature: Optional pipekit.Signature describing the shape / dtype of the state carrier. None if the model doesn't track named dimensions.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class ForwardModel(Protocol):
    """Advance a model state forward in time by ``dt``.

    Implementations: domain forward models (chemistry transport,
    ocean state, plume dispersion), neural emulators, hybrid
    physics + ML hybrids. Algorithm libraries provide adapters
    that satisfy this protocol structurally.

    Members:
        step(state, dt): Return the state advanced by ``dt``.
        dt: Default integration step.
        state_signature: Optional `pipekit.Signature` describing the
            shape / dtype of the state carrier. ``None`` if the model
            doesn't track named dimensions.
    """

    def step(self, state: Any, dt: float) -> Any: ...

    @property
    def dt(self) -> float: ...

    @property
    def state_signature(self) -> Any: ...

`IterativeProcess` ¶

Bases: Protocol

An iterative ensemble inversion / sampling process (EKI, EKS, UKI).

The parameter-estimation counterpart of the sequential-DA triple: instead of cycling through time, the ensemble iterates in pseudo-time toward the posterior over parameters. pipekit owns the loop (initialise, evaluate the forward model on each particle, update, repeat per the StepScheduler); algorithm libraries supply the process.

Members

init(particles, obs, noise_cov): Build the initial iteration state from the prior ensemble, the observations, and the observation-noise covariance. update(state, forward_evals, **kwargs): Apply one iteration given the forward-model evaluations of the current particles; return the updated state.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class IterativeProcess(Protocol):
    """An iterative ensemble inversion / sampling process (EKI, EKS, UKI).

    The parameter-estimation counterpart of the sequential-DA triple:
    instead of cycling through time, the ensemble iterates in
    pseudo-time toward the posterior over parameters. pipekit owns the
    loop (initialise, evaluate the forward model on each particle,
    update, repeat per the `StepScheduler`); algorithm libraries supply
    the process.

    Members:
        init(particles, obs, noise_cov): Build the initial iteration
            state from the prior ensemble, the observations, and the
            observation-noise covariance.
        update(state, forward_evals, **kwargs): Apply one iteration
            given the forward-model evaluations of the current
            particles; return the updated state.
    """

    def init(self, particles: Any, obs: Any, noise_cov: Any) -> Any: ...

    def update(self, state: Any, forward_evals: Any, **kwargs: Any) -> Any: ...

`LatentMap` ¶

Bases: Protocol

An encode/decode pair between full state space and a latent space.

The stateless subset of ReducedOrderModel — no latent dynamics, just the projection maps. This is the seam latent-space assimilation methods need (encode the ensemble, assimilate in latent space, decode the analysis) and what autoencoder priors naturally expose. Implementations: POD/EOF projections, autoencoders, or an identity map for testing.

Members

encode(state): Full state -> latent coordinates z. decode(coords): Latent coordinates -> full state.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class LatentMap(Protocol):
    """An encode/decode pair between full state space and a latent space.

    The stateless subset of `ReducedOrderModel` — no latent dynamics,
    just the projection maps. This is the seam latent-space assimilation
    methods need (encode the ensemble, assimilate in latent space,
    decode the analysis) and what autoencoder priors naturally expose.
    Implementations: POD/EOF projections, autoencoders, or an identity
    map for testing.

    Members:
        encode(state): Full state -> latent coordinates ``z``.
        decode(coords): Latent coordinates -> full state.
    """

    def encode(self, state: Any) -> Any: ...

    def decode(self, coords: Any) -> Any: ...

`ObservationNoise` ¶

Bases: Protocol

An observation-error model that exposes covariance and sampling.

Richer than the opaque obs_err_cov operand of AnalysisStep: deterministic analyses (ETKF, 3D-Var) need only the covariance, while stochastic ones (perturbed-observation EnKF) and synthetic-data generation also need draws from the error distribution. Bundling both keeps the two views consistent by construction.

Members

covariance(): The error covariance — a matrix, a linear operator, or any object the consuming AnalysisStep understands. sample(key, shape): Draw error samples of the given batch shape. key is the RNG seed/key (e.g. a JAX PRNG key); implementations on a global RNG may ignore it.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class ObservationNoise(Protocol):
    """An observation-error model that exposes covariance *and* sampling.

    Richer than the opaque ``obs_err_cov`` operand of `AnalysisStep`:
    deterministic analyses (ETKF, 3D-Var) need only the covariance,
    while stochastic ones (perturbed-observation EnKF) and synthetic-data
    generation also need draws from the error distribution. Bundling
    both keeps the two views consistent by construction.

    Members:
        covariance(): The error covariance — a matrix, a linear
            operator, or any object the consuming `AnalysisStep`
            understands.
        sample(key, shape): Draw error samples of the given batch
            shape. ``key`` is the RNG seed/key (e.g. a JAX PRNG key);
            implementations on a global RNG may ignore it.
    """

    def covariance(self) -> Any: ...

    def sample(self, key: Any, shape: tuple[int, ...]) -> Any: ...

`ObservationOperator` ¶

Bases: Protocol

Map model state → predicted observations.

The H operator in classical data-assimilation notation: H(x) produces "what would the observations look like if the state were x?". The innovation in DA is then obs - H(forecast).

Members

call(state): Return predicted observations for state. linearize(state): Optional tangent-linear operator at state (returns a callable / matrix). Implementations that don't expose a linearisation may raise NotImplementedError.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class ObservationOperator(Protocol):
    """Map model state → predicted observations.

    The H operator in classical data-assimilation notation:
    ``H(x)`` produces "what would the observations look like if the
    state were ``x``?". The innovation in DA is then ``obs - H(forecast)``.

    Members:
        __call__(state): Return predicted observations for ``state``.
        linearize(state): Optional tangent-linear operator at
            ``state`` (returns a callable / matrix). Implementations
            that don't expose a linearisation may raise
            ``NotImplementedError``.
    """

    def __call__(self, state: Any) -> Any: ...

    def linearize(self, state: Any) -> Any: ...

`ReducedBasis` ¶

Bases: Protocol

Map a reduced control vector to a state increment (with a prior).

Variational inverse problems (4D-Var on SSH, …) parameterise the control vector X not in full state space but in a reduced basis Φ (Gaussian RBF, wavelet, MIOST), each coefficient carrying a diagonal prior Q. The basis maps coefficients to a state increment, and its prior supplies the background term ½ Xᵀ Q⁻¹ X.

Concrete implementations live in algorithm libraries (or pipekit_cycle.basis) and satisfy this protocol structurally.

Members

operg(t, X, state): Apply the basis — coefficients X to a state increment Φ(t) X at time t. prior_inv(X): Apply the diagonal prior inverse Q⁻¹ X (the 4D-Var background term operand). nbasis: Length of the control vector X.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class ReducedBasis(Protocol):
    """Map a reduced control vector to a state increment (with a prior).

    Variational inverse problems (4D-Var on SSH, …) parameterise the
    control vector ``X`` not in full state space but in a reduced basis
    ``Φ`` (Gaussian RBF, wavelet, MIOST), each coefficient carrying a
    diagonal prior ``Q``. The basis maps coefficients to a state
    increment, and its prior supplies the background term ``½ Xᵀ Q⁻¹ X``.

    Concrete implementations live in algorithm libraries (or
    `pipekit_cycle.basis`) and satisfy this protocol structurally.

    Members:
        operg(t, X, state): Apply the basis — coefficients ``X`` to a
            state increment ``Φ(t) X`` at time ``t``.
        prior_inv(X): Apply the diagonal prior inverse ``Q⁻¹ X`` (the
            4D-Var background term operand).
        nbasis: Length of the control vector ``X``.
    """

    def operg(self, t: float, X: Any, state: Any = None) -> Any: ...

    def prior_inv(self, X: Any) -> Any: ...

    @property
    def nbasis(self) -> int: ...

`ReducedOrderModel` ¶

Bases: Protocol

A Galerkin / POD reduced-order model: encode/decode + latent dynamics.

A Galerkin ROM projects the full state onto a low-dimensional trial subspace and advances the reduced coordinates with the subspace-projected dynamics. Bundling the projection maps with the reduced step lets one object serve both uses:

As a fast ForwardModel on full states — wrap x -> decode(step(encode(x), dt)).
For data assimilation in reduced space — encode the state (and observations), assimilate the small latent vector, then decode.

Concrete ROMs (POD/EOF, balanced truncation, autoencoder latents) live in algorithm libraries and satisfy this structurally; the latent dynamics may itself be a ForwardModel or a learned emulator.

Members

encode(state): Full state -> reduced coordinates z (e.g. z = Psi^T (x - x_ref) for a (Petrov-)Galerkin trial/test basis). decode(coords): Reduced coordinates -> full state (x = x_ref + Phi z). step(coords, dt): Advance the reduced coordinates by dt with the projected reduced dynamics. latent_dim: Dimension of the reduced coordinate vector.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class ReducedOrderModel(Protocol):
    """A Galerkin / POD reduced-order model: encode/decode + latent dynamics.

    A Galerkin ROM projects the full state onto a low-dimensional trial
    subspace and advances the *reduced* coordinates with the
    subspace-projected dynamics. Bundling the projection maps with the
    reduced step lets one object serve both uses:

    - **As a fast `ForwardModel`** on full states — wrap
      ``x -> decode(step(encode(x), dt))``.
    - **For data assimilation in reduced space** — ``encode`` the state (and
      observations), assimilate the small latent vector, then ``decode``.

    Concrete ROMs (POD/EOF, balanced truncation, autoencoder latents) live
    in algorithm libraries and satisfy this structurally; the latent
    dynamics may itself be a `ForwardModel` or a learned emulator.

    Members:
        encode(state): Full state -> reduced coordinates ``z`` (e.g.
            ``z = Psi^T (x - x_ref)`` for a (Petrov-)Galerkin trial/test
            basis).
        decode(coords): Reduced coordinates -> full state
            (``x = x_ref + Phi z``).
        step(coords, dt): Advance the reduced coordinates by ``dt`` with the
            projected reduced dynamics.
        latent_dim: Dimension of the reduced coordinate vector.
    """

    def encode(self, state: Any) -> Any: ...

    def decode(self, coords: Any) -> Any: ...

    def step(self, coords: Any, dt: float) -> Any: ...

    @property
    def latent_dim(self) -> int: ...

`StepScheduler` ¶

Bases: Protocol

Step-size policy for an IterativeProcess.

Iterative ensemble methods (EKI, EKS, UKI) advance in pseudo-time; the scheduler chooses each increment — fixed, data-misfit controlled, or stability-bounded — from the current iteration state.

Members

get_dt(state): Return the scalar pseudo-time step for the next iteration given the process state (particles, forward evaluations, misfit, …).

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class StepScheduler(Protocol):
    """Step-size policy for an `IterativeProcess`.

    Iterative ensemble methods (EKI, EKS, UKI) advance in pseudo-time;
    the scheduler chooses each increment — fixed, data-misfit
    controlled, or stability-bounded — from the current iteration
    state.

    Members:
        get_dt(state): Return the scalar pseudo-time step for the next
            iteration given the process state (particles, forward
            evaluations, misfit, …).
    """

    def get_dt(self, state: Any) -> Any: ...

`SupportsAdjoint` ¶

Bases: Protocol

A differentiable component that exposes its adjoint strategy.

Components whose forward pass is a rollout (a ForwardModel backed by an ODE/SDE solve, a differentiable analysis loop, …) advertise how gradients flow through them via a declarative spec from pipekit_cycle.adjoints. Cycles never interpret the spec — they pass components through untouched — so strategy selection is pure configuration.

Members

adjoint: The current adjoint spec (an pipekit_cycle.adjoints.AdjointSpec member, or any object the component's execution layer understands). with_adjoint(spec): Return a copy of the component configured with spec; the original is unchanged.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class SupportsAdjoint(Protocol):
    """A differentiable component that exposes its adjoint strategy.

    Components whose forward pass is a rollout (a `ForwardModel`
    backed by an ODE/SDE solve, a differentiable analysis loop, …)
    advertise *how gradients flow* through them via a declarative spec
    from `pipekit_cycle.adjoints`. Cycles never interpret the spec —
    they pass components through untouched — so strategy selection is
    pure configuration.

    Members:
        adjoint: The current adjoint spec (an
            `pipekit_cycle.adjoints.AdjointSpec` member, or any object
            the component's execution layer understands).
        with_adjoint(spec): Return a copy of the component configured
            with ``spec``; the original is unchanged.
    """

    @property
    def adjoint(self) -> Any: ...

    def with_adjoint(self, spec: Any) -> SupportsAdjoint: ...

`TangentLinearModel` ¶

Bases: Protocol

Tangent-linear and adjoint of a ForwardModel's dynamics.

The M' / M* seam for methods that need the linearised model explicitly — incremental 4D-Var inner loops and reduced-rank Kalman error-subspace propagation — rather than relying on automatic differentiation of ForwardModel.step. Implementations are supplied by algorithm libraries (or wrap jax.jvp / jax.vjp of a differentiable model).

Members

tangent(state, dx, dt): Apply the tangent-linear M'(state) to a state perturbation dx, advancing it by dt. adjoint(state, dz, dt): Apply the adjoint M*(state) to an adjoint (dual) variable dz over dt.

Source code in packages/pipekit-cycle/src/pipekit_cycle/protocols.py

@runtime_checkable
class TangentLinearModel(Protocol):
    """Tangent-linear and adjoint of a `ForwardModel`'s dynamics.

    The ``M'`` / ``M*`` seam for methods that need the linearised model
    explicitly — incremental 4D-Var inner loops and reduced-rank Kalman
    error-subspace propagation — rather than relying on automatic
    differentiation of `ForwardModel.step`. Implementations are supplied by
    algorithm libraries (or wrap ``jax.jvp`` / ``jax.vjp`` of a
    differentiable model).

    Members:
        tangent(state, dx, dt): Apply the tangent-linear ``M'(state)`` to a
            state perturbation ``dx``, advancing it by ``dt``.
        adjoint(state, dz, dt): Apply the adjoint ``M*(state)`` to an
            adjoint (dual) variable ``dz`` over ``dt``.
    """

    def tangent(self, state: Any, dx: Any, dt: float) -> Any: ...

    def adjoint(self, state: Any, dz: Any, dt: float) -> Any: ...

Adjoint strategy specs¶

A shared, declarative vocabulary for how gradients flow through the three nested rollouts of a differentiable DA system — the dynamics integration, the inner variational solve, and the assimilation cycle. Each execution layer interprets a spec into its native mechanism (diffrax adjoints via pipekit_jax.to_diffrax_adjoint, optimistix adjoints in vardax, scan transforms in filterax):

Spec	dynamics (diffrax)	inner solve (optimistix)	cycle (scan)
`DirectAdjoint`	store-everything	unrolled	plain scan — exact, O(T) memory
`RecursiveCheckpointAdjoint`	treeverse ✅ default	checkpointed	checkpointed scan — exact, O(√T)
`TruncatedAdjoint(k)`	—	k-step / one-step	truncated scan — biased, O(1) in T, chaos-tolerant
`ImplicitAdjoint`	steady-state only	IFT at the fixed point	—
`BacksolveAdjoint`	continuous adjoint ⚠ unstable for chaotic/stiff dynamics	—	—

`pipekit_cycle.adjoints` ¶

Adjoint strategy specs — a shared vocabulary for gradient flow.

Differentiating a data-assimilation system involves up to three nested rollouts: the dynamics integration (one forecast window), the inner variational solve (4D-Var minimisation), and the outer assimilation cycle. Each layer realises reverse-mode differentiation through a different mechanism (diffrax adjoints, optimistix adjoints, scan transforms), but the strategy — how gradients should flow — is the same small set of concepts everywhere.

The classes here are that vocabulary: plain frozen dataclasses, pure Python, serialisable like any other pipekit config knob. They carry no behaviour; each execution layer interprets a spec into its native mechanism (see pipekit_jax.to_diffrax_adjoint for the diffrax mapping, and the algorithm libraries — filterax, vardax — for the cycle- and solver-layer interpretations). Names mirror diffrax's adjoint classes where the concepts coincide, so the vocabulary is familiar; TruncatedAdjoint is the one concept diffrax does not ship.

Strategy-by-layer summary:

Spec	dynamics	inner solve	cycle
`DirectAdjoint`	store all	unrolled	plain scan
`RecursiveCheckpointAdjoint`	treeverse	checkpointed	ckpt. scan
`TruncatedAdjoint`	—	k-step	trunc. scan
`ImplicitAdjoint`	steady-state	IFT	—
`BacksolveAdjoint`	continuous	—	—

`AdjointSpec = DirectAdjoint | RecursiveCheckpointAdjoint | TruncatedAdjoint | ImplicitAdjoint | BacksolveAdjoint` `module-attribute` ¶

Union of the adjoint strategy specs.

`BacksolveAdjoint` `dataclass` ¶

Continuous adjoint: solve the adjoint ODE backwards (O(1) memory).

Warning

Optimise-then-discretise — the result is not the exact gradient of the discrete forward solve, and reconstructing the forward trajectory by integrating backwards diverges for chaotic or stiff dynamics (the usual case in DA). Provided for completeness at the dynamics layer; never a default.

Source code in packages/pipekit-cycle/src/pipekit_cycle/adjoints.py

@dataclass(frozen=True)
class BacksolveAdjoint:
    """Continuous adjoint: solve the adjoint ODE backwards (O(1) memory).

    Warning:
        Optimise-then-discretise — the result is *not* the exact
        gradient of the discrete forward solve, and reconstructing the
        forward trajectory by integrating backwards diverges for
        chaotic or stiff dynamics (the usual case in DA). Provided for
        completeness at the dynamics layer; never a default.
    """

`DirectAdjoint` `dataclass` ¶

Exact gradients by storing every step of the rollout.

Reverse-mode memory grows linearly with the rollout length; the right choice for short rollouts where recomputation overhead is not worth paying.

Source code in packages/pipekit-cycle/src/pipekit_cycle/adjoints.py

@dataclass(frozen=True)
class DirectAdjoint:
    """Exact gradients by storing every step of the rollout.

    Reverse-mode memory grows linearly with the rollout length; the
    right choice for short rollouts where recomputation overhead is
    not worth paying.
    """

`ImplicitAdjoint` `dataclass` ¶

Differentiate a converged fixed point via the implicit function theorem.

O(1) memory and independent of the solve path, at the price of one linear solve with the (Gauss-Newton) Hessian. Only sound when the inner solve actually converged — an unconverged iterate gets the gradient of a point it never reached. Applies to the inner-solve layer (and steady-state dynamics solves); meaningless for a non-fixed-point cycle.

Source code in packages/pipekit-cycle/src/pipekit_cycle/adjoints.py

@dataclass(frozen=True)
class ImplicitAdjoint:
    """Differentiate a converged fixed point via the implicit function theorem.

    O(1) memory and independent of the solve path, at the price of one
    linear solve with the (Gauss-Newton) Hessian. Only sound when the
    inner solve actually converged — an unconverged iterate gets the
    gradient of a point it never reached. Applies to the inner-solve
    layer (and steady-state dynamics solves); meaningless for a
    non-fixed-point cycle.
    """

`RecursiveCheckpointAdjoint` `dataclass` ¶

Exact gradients with treeverse checkpointing (the recommended default).

Stores only checkpoints snapshots and recomputes segments during the backward pass (Griewank 1992): memory is logarithmic to square-root in the rollout length at the cost of extra forward work. checkpoints=None lets the interpreter pick its default schedule.

Attributes:

Name	Type	Description
`checkpoints`	`int \| None`	Number of checkpoints to store, or `None` for the interpreting layer's default schedule.

Source code in packages/pipekit-cycle/src/pipekit_cycle/adjoints.py

@dataclass(frozen=True)
class RecursiveCheckpointAdjoint:
    """Exact gradients with treeverse checkpointing (the recommended default).

    Stores only ``checkpoints`` snapshots and recomputes segments
    during the backward pass (Griewank 1992): memory is logarithmic to
    square-root in the rollout length at the cost of extra forward
    work. ``checkpoints=None`` lets the interpreter pick its default
    schedule.

    Attributes:
        checkpoints: Number of checkpoints to store, or ``None`` for
            the interpreting layer's default schedule.
    """

    checkpoints: int | None = None

`TruncatedAdjoint` `dataclass` ¶

Biased gradients: flow stops more than k steps back.

Steps before the trailing-k window run under stop_gradient, giving O(1) memory in the rollout length. k=1 is the one-step adjoint (Bolte et al. 2023) — exact when the rollout has converged to a fixed point, and the ROAD-EnKF training regime when applied at the cycle layer. For chaotic dynamics, truncation is gradient regularisation: the exact adjoint norm grows exponentially with the horizon, so the unbiased gradient is useless at long horizons regardless of how it is computed.

Attributes:

Name	Type	Description
`k`	`int`	Number of trailing steps that propagate gradients. Must be at least 1.

Source code in packages/pipekit-cycle/src/pipekit_cycle/adjoints.py

@dataclass(frozen=True)
class TruncatedAdjoint:
    """Biased gradients: flow stops more than ``k`` steps back.

    Steps before the trailing-``k`` window run under ``stop_gradient``,
    giving O(1) memory in the rollout length. ``k=1`` is the one-step
    adjoint (Bolte et al. 2023) — *exact* when the rollout has
    converged to a fixed point, and the ROAD-EnKF training regime when
    applied at the cycle layer. For chaotic dynamics, truncation is
    gradient *regularisation*: the exact adjoint norm grows
    exponentially with the horizon, so the unbiased gradient is
    useless at long horizons regardless of how it is computed.

    Attributes:
        k: Number of trailing steps that propagate gradients. Must be
            at least 1.
    """

    k: int = 1

Data-assimilation cycles¶

`pipekit_cycle.da` ¶

Data-assimilation cycles.

Three composable shapes:

DACycle — classic forecast → observe → analyse loop.
EnsembleDACycle — same shape, but the carrier is an ensemble of members; the analysis step is an ensemble update (EnKF, ETKF, …).
SmootherCycle — windowed 4D-Var-shaped variant; the analysis step receives the full window of observations at once.

Each DACycle returns (analysis_carrier, state). When save_history=True the per-cycle (forecast, analysis) pairs accumulate on self.history.

See master plan Report 10, section 2.4.

`DACycle` ¶

Bases: StatefulOperator

Classic forecast-analysis cycle.

Each call advances n_steps cycles. Within each cycle:

Forecast: forecast = forward_model.step(state_carrier, dt).
Observation: obs = obs_source(step_index) if obs_source is set, else skip the update.
Analyse: carrier = analysis_step(forecast, obs, obs_op=obs_op, obs_err_cov=state.obs_err_cov).

The CarryState (typically DAState) carries obs_err_cov and bookkeeping (t, cycle_count). At least: any object with attributes t, cycle_count, obs_err_cov works — duck-typed.

Parameters:

Name	Type	Description	Default
`forward_model`	`ForwardModel`	Object satisfying `ForwardModel`.	required
`obs_op`	`ObservationOperator`	Object satisfying `ObservationOperator`.	required
`analysis_step`	`AnalysisStep`	Object satisfying `AnalysisStep`.	required
`obs_source`	`Operator \| None`	Optional `Operator` invoked as `obs_source(step_index)` to obtain observations for the current cycle. If `None`, the analysis step is skipped and the forecast is propagated forward.	`None`
`n_steps`	`int`	Number of cycles per call.	`1`
`save_history`	`bool`	Append `(forecast, analysis)` pairs to `self.history` each cycle.	`True`

Returns from _apply(carrier, state): (analysis_carrier, updated_state).

Raises:

Type	Description
`TypeError`	If a component doesn't satisfy its protocol (`ForwardModel` / `ObservationOperator` / `AnalysisStep`), or `obs_source` is neither an `Operator` nor `None`.
`ValueError`	If `n_steps < 0`.

Source code in packages/pipekit-cycle/src/pipekit_cycle/da.py

class DACycle(StatefulOperator):
    """Classic forecast-analysis cycle.

    Each call advances ``n_steps`` cycles. Within each cycle:

    1. **Forecast**: ``forecast = forward_model.step(state_carrier, dt)``.
    2. **Observation**: ``obs = obs_source(step_index)`` if
       ``obs_source`` is set, else skip the update.
    3. **Analyse**: ``carrier = analysis_step(forecast, obs,
       obs_op=obs_op, obs_err_cov=state.obs_err_cov)``.

    The CarryState (typically `DAState`) carries
    ``obs_err_cov`` and bookkeeping (``t``, ``cycle_count``). At
    least: any object with attributes ``t``, ``cycle_count``,
    ``obs_err_cov`` works — duck-typed.

    Args:
        forward_model: Object satisfying `ForwardModel`.
        obs_op: Object satisfying `ObservationOperator`.
        analysis_step: Object satisfying `AnalysisStep`.
        obs_source: Optional `Operator` invoked as
            ``obs_source(step_index)`` to obtain observations for the
            current cycle. If ``None``, the analysis step is skipped
            and the forecast is propagated forward.
        n_steps: Number of cycles per call.
        save_history: Append ``(forecast, analysis)`` pairs to
            ``self.history`` each cycle.

    Returns from ``_apply(carrier, state)``:
        ``(analysis_carrier, updated_state)``.

    Raises:
        TypeError: If a component doesn't satisfy its protocol
            (`ForwardModel` / `ObservationOperator` / `AnalysisStep`),
            or ``obs_source`` is neither an `Operator` nor ``None``.
        ValueError: If ``n_steps < 0``.
    """

    __config_mixin_auto__ = False

    def __init__(
        self,
        forward_model: ForwardModel,
        obs_op: ObservationOperator,
        analysis_step: AnalysisStep,
        obs_source: Operator | None = None,
        n_steps: int = 1,
        save_history: bool = True,
    ) -> None:
        if not isinstance(forward_model, ForwardModel):
            raise TypeError(
                "DACycle.forward_model must satisfy ForwardModel, "
                f"got {type(forward_model).__name__}."
            )
        if not isinstance(obs_op, ObservationOperator):
            raise TypeError(
                "DACycle.obs_op must satisfy ObservationOperator, "
                f"got {type(obs_op).__name__}."
            )
        if not isinstance(analysis_step, AnalysisStep):
            raise TypeError(
                "DACycle.analysis_step must satisfy AnalysisStep, "
                f"got {type(analysis_step).__name__}."
            )
        if obs_source is not None and not isinstance(obs_source, Operator):
            raise TypeError(
                "DACycle.obs_source must be an Operator or None, "
                f"got {type(obs_source).__name__}."
            )
        if n_steps < 0:
            raise ValueError(f"DACycle.n_steps must be >= 0, got {n_steps}.")
        self.forward_model = forward_model
        self.obs_op = obs_op
        self.analysis_step = analysis_step
        self.obs_source = obs_source
        self.n_steps = n_steps
        self.save_history = save_history
        self.history: list[tuple[Any, Any]] = []

    def _apply(self, carrier: Any, state: Any) -> tuple[Any, Any]:
        self.history = []
        for i in range(self.n_steps):
            forecast = self.forward_model.step(carrier, self.forward_model.dt)
            if self.obs_source is not None:
                obs = self.obs_source(i)
                carrier = self.analysis_step(
                    forecast,
                    obs,
                    obs_op=self.obs_op,
                    obs_err_cov=getattr(state, "obs_err_cov", None),
                )
            else:
                carrier = forecast
            if self.save_history:
                self.history.append((forecast, carrier))
            state = _advance_da_state(state, dt=self.forward_model.dt)
        return carrier, state

    def get_config(self) -> dict[str, Any]:
        return {
            "forward_model": _component_config(self.forward_model),
            "obs_op": _component_config(self.obs_op),
            "analysis_step": _component_config(self.analysis_step),
            "obs_source": _component_config(self.obs_source),
            "n_steps": self.n_steps,
            "save_history": self.save_history,
        }

`EnsembleDACycle` ¶

Bases: StatefulOperator

Ensemble-based DA — EnKF, ETKF, LETKF, particle filters, …

Differences from DACycle:

The carrier is a list (or tuple) of n_members ensemble members. Each member is propagated through the forward model independently.
The analysis step receives the full ensemble (forecast_members) plus observations; it returns the updated ensemble.

Parameters:

Name	Type	Description	Default
`forward_model`	`ForwardModel`	Same protocol as `DACycle`. Applied per-member.	required
`obs_op`	`ObservationOperator`	Same protocol as `DACycle`.	required
`analysis_step`	`AnalysisStep`	An ensemble analysis (e.g. filterx.EnKFAnalysis). Receives `(forecast_members, obs, obs_op=..., obs_err_cov=...)`.	required
`obs_source`	`Operator \| None`	Optional observation source.	`None`
`n_steps`	`int`	Number of cycles per call.	`1`
`n_members`	`int`	Expected ensemble size (keyword-only; length-checked at call time).	required

Raises:

Type	Description
`TypeError`	If a component doesn't satisfy its protocol, or `obs_source` is neither an `Operator` nor `None`.
`ValueError`	If `n_members < 1` or `n_steps < 0` (from the constructor), or if the input ensemble's length differs from `n_members` (at call time).

Source code in packages/pipekit-cycle/src/pipekit_cycle/da.py

class EnsembleDACycle(StatefulOperator):
    """Ensemble-based DA — EnKF, ETKF, LETKF, particle filters, …

    Differences from `DACycle`:

    - The carrier is a list (or tuple) of ``n_members`` ensemble
      members. Each member is propagated through the forward model
      independently.
    - The analysis step receives the full ensemble (``forecast_members``)
      plus observations; it returns the updated ensemble.

    Args:
        forward_model: Same protocol as `DACycle`. Applied per-member.
        obs_op: Same protocol as `DACycle`.
        analysis_step: An ensemble analysis (e.g. filterx.EnKFAnalysis).
            Receives ``(forecast_members, obs, obs_op=..., obs_err_cov=...)``.
        obs_source: Optional observation source.
        n_steps: Number of cycles per call.
        n_members: Expected ensemble size (keyword-only;
            length-checked at call time).

    Raises:
        TypeError: If a component doesn't satisfy its protocol, or
            ``obs_source`` is neither an `Operator` nor ``None``.
        ValueError: If ``n_members < 1`` or ``n_steps < 0`` (from the
            constructor), or if the input ensemble's length differs
            from ``n_members`` (at call time).
    """

    __config_mixin_auto__ = False

    def __init__(
        self,
        forward_model: ForwardModel,
        obs_op: ObservationOperator,
        analysis_step: AnalysisStep,
        obs_source: Operator | None = None,
        n_steps: int = 1,
        *,
        n_members: int,
    ) -> None:
        if not isinstance(forward_model, ForwardModel):
            raise TypeError("EnsembleDACycle.forward_model must satisfy ForwardModel.")
        if not isinstance(obs_op, ObservationOperator):
            raise TypeError("EnsembleDACycle.obs_op must satisfy ObservationOperator.")
        if not isinstance(analysis_step, AnalysisStep):
            raise TypeError("EnsembleDACycle.analysis_step must satisfy AnalysisStep.")
        if obs_source is not None and not isinstance(obs_source, Operator):
            raise TypeError("EnsembleDACycle.obs_source must be Operator or None.")
        if n_members < 1:
            raise ValueError(
                f"EnsembleDACycle.n_members must be >= 1, got {n_members}."
            )
        if n_steps < 0:
            raise ValueError(f"EnsembleDACycle.n_steps must be >= 0, got {n_steps}.")
        self.forward_model = forward_model
        self.obs_op = obs_op
        self.analysis_step = analysis_step
        self.obs_source = obs_source
        self.n_steps = n_steps
        self.n_members = n_members

    def _apply(  # ty: ignore[invalid-method-override]
        self,
        members: list[Any],
        state: Any,
    ) -> tuple[list[Any], Any]:
        if len(members) != self.n_members:
            raise ValueError(
                f"EnsembleDACycle expected {self.n_members} members, "
                f"got {len(members)}."
            )
        for i in range(self.n_steps):
            forecasts = [
                self.forward_model.step(m, self.forward_model.dt) for m in members
            ]
            if self.obs_source is not None:
                obs = self.obs_source(i)
                members = list(
                    self.analysis_step(
                        forecasts,
                        obs,
                        obs_op=self.obs_op,
                        obs_err_cov=getattr(state, "obs_err_cov", None),
                    )
                )
            else:
                members = forecasts
            state = _advance_da_state(state, dt=self.forward_model.dt)
        return members, state

    def get_config(self) -> dict[str, Any]:
        return {
            "forward_model": _component_config(self.forward_model),
            "obs_op": _component_config(self.obs_op),
            "analysis_step": _component_config(self.analysis_step),
            "obs_source": _component_config(self.obs_source),
            "n_steps": self.n_steps,
            "n_members": self.n_members,
        }

`SmootherCycle` ¶

Bases: StatefulOperator

4D-Var-shaped windowed smoother.

Each call processes window forward steps, accumulating observations, then runs the analysis step on the whole window at once. This matches the 4D-Var control-variable formulation where the cost function is minimised over the full window jointly.

Parameters:

Name	Type	Description	Default
`forward_model`	`ForwardModel`	Same protocol as `DACycle`.	required
`obs_op`	`ObservationOperator`	Same protocol as `DACycle`.	required
`analysis_step`	`AnalysisStep`	A smoother / variational solver. Called as `analysis_step(window_forecasts, window_obs, obs_op=..., obs_err_cov=...)` where the window arguments are lists.	required
`window`	`int`	Number of forward steps per analysis window.	required
`stride`	`int`	Number of windows to take per call (default 1).	`1`
`obs_source`	`Operator \| None`	Optional observation source; invoked per step inside the window. `None` skips the analysis step.	`None`

Raises:

Type	Description
`TypeError`	If a component doesn't satisfy its protocol, or `obs_source` is neither an `Operator` nor `None`.
`ValueError`	If `window < 1` or `stride < 1`.

Source code in packages/pipekit-cycle/src/pipekit_cycle/da.py

class SmootherCycle(StatefulOperator):
    """4D-Var-shaped windowed smoother.

    Each call processes ``window`` forward steps, accumulating
    observations, then runs the analysis step on the whole window
    at once. This matches the 4D-Var control-variable formulation
    where the cost function is minimised over the full window
    jointly.

    Args:
        forward_model: Same protocol as `DACycle`.
        obs_op: Same protocol as `DACycle`.
        analysis_step: A smoother / variational solver. Called as
            ``analysis_step(window_forecasts, window_obs,
            obs_op=..., obs_err_cov=...)`` where the window arguments
            are lists.
        window: Number of forward steps per analysis window.
        stride: Number of windows to take per call (default 1).
        obs_source: Optional observation source; invoked per step
            inside the window. ``None`` skips the analysis step.

    Raises:
        TypeError: If a component doesn't satisfy its protocol, or
            ``obs_source`` is neither an `Operator` nor ``None``.
        ValueError: If ``window < 1`` or ``stride < 1``.
    """

    __config_mixin_auto__ = False

    def __init__(
        self,
        forward_model: ForwardModel,
        obs_op: ObservationOperator,
        analysis_step: AnalysisStep,
        window: int,
        stride: int = 1,
        obs_source: Operator | None = None,
    ) -> None:
        if not isinstance(forward_model, ForwardModel):
            raise TypeError("SmootherCycle.forward_model must satisfy ForwardModel.")
        if not isinstance(obs_op, ObservationOperator):
            raise TypeError("SmootherCycle.obs_op must satisfy ObservationOperator.")
        if not isinstance(analysis_step, AnalysisStep):
            raise TypeError("SmootherCycle.analysis_step must satisfy AnalysisStep.")
        if obs_source is not None and not isinstance(obs_source, Operator):
            raise TypeError("SmootherCycle.obs_source must be Operator or None.")
        if window < 1:
            raise ValueError(f"SmootherCycle.window must be >= 1, got {window}.")
        if stride < 1:
            raise ValueError(f"SmootherCycle.stride must be >= 1, got {stride}.")
        self.forward_model = forward_model
        self.obs_op = obs_op
        self.analysis_step = analysis_step
        self.obs_source = obs_source
        self.window = window
        self.stride = stride

    def _apply(self, carrier: Any, state: Any) -> tuple[Any, Any]:
        global_step = 0
        for _ in range(self.stride):
            forecasts: list[Any] = []
            observations: list[Any] = []
            for _ in range(self.window):
                carrier = self.forward_model.step(carrier, self.forward_model.dt)
                forecasts.append(carrier)
                if self.obs_source is not None:
                    observations.append(self.obs_source(global_step))
                global_step += 1
                state = _advance_da_state(state, dt=self.forward_model.dt)
            if self.obs_source is not None:
                carrier = self.analysis_step(
                    forecasts,
                    observations,
                    obs_op=self.obs_op,
                    obs_err_cov=getattr(state, "obs_err_cov", None),
                )
        return carrier, state

    def get_config(self) -> dict[str, Any]:
        return {
            "forward_model": _component_config(self.forward_model),
            "obs_op": _component_config(self.obs_op),
            "analysis_step": _component_config(self.analysis_step),
            "obs_source": _component_config(self.obs_source),
            "window": self.window,
            "stride": self.stride,
        }

Observation operators¶

`pipekit_cycle.obs` ¶

Generic observation operators.

Each is a pipekit.Operator that also satisfies the ObservationOperator Protocol so it can be passed wherever a protocol-typed slot is expected (e.g. into DACycle.obs_op).

IdentityObs — direct measurement (H = I).
LinearObs(H) — fixed linear map, H @ state. Carrier-agnostic via duck-typed @.
CallableObs(fn) — arbitrary callable (e.g. radiative-transfer forward). Carries forbid_in_yaml = True.
CompositeObs(components) — h_k ∘ … ∘ h_1 applied left-to-right.

See master plan Report 10, section 2.5.

`CallableObs` ¶

Bases: Operator

Arbitrary callable observation operator.

The escape hatch for nonlinear observation operators (radiative transfer, instrument-response convolution, …). Carries forbid_in_yaml = True because the callable doesn't round-trip through JSON / YAML.

Parameters:

Name	Type	Description	Default
`fn`	`Callable[[Any], Any]`	Callable `state → predicted_obs`.	required
`linearize_fn`	`Callable[[Any], Any] \| None`	Optional callable `state → linear operator`. If absent, `linearize` raises `NotImplementedError`.	`None`
`name`	`str \| None`	Optional name for `__repr__`.	`None`

Source code in packages/pipekit-cycle/src/pipekit_cycle/obs.py

class CallableObs(Operator):
    """Arbitrary callable observation operator.

    The escape hatch for nonlinear observation operators (radiative
    transfer, instrument-response convolution, …). Carries
    ``forbid_in_yaml = True`` because the callable doesn't round-trip
    through JSON / YAML.

    Args:
        fn: Callable ``state → predicted_obs``.
        linearize_fn: Optional callable ``state → linear operator``.
            If absent, `linearize` raises `NotImplementedError`.
        name: Optional name for ``__repr__``.
    """

    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(
        self,
        fn: Callable[[Any], Any],
        linearize_fn: Callable[[Any], Any] | None = None,
        name: str | None = None,
    ) -> None:
        self.fn = fn
        self.linearize_fn = linearize_fn
        self.name = name or getattr(fn, "__name__", "CallableObs")

    def _apply(self, state: Any) -> Any:
        return self.fn(state)

    def linearize(self, state: Any) -> Any:
        if self.linearize_fn is None:
            raise NotImplementedError(f"{self.name}: no linearisation supplied.")
        return self.linearize_fn(state)

    def get_config(self) -> dict[str, Any]:
        return {"name": self.name}

    def __repr__(self) -> str:
        return f"CallableObs(name={self.name!r})"

`CompositeObs` ¶

Bases: Operator

Compose multiple observation operators: H_k ∘ … ∘ H_1.

Apply components[0] first, then components[1], etc.

Parameters:

Name	Type	Description	Default
`components`	`tuple[Operator, ...]`	One or more `Operator` instances. Each must accept the previous component's output as its input.	required

Raises:

Type	Description
`TypeError`	if any component isn't an `Operator`.
`ValueError`	if `components` is empty.

Source code in packages/pipekit-cycle/src/pipekit_cycle/obs.py

class CompositeObs(Operator):
    """Compose multiple observation operators: ``H_k ∘ … ∘ H_1``.

    Apply ``components[0]`` first, then ``components[1]``, etc.

    Args:
        components: One or more `Operator` instances. Each must accept
            the previous component's output as its input.

    Raises:
        TypeError: if any component isn't an `Operator`.
        ValueError: if ``components`` is empty.
    """

    __config_mixin_auto__ = False

    def __init__(self, components: tuple[Operator, ...]) -> None:
        if not components:
            raise ValueError("CompositeObs requires at least one component.")
        for i, op in enumerate(components):
            if not isinstance(op, Operator):
                raise TypeError(
                    f"CompositeObs.components[{i}] must be an Operator, "
                    f"got {type(op).__name__}."
                )
        self.components = tuple(components)

    def _apply(self, state: Any) -> Any:
        out = state
        for op in self.components:
            out = op(out)
        return out

    def linearize(self, state: Any) -> Any:
        """Compose component linearisations (chain rule).

        Each component must expose a `linearize`. If a non-linearisable
        component is present, raises `NotImplementedError`.
        """
        lins: list[Operator] = []
        carrier = state
        for op in self.components:
            linearize = getattr(op, "linearize", None)
            if linearize is None:
                raise NotImplementedError(
                    f"CompositeObs: component {type(op).__name__} has no "
                    "`linearize` — cannot build composite tangent linear."
                )
            lins.append(linearize(carrier))
            carrier = op(carrier)
        return CompositeObs(tuple(lins))

    def get_config(self) -> dict[str, Any]:
        return {"components": [nested_config(op) for op in self.components]}

`linearize(state)` ¶

Compose component linearisations (chain rule).

Each component must expose a linearize. If a non-linearisable component is present, raises NotImplementedError.

Source code in packages/pipekit-cycle/src/pipekit_cycle/obs.py

def linearize(self, state: Any) -> Any:
    """Compose component linearisations (chain rule).

    Each component must expose a `linearize`. If a non-linearisable
    component is present, raises `NotImplementedError`.
    """
    lins: list[Operator] = []
    carrier = state
    for op in self.components:
        linearize = getattr(op, "linearize", None)
        if linearize is None:
            raise NotImplementedError(
                f"CompositeObs: component {type(op).__name__} has no "
                "`linearize` — cannot build composite tangent linear."
            )
        lins.append(linearize(carrier))
        carrier = op(carrier)
    return CompositeObs(tuple(lins))

`IdentityObs` ¶

Bases: Operator

Direct measurement: H(x) = x.

The state is observable as-is — useful for synthetic-twin tests and as the default observation operator when one is needed but no real mapping applies.

Source code in packages/pipekit-cycle/src/pipekit_cycle/obs.py

class IdentityObs(Operator):
    """Direct measurement: ``H(x) = x``.

    The state is observable as-is — useful for synthetic-twin tests
    and as the default observation operator when one is needed but
    no real mapping applies.
    """

    def _apply(self, state: Any) -> Any:
        return state

    def linearize(self, state: Any) -> Operator:
        """Identity is its own tangent linear (``state`` is intentionally unused)."""
        return self

`linearize(state)` ¶

Identity is its own tangent linear (state is intentionally unused).

Source code in packages/pipekit-cycle/src/pipekit_cycle/obs.py

def linearize(self, state: Any) -> Operator:
    """Identity is its own tangent linear (``state`` is intentionally unused)."""
    return self

`LinearObs` ¶

Bases: Operator

Linear observation map: H(x) = H @ x.

Parameters:

Name	Type	Description	Default
`H`	`Any`	Any object supporting the `@` operator with the carrier. Typically a numpy / JAX array; pipekit-cycle stays backend-agnostic by relying on `__matmul__`.	required

Carries forbid_in_yaml = True because the matrix H is a runtime array that doesn't round-trip through JSON / YAML; get_config() reports only its shape as a debug payload.

Source code in packages/pipekit-cycle/src/pipekit_cycle/obs.py

class LinearObs(Operator):
    """Linear observation map: ``H(x) = H @ x``.

    Args:
        H: Any object supporting the ``@`` operator with the carrier.
            Typically a numpy / JAX array; pipekit-cycle stays
            backend-agnostic by relying on ``__matmul__``.

    Carries ``forbid_in_yaml = True`` because the matrix ``H`` is a
    runtime array that doesn't round-trip through JSON / YAML;
    ``get_config()`` reports only its shape as a debug payload.
    """

    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(self, H: Any) -> None:
        self.H = H

    def _apply(self, state: Any) -> Any:
        return self.H @ state

    def linearize(self, state: Any) -> LinearObs:
        """The map is already linear — returns ``self`` (``state`` unused)."""
        return self

    def get_config(self) -> dict[str, Any]:
        return {"H_shape": _shape_or_repr(self.H)}

`linearize(state)` ¶

The map is already linear — returns self (state unused).

Source code in packages/pipekit-cycle/src/pipekit_cycle/obs.py

def linearize(self, state: Any) -> LinearObs:
    """The map is already linear — returns ``self`` (``state`` unused)."""
    return self

Forward-model adapters¶

`pipekit_cycle.forward` ¶

Generic ForwardModel adapters.

Each class is a pipekit.Operator and satisfies the ForwardModel Protocol so it can be passed wherever a protocol-typed slot accepts forward models (e.g. DACycle.forward_model).

CallableForward(fn, dt) — wrap any callable as a ForwardModel.
CompositeForward(components, dt) — apply each component in turn; total advance is dt (each component sees its own portion).
NeuralForward(model_op, dt) — wrap a neural emulator operator (typically a pipekit-array.ModelOp) as a forward model.

See master plan Report 10, section 2.6.

`CallableForward` ¶

Bases: Operator

Wrap a plain callable as a forward model.

Parameters:

Name	Type	Description	Default
`fn`	`Callable[[Any, float], Any]`	Callable `(state, dt) → new_state`.	required
`dt`	`float`	Default integration step in user units (seconds, hours, …).	required
`state_signature`	`Any`	Optional `pipekit.Signature` describing the state carrier. `None` by default.	`None`
`name`	`str \| None`	Optional name for `__repr__`.	`None`

Carries forbid_in_yaml = True — the callable doesn't round-trip.

Source code in packages/pipekit-cycle/src/pipekit_cycle/forward.py

class CallableForward(Operator):
    """Wrap a plain callable as a forward model.

    Args:
        fn: Callable ``(state, dt) → new_state``.
        dt: Default integration step in user units (seconds, hours, …).
        state_signature: Optional `pipekit.Signature` describing the
            state carrier. ``None`` by default.
        name: Optional name for ``__repr__``.

    Carries ``forbid_in_yaml = True`` — the callable doesn't round-trip.
    """

    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(
        self,
        fn: Callable[[Any, float], Any],
        dt: float,
        state_signature: Any = None,
        name: str | None = None,
    ) -> None:
        self.fn = fn
        self._dt = float(dt)
        self._state_signature = state_signature
        self.name = name or getattr(fn, "__name__", "CallableForward")

    def step(self, state: Any, dt: float) -> Any:
        return self.fn(state, dt)

    def _apply(self, state: Any) -> Any:
        """Calling the operator advances by the default ``dt``."""
        return self.step(state, self._dt)

    @property
    def dt(self) -> float:
        return self._dt

    @property
    def state_signature(self) -> Any:
        return self._state_signature

    def get_config(self) -> dict[str, Any]:
        return {"dt": self._dt, "name": self.name}

    def __repr__(self) -> str:
        return f"CallableForward(name={self.name!r}, dt={self._dt})"

`CompositeForward` ¶

Bases: Operator

Compose multiple forward models — F_k ∘ … ∘ F_1.

Each component advances state by its own internal dt. The composite's dt is the sum (the total advance per step call). The components run left-to-right per step.

Parameters:

Name	Type	Description	Default
`components`	`tuple[Any, ...]`	Tuple of objects satisfying `ForwardModel`. Must be non-empty.	required

Raises:

Type	Description
`ValueError`	if `components` is empty.
`TypeError`	if any component doesn't satisfy `ForwardModel`.

Source code in packages/pipekit-cycle/src/pipekit_cycle/forward.py

class CompositeForward(Operator):
    """Compose multiple forward models — ``F_k ∘ … ∘ F_1``.

    Each component advances state by its own internal ``dt``. The
    composite's ``dt`` is the sum (the total advance per ``step`` call).
    The components run left-to-right per step.

    Args:
        components: Tuple of objects satisfying `ForwardModel`. Must
            be non-empty.

    Raises:
        ValueError: if ``components`` is empty.
        TypeError: if any component doesn't satisfy `ForwardModel`.
    """

    __config_mixin_auto__ = False

    def __init__(self, components: tuple[Any, ...]) -> None:
        if not components:
            raise ValueError("CompositeForward requires at least one component.")
        for i, c in enumerate(components):
            if not isinstance(c, ForwardModel):
                raise TypeError(
                    f"CompositeForward.components[{i}] does not satisfy "
                    "ForwardModel (needs `step`, `dt`, and `state_signature`)."
                )
        self.components = tuple(components)

    def step(self, state: Any, dt: float) -> Any:
        out = state
        for c in self.components:
            out = c.step(out, c.dt)
        return out

    def _apply(self, state: Any) -> Any:
        return self.step(state, self.dt)

    @property
    def dt(self) -> float:
        return float(sum(c.dt for c in self.components))

    @property
    def state_signature(self) -> Any:
        for c in self.components:
            sig = getattr(c, "state_signature", None)
            if sig is not None:
                return sig
        return None

    def get_config(self) -> dict[str, Any]:
        return {"components": [_component_config(c) for c in self.components]}

`NeuralForward` ¶

Bases: Operator

Wrap a neural-emulator operator (e.g. ModelOp) as a ForwardModel.

The wrapped model_op should be a pipekit.Operator that accepts the state carrier and returns the advanced state. dt is declared explicitly — neural emulators have a fixed training cadence that callers must respect.

Parameters:

Name	Type	Description	Default
`model_op`	`Operator`	The neural emulator `Operator`. Loaded by content-hash from a model registry in typical use.	required
`dt`	`float`	The (fixed) integration step the emulator was trained for.	required
`state_signature`	`Any`	Optional `pipekit.Signature` for the state.	`None`

Raises:

Type	Description
`TypeError`	if `model_op` isn't an `Operator`.

Source code in packages/pipekit-cycle/src/pipekit_cycle/forward.py

class NeuralForward(Operator):
    """Wrap a neural-emulator operator (e.g. `ModelOp`) as a `ForwardModel`.

    The wrapped ``model_op`` should be a `pipekit.Operator` that accepts
    the state carrier and returns the advanced state. ``dt`` is
    declared explicitly — neural emulators have a fixed training
    cadence that callers must respect.

    Args:
        model_op: The neural emulator `Operator`. Loaded by
            content-hash from a model registry in typical use.
        dt: The (fixed) integration step the emulator was trained for.
        state_signature: Optional `pipekit.Signature` for the state.

    Raises:
        TypeError: if ``model_op`` isn't an `Operator`.
    """

    __config_mixin_auto__ = False

    def __init__(
        self,
        model_op: Operator,
        dt: float,
        state_signature: Any = None,
    ) -> None:
        if not isinstance(model_op, Operator):
            raise TypeError(
                "NeuralForward.model_op must be an Operator, "
                f"got {type(model_op).__name__}."
            )
        self.model_op = model_op
        self._dt = float(dt)
        self._state_signature = state_signature

    def step(self, state: Any, dt: float) -> Any:
        # `dt` is accepted for protocol conformance; the emulator's
        # cadence is fixed at construction.
        return self.model_op(state)

    def _apply(self, state: Any) -> Any:
        return self.step(state, self._dt)

    @property
    def dt(self) -> float:
        return self._dt

    @property
    def state_signature(self) -> Any:
        return self._state_signature

    def get_config(self) -> dict[str, Any]:
        return {"model_op": nested_config(self.model_op), "dt": self._dt}

Carry-state subclasses¶

`pipekit_cycle.state` ¶

Canonical CarryState subclasses for the cycle patterns.

Each cycle pattern brings its own state shape — pipekit-cycle ships the three most common ones. Algorithm libraries are free to subclass pipekit.CarryState directly with their own fields.

DAState — clock + cycle bookkeeping for DACycle / EnsembleDACycle.
IterationState — residual + iteration counter for Recurrence.
WindowState — sliding-window bookkeeping for WindowedCycle / SmootherCycle.

See master plan Report 10, section 6.2.

`DAState` ¶

Bases: CarryState

State threaded through a DA cycle.

Attributes:

Name	Type	Description
`t`		Current model time (seconds, or any user-defined unit).
`cycle_count`		Number of completed cycles.
`obs_err_cov`		Observation-error covariance (handed to the analysis step). Carrier-agnostic — typically a numpy / JAX array or a covariance object from the algorithm library.
`extras`		Free-form dict for algorithm-specific state.

Source code in packages/pipekit-cycle/src/pipekit_cycle/state.py

class DAState(CarryState):
    """State threaded through a DA cycle.

    Attributes:
        t: Current model time (seconds, or any user-defined unit).
        cycle_count: Number of completed cycles.
        obs_err_cov: Observation-error covariance (handed to the
            analysis step). Carrier-agnostic — typically a numpy /
            JAX array or a covariance object from the algorithm
            library.
        extras: Free-form dict for algorithm-specific state.
    """

    def __init__(
        self,
        t: float = 0.0,
        cycle_count: int = 0,
        obs_err_cov: Any = None,
        extras: dict[str, Any] | None = None,
    ) -> None:
        self.t = t
        self.cycle_count = cycle_count
        self.obs_err_cov = obs_err_cov
        self.extras = dict(extras) if extras is not None else {}

`IterationState` ¶

Bases: CarryState

State threaded through a Recurrence.

Attributes:

Name	Type	Description
`iter_count`		Number of completed iterations.
`residual_norm`		Most recent residual norm (convergence check).
`extras`		Free-form dict for algorithm-specific bookkeeping (learning rate schedule, line-search trace, …).

Source code in packages/pipekit-cycle/src/pipekit_cycle/state.py

class IterationState(CarryState):
    """State threaded through a `Recurrence`.

    Attributes:
        iter_count: Number of completed iterations.
        residual_norm: Most recent residual norm (convergence check).
        extras: Free-form dict for algorithm-specific bookkeeping
            (learning rate schedule, line-search trace, …).
    """

    def __init__(
        self,
        iter_count: int = 0,
        residual_norm: float = float("inf"),
        extras: dict[str, Any] | None = None,
    ) -> None:
        self.iter_count = iter_count
        self.residual_norm = residual_norm
        self.extras = dict(extras) if extras is not None else {}

`WindowState` ¶

Bases: CarryState

State threaded through a windowed / smoother cycle.

Attributes:

Name	Type	Description
`window_index`		Zero-based index of the current window.
`last_analysis`		Most recent analysis carrier (the bridge between windows). `None` before the first analysis is produced.
`extras`		Free-form dict for algorithm-specific bookkeeping (windowed obs accumulator, adjoint trajectories, …).

Source code in packages/pipekit-cycle/src/pipekit_cycle/state.py

class WindowState(CarryState):
    """State threaded through a windowed / smoother cycle.

    Attributes:
        window_index: Zero-based index of the current window.
        last_analysis: Most recent analysis carrier (the bridge between
            windows). ``None`` before the first analysis is produced.
        extras: Free-form dict for algorithm-specific bookkeeping
            (windowed obs accumulator, adjoint trajectories, …).
    """

    def __init__(
        self,
        window_index: int = 0,
        last_analysis: Any = None,
        extras: dict[str, Any] | None = None,
    ) -> None:
        self.window_index = window_index
        self.last_analysis = last_analysis
        self.extras = dict(extras) if extras is not None else {}

pipekit-cycle¶

Cycle wrappers¶

pipekit_cycle.cycle ¶

Cycle ¶

EnsembleCycle ¶

Recurrence ¶

WindowedCycle ¶

Protocols¶

pipekit_cycle.protocols ¶

AnalysisStep ¶

CovarianceLocalizer ¶

EnsembleInflator ¶

ErrorSubspace ¶

ForwardModel ¶

IterativeProcess ¶

LatentMap ¶

ObservationNoise ¶

ObservationOperator ¶

ReducedBasis ¶

ReducedOrderModel ¶

StepScheduler ¶

SupportsAdjoint ¶

TangentLinearModel ¶

Adjoint strategy specs¶

pipekit_cycle.adjoints ¶

AdjointSpec = DirectAdjoint | RecursiveCheckpointAdjoint | TruncatedAdjoint | ImplicitAdjoint | BacksolveAdjoint module-attribute ¶

BacksolveAdjoint dataclass ¶

DirectAdjoint dataclass ¶

ImplicitAdjoint dataclass ¶

RecursiveCheckpointAdjoint dataclass ¶

TruncatedAdjoint dataclass ¶

Data-assimilation cycles¶

pipekit_cycle.da ¶

DACycle ¶

EnsembleDACycle ¶

SmootherCycle ¶

Observation operators¶

pipekit_cycle.obs ¶

CallableObs ¶

CompositeObs ¶

linearize(state) ¶

IdentityObs ¶

linearize(state) ¶

LinearObs ¶

linearize(state) ¶

Forward-model adapters¶

pipekit_cycle.forward ¶

CallableForward ¶

CompositeForward ¶

NeuralForward ¶

Carry-state subclasses¶

pipekit_cycle.state ¶

DAState ¶

IterationState ¶

WindowState ¶

`pipekit_cycle.cycle` ¶

`Cycle` ¶

`EnsembleCycle` ¶

`Recurrence` ¶

`WindowedCycle` ¶

`pipekit_cycle.protocols` ¶

`AnalysisStep` ¶

`CovarianceLocalizer` ¶

`EnsembleInflator` ¶

`ErrorSubspace` ¶

`ForwardModel` ¶

`IterativeProcess` ¶

`LatentMap` ¶

`ObservationNoise` ¶

`ObservationOperator` ¶

`ReducedBasis` ¶

`ReducedOrderModel` ¶

`StepScheduler` ¶

`SupportsAdjoint` ¶

`TangentLinearModel` ¶

`pipekit_cycle.adjoints` ¶

`AdjointSpec = DirectAdjoint | RecursiveCheckpointAdjoint | TruncatedAdjoint | ImplicitAdjoint | BacksolveAdjoint` `module-attribute` ¶

`BacksolveAdjoint` `dataclass` ¶

`DirectAdjoint` `dataclass` ¶

`ImplicitAdjoint` `dataclass` ¶

`RecursiveCheckpointAdjoint` `dataclass` ¶

`TruncatedAdjoint` `dataclass` ¶

`pipekit_cycle.da` ¶

`DACycle` ¶

`EnsembleDACycle` ¶

`SmootherCycle` ¶

`pipekit_cycle.obs` ¶

`CallableObs` ¶

`CompositeObs` ¶

`linearize(state)` ¶

`IdentityObs` ¶

`linearize(state)` ¶

`LinearObs` ¶

`linearize(state)` ¶

`pipekit_cycle.forward` ¶

`CallableForward` ¶

`CompositeForward` ¶

`NeuralForward` ¶

`pipekit_cycle.state` ¶

`DAState` ¶

`IterationState` ¶

`WindowState` ¶