Report 13 — `statecatalog`: catalog for model states

   Domain libraries           geotoolz │ xr_toolz
                                  ▲
                                  │
   Infrastructure        ┌─ pipekit-cycle ─┐    Reads/writes
                         │  (Report 10)   │  ──▶ statecatalog  ◄── (this report)
                         └────────────────┘
                                  ▲                  ▲
                                  │                  │
   Framework                  pipekit          ┌─ georeader ─┐
                                                │   (reads     │
                                                │   state data) │
                                                └──────────────┘

statecatalog/
  __init__.py             # public re-exports
  _src/
    slice.py              # StateSlice — the wire format (the headline)
    base.py               # StateRow, StateCatalog Protocol
    memory.py             # InMemoryStateCatalog (in-memory pandas-backed)
    duckdb_backend.py     # DuckDBStateCatalog (extras-gated)
    queries.py            # Common query patterns: at_time, lead_window, lineage, ensemble
    streaming.py          # Streaming-write infrastructure (parallel to geocatalog's)
    parquet.py            # GeoParquet round-trip (with state-extension fields)
    integrations/
      pipekit.py          # CatalogReader / CatalogWriter Operators (extras-gated)
      cycle.py            # Direct integration with pipekit-cycle (extras-gated)

@dataclass(frozen=True)
class StateSlice:
    """A bounded model-state request.
    
    Parallel to geocatalog.GeoSlice but for model outputs rather than
    observations. Carries the additional dimensions that model state
    requires.
    """
    # Spatial
    bbox: tuple[float, float, float, float]
    crs: str
    
    # Temporal — the key asymmetry vs GeoSlice
    valid_time: datetime              # what time the state is "about"
    run_time: datetime                # when the model run was launched
    
    # Model identity
    model_config_hash: str            # identifies the producing model + config
    model_name: str | None = None     # human-readable convenience tag
    
    # Ensemble
    ensemble_member: int | None = None        # None for deterministic runs
    ensemble_total: int | None = None         # total members in this ensemble
    
    # Lineage — the second key asymmetry
    parent_state_uri: str | None = None        # the state this run was launched from
    
    # Storage
    state_uri: str                    # where the actual state lives
    
    # Free-form metadata
    metadata: dict = field(default_factory=dict)
    
    @property
    def lead_time(self) -> timedelta:
        return self.valid_time - self.run_time
    
    @property
    def is_analysis(self) -> bool:
        """Analysis = valid_time == run_time (zero lead)."""
        return self.lead_time == timedelta(0)
    
    @property
    def is_forecast(self) -> bool:
        return self.lead_time > timedelta(0)
    
    @property
    def is_ensemble(self) -> bool:
        return self.ensemble_member is not None
    
    def to_dict(self) -> dict: ...
    @classmethod
    def from_dict(cls, d) -> "StateSlice": ...

@dataclass(frozen=True)
class StateRow:
    """Backend-neutral row view. Carries enough to construct a StateSlice."""
    bbox: tuple
    crs: str
    valid_time: datetime
    run_time: datetime
    model_config_hash: str
    state_uri: str
    ensemble_member: int | None
    parent_state_uri: str | None
    metadata: dict

@runtime_checkable
class StateCatalog(Protocol):
    """Queryable index over model states.
    
    Implementations: InMemoryStateCatalog, DuckDBStateCatalog.
    Parallel to geocatalog.GeoCatalog Protocol.
    """
    @property
    def crs(self) -> str: ...
    @property
    def n_states(self) -> int: ...
    
    def iter_slices(self, **filters) -> Iterator[StateSlice]: ...
    def query(self, **filters) -> "StateCatalog": ...   # returns a subset catalog
    def to_pandas(self) -> pd.DataFrame: ...

def at_time(cat: StateCatalog, valid_time: datetime, *, 
            model_config_hash: str | None = None,
            ensemble_member: int | None = None) -> StateCatalog:
    """All states valid at this time. Picks across forecast lead times."""

def latest_analysis(cat: StateCatalog, *, 
                    before: datetime | None = None) -> StateSlice | None:
    """The most recent analysis (lead_time=0) before a given time."""

def forecast_chain(cat: StateCatalog, run_time: datetime, *,
                   model_config_hash: str,
                   ensemble_member: int | None = None) -> list[StateSlice]:
    """All forecast states from a given run, ordered by lead time."""

def lead_window(cat: StateCatalog, run_time: datetime, 
                lead_range: tuple[timedelta, timedelta], **filters) -> StateCatalog:
    """Forecasts with lead in [lead_min, lead_max]."""

def ensemble_at(cat: StateCatalog, valid_time: datetime, run_time: datetime,
                model_config_hash: str) -> list[StateSlice]:
    """All ensemble members for a given (valid_time, run_time) pair."""

def lineage(cat: StateCatalog, state_uri: str, *, 
            depth: int = -1) -> list[StateSlice]:
    """Trace the ancestral chain back through parent_state_uri.
    
    depth=-1 means trace to the root.
    Used to answer 'what observations / earlier states contributed to this?'
    """

def assimilation_window(cat: StateCatalog, run_time: datetime,
                        obs_catalog: "GeoCatalog",
                        window: timedelta) -> "GeoCatalog":
    """Given a model run, return the observation catalog filtered to the
    obs that were assimilated. Bridges statecatalog ↔ geocatalog.
    """

class InMemoryStateCatalog:
    """pandas-backed. Suitable for small-to-medium catalogs (< 1M rows).
    Used for testing, dev, and small research workflows.
    """
    df: pd.DataFrame
    
    def __init__(self, df: pd.DataFrame): ...
    @classmethod
    def empty(cls) -> "InMemoryStateCatalog": ...
    
    def append(self, slice: StateSlice): ...
    def append_many(self, slices: list[StateSlice]): ...

class DuckDBStateCatalog:
    """SQL-backed via DuckDB + GeoParquet. Suitable for very large catalogs
    (10M+ rows). Extras-gated: pip install statecatalog[duckdb].
    
    The schema includes a state_extension namespace in the GeoParquet
    metadata, parallel to geocatalog's catalog-extension.
    """
    parquet_uri: str
    
    def __init__(self, parquet_uri: str): ...
    @classmethod
    def open(cls, parquet_uri: str) -> "DuckDBStateCatalog": ...

class CycleStateWriter(StatefulOperator):
    """Write each step of a Cycle's history to a StateCatalog.
    
    Used as a callback inside DACycle / Cycle to persist state outputs.
    """
    state_catalog: StateCatalog
    model_config_hash: str

class CycleStateReader(Operator):
    """Read initial state for a cycle from a StateCatalog.
    
    Common pattern: 'start my cycle from the latest analysis before 06Z.'
    """
    state_catalog: StateCatalog
    query: Callable[[StateCatalog], StateSlice]

import pipekit as pk
import statecatalog as sc
import pipekit_cycle as pc

cat = sc.DuckDBStateCatalog.open("s3://reanalysis/methane_v3/states.parquet")
writer = sc.integrations.cycle.CycleStateWriter(
    state_catalog=cat,
    model_config_hash=pk.config_hash(forecast_op),
)

# A 72-hour ensemble forecast — each step writes to the catalog
forecast = pc.EnsembleCycle(
    step_op=pc.Sequential([
        chemistry_forward,
        writer,                       # ← writes each step's state
    ]),
    n_steps=72,
    n_members=40,
)

trajectory = forecast(initial_state, initial_carry_state)
# state catalog now has 72 × 40 = 2880 new StateSlice rows

import statecatalog.queries as q
from datetime import datetime

# What's the most recent analysis state available before 06Z today?
slice = q.latest_analysis(
    cat,
    before=datetime(2026, 5, 17, 6),
)

# Load the actual state through whatever reader is appropriate
state = georeader.read_from_url(slice.state_uri)

reader = sc.integrations.cycle.CycleStateReader(
    state_catalog=cat,
    query=lambda c: q.latest_analysis(c, before=datetime.now()),
)

forecast = pc.Sequential([
    reader,                           # load initial state from catalog
    chemistry_forward,
    pc.Cycle(step_op=chemistry_forward, n_steps=72),
    writer,                           # write forecast states back
])

forecast(carrier=None)

# For each forecast step in a run, find the verifying analysis and compute RMSE
forecasts = q.forecast_chain(cat, run_time=datetime(2026, 5, 17, 0), 
                              model_config_hash=H)

for fcst in forecasts:
    # Find the analysis valid at the forecast's valid_time
    truth = q.latest_analysis(cat, before=fcst.valid_time)
    if truth is None: continue
    
    fcst_state = georeader.read_from_url(fcst.state_uri)
    truth_state = georeader.read_from_url(truth.state_uri)
    rmse = compute_rmse(fcst_state, truth_state)
    print(f"Lead {fcst.lead_time}: RMSE = {rmse}")

# What contributed to this analysis?
ancestors = q.lineage(cat, state_uri="s3://.../analysis_2026-05-17_06Z.nc",
                       depth=-1)

# 'ancestors' is a chain back through previous forecasts and analyses
for s in ancestors:
    print(f"{s.run_time}+{s.lead_time}: {s.state_uri}")

# Also: what obs went into the analysis step?
obs_used = q.assimilation_window(
    cat=cat,
    run_time=datetime(2026, 5, 17, 6),
    obs_catalog=obs_cat,
    window=timedelta(hours=3),
)
# 'obs_used' is a filtered geocatalog.GeoCatalog

import pipekit_train as pt

# Pull historical forecasts (a year's worth of methane forecasts) for emulator training
training_slices = list(q.lead_window(
    cat,
    run_time=datetime(2025, 1, 1),
    lead_range=(timedelta(0), timedelta(hours=72)),
    model_config_hash=H,
).iter_slices())

# Build a TrainingDataset over them
dataset = pt.SimulationDataset(
    forward_model=NullForward(),    # we're reading pre-computed trajectories
    prior=None,
    n_samples=len(training_slices),
    cycle=None,
    # custom mode: iterate state catalog rather than re-simulate
)
# (Or: add a `pt.StateCatalogDataset(state_catalog=cat, ...)` to pipekit-train)

obs_assimilated = q.assimilation_window(state_cat, run_time, obs_cat, window)

def states_that_used(state_cat: StateCatalog, obs_slice: "GeoSlice") -> list[StateSlice]:
    """Find analyses whose assimilation window contained this obs."""

[project]
name = "statecatalog"
version = "0.1.0"
dependencies = [
    "pandas>=2.0",
    "pyarrow>=17.0",
    "pyproj>=3.6",
]

[project.optional-dependencies]
duckdb  = ["duckdb>=1.0"]                       # DuckDBStateCatalog
pipekit = ["pipekit>=0.1", "pipekit-cycle>=0.1"]  # cycle integration
geocatalog = ["geocatalog>=0.1"]               # cross-catalog bridges
streaming  = ["pyarrow>=17.0"]
all     = ["statecatalog[duckdb,pipekit,geocatalog]"]