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.

flowchart TD
    OBSERVE["OBSERVE<br/>raw multi-sensor data<br/>(satellite · radar · stations · buoys)"]
    ASSIMILATE["ASSIMILATE — backward · estimation<br/>optimize the STATE to match PAST obs"]
    FORECAST["FORECAST<br/>roll the differentiable model forward"]
    DECIDE["DECIDE — forward · control<br/>optimize future ACTIONS toward an OBJECTIVE"]
    ACT["ACT<br/>apply the first action only (receding horizon)"]

    OBSERVE --> ASSIMILATE --> FORECAST --> DECIDE --> ACT
    ACT -->|world advances · repeat each cycle| OBSERVE

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 is designed to build on datarax, whose pipeline stages are themselves Flax NNX modules: the windowed series and the padded station observations are datarax sources, and normalization is a datarax operator stage, so a training epoch is a datarax Pipeline.

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
• substrax — device meshes, batch placement, checkpoints

Its sibling cellifex applies the same substrate to cellular dynamics.

Install

The Avitai siblings (opifex, datarax, calibrax, substrax) are core dependencies, installed from PyPI.

From source (recommended — auto-detects CUDA / Metal / CPU):

git clone https://github.com/avitai/fluctifex.git
cd fluctifex
./setup.sh && source ./activate.sh
uv run pytest

Force a backend or add optional extras as needed:

./setup.sh --backend cpu    # or cuda12, metal
uv sync --extra data        # xarray / zarr / netCDF4 for Earth-system formats

Quickstart

import jax
from flax import nnx

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

model = build_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.