fluctifex

Differentiable Earth-system modeling in JAX — assimilate raw observations, forecast, and control.

fluctifex — Latin fluctus (wave, flow) + -fex (maker), by analogy with opifex, artifex. The name points at the physics, not the weather: continuum flow, which is as true of the ocean and the climate as of the atmosphere.

Status: alpha (v0.1.0). APIs are unstable and will change. This is research infrastructure.

The thesis

Most machine-learning weather models are very good correlational forecasters. They train on reanalysis — itself the output of a physical assimilation system — and map one gridded state to the next. That is useful, and it is rung one.

fluctifex is built around a different claim: if the entire pipeline is differentiable, from raw sensor readings through to the predicted field, then estimation and control fall out of the same machinery.

• Differentiating the model backward through time against past observations is data assimilation (4D-Var) — the adjoint comes free from autodiff instead of being hand-written.
• Differentiating it forward against a future objective is control (differentiable MPC) — the intervention gradient falls out of the same rollout.

Estimation and control are mirror images: same forward model, opposite directions in time.

   ┌─▶ OBSERVE ............ raw multi-sensor data (satellite, radar, stations, buoys)
   │       ▼
   │   ASSIMILATE .......... optimize the STATE to match PAST obs      [backward · estimation]
   │       ▼
   │   FORECAST ............ roll the differentiable model forward
   │       ▼
   │   DECIDE .............. optimize future ACTIONS to hit an OBJECTIVE [forward · control]
   │       ▼
   │   ACT ................. apply the first action only (receding horizon)
   │       ▼
   └──────── world advances → repeat each cycle ────────

Every arrow is differentiable. The same jax.grad supplies the adjoint going backward and the action gradient going forward.

Why differentiability has to start at the raw data

If the pipeline only becomes differentiable after preprocessing, then assimilation can only ever reach as far back as the preprocessed state — which means you inherit whatever a conventional assimilation system already decided. Beginning at raw, heterogeneous, gappy sensor readings is what makes end-to-end assimilation possible at all. That is why the data layer here is built on datarax, whose pipeline stages are themselves Flax NNX modules.

Architecture

Ecosystem

fluctifex is the Earth-system domain layer of the Avitai stack. It composes rather than reimplements:

• opifex — neural operators (SFNO, FNO, …), PINNs, UQ, assimilation primitives, differentiable MPC
• datarax — differentiable, DAG-structured, shardable data pipelines
• calibrax — metrics, calibration, FLOP/roofline profiling
• artifex — generative models (diffusion ensembles)

Its sibling cellifex applies the same substrate to cellular dynamics.

Install

pip install fluctifex                 # core (JAX + Flax only)
pip install "fluctifex[avitai]"       # + opifex / datarax / calibrax
pip install "fluctifex[all]"          # everything, incl. diffusion and docs

From source:

git clone https://github.com/avitai/fluctifex.git
cd fluctifex && ./setup.sh && source ./activate.sh
pytest tests/ -v

Quickstart

import jax
from flax import nnx

from fluctifex.models import Forecaster, ForecasterConfig
from fluctifex.training import rollout_loss

model = Forecaster(ForecasterConfig(in_channels=8, out_channels=8), rngs=nnx.Rngs(0))

state = jax.random.normal(jax.random.key(0), (2, 8, 32, 64))   # (B, C, H, W)
next_state = model(state)

# Train through the rollout, not on single steps — a-posteriori stability.
loss = rollout_loss(model, state, targets, steps=4)

Assimilate sparse observations into a state estimate:

from fluctifex.assimilation import four_dvar

analysis = four_dvar(
    model, background=x_b, observations=y, observation_operator=H, num_steps=50
)

Then plan an intervention through the forecast:

from fluctifex.control import DifferentiableMPC, MPCConfig

mpc = DifferentiableMPC(MPCConfig(horizon=8), dynamics=model)
action = mpc.solve(analysis, objective)      # apply only the first step, then re-plan

Design commitments

1. Differentiable end to end, beginning at raw observations.
2. Composition over reimplementation — the Avitai siblings do the heavy lifting.
3. Stability is a first-class result, not a footnote: train through rollouts, report drift.
4. Uncertainty gates action. A forecast does not drive a recommendation unless its uncertainty is calibrated.
5. Honest baselines. Persistence, climatology, and matched-cost references, always.

Honest limitations

• Alpha software; APIs will break between minor versions.
• No pretrained weights are shipped yet.
• Single-host SPMD is exercised; multi-host is untested.
• Physical intervention (control) is a research capability, validated only in simulation. Nothing here is fit for real-world atmospheric intervention, and it should not be represented as such.

License

MIT — see LICENSE.