Skip to content

ekalmX Design Doc

Differentiable ensemble Kalman methods for filtering, inversion, and joint estimation in JAX/Equinox.

Structure

filterX/
├── README.md              # This file
├── vision.md              # Motivation, user stories, design principles, identity
├── architecture.md        # Three-layer stack, protocols, data types, package layout, dependencies
├── boundaries.md          # Ownership, ecosystem, scope, testing strategy, roadmap
├── decisions.md           # 10 ADRs (D1–D10)
├── api/
│   ├── README.md          # Surface inventory, notation, performance notes
│   ├── primitives.md      # Layer 0 — ensemble statistics, gain, localization, inflation, patches
│   ├── components.md      # Layer 1 — protocols, filters, processes, smoothers, schedulers
│   └── models.md          # Layer 2 — LETKF/ETKF/EnSRF, EKI/EKS/UKI, optax transforms
├── features/
│   ├── filters.md         # Sequential EnKF variants (ETKF, LETKF, EnSRF, ESTKF, etc.)
│   ├── processes.md       # EKP algorithms (EKI, EKS, UKI, GNKI, SparseEKI, TEKI)
│   ├── smoothers.md       # Ensemble smoothers (EnKS, RTS, fixed-lag, IES)
│   ├── localization_inflation.md  # Localization tapers + inflation strategies (incl. GeoLocalizer)
│   ├── diagnostics.md     # DA evaluation metrics (rank histogram, Desroziers, CRPS, etc.)
│   ├── differentiable_da.md      # Deep dive: backprop through filters
│   ├── optax_ekp.md       # Deep dive: EKP as optax GradientTransformations
│   ├── multi_instrument.md       # JointObsOperator + SequentialAssimilation (D12)
│   └── state_persistence.md      # Serialization contract + warm-start patterns (D15)
├── integrations/
│   ├── README.md          # Index and reading order for ecosystem integration docs
│   ├── pipekit.md         # Structural Protocol satisfaction with pipekit-cycle (D11)
│   ├── plumax.md          # Multi-instrument methane attribution recipe (Tier IV)
│   └── geostack.md        # CarrierAdapter for coordax / GeoTensor particles (D13)
├── examples/
│   ├── README.md          # Index and reading order
│   ├── primitives.md      # Layer 0 usage patterns
│   ├── components.md      # Layer 1 composition patterns
│   ├── models.md          # Layer 2 high-level API patterns
│   └── integration.md     # Layer 3 cross-library patterns (optax, somax, gaussx, geo_toolz)
├── research/
│   ├── README.md          # Index for research docs
│   ├── research_enskf.md  # Audit of existing kalman_filter/ code (~5,000 lines)
│   ├── ens_kalman_process.md  # EKP theory: math, process zoo, BLR connection
│   └── operational_da.md  # Streaming / alert-service patterns
└── decisions.md           # 16 ADRs (D1–D16)

Reading Order

  1. vision.md — understand the why
  2. architecture.md — understand the three-layer stack
  3. boundaries.md — understand the scope
  4. api/README.md — scan the surface
  5. api/primitives.md → components.md → models.md — drill into detail
  6. features/filters.md → processes.md → smoothers.md — algorithm deep dives
  7. features/localization_inflation.md → diagnostics.md — tuning and evaluation
  8. features/differentiable_da.md → optax_ekp.md — unique capabilities
  9. features/multi_instrument.md → features/state_persistence.md — operational extensions
  10. examples/primitives.md → components.md → models.md → integration.md — see it in action
  11. integrations/pipekit.md → integrations/plumax.md → integrations/geostack.md — ecosystem integration recipes
  12. decisions.md — understand the tradeoffs (D1–D16)
  13. research/ — background theory, code audits, operational DA