Metadata-Version: 2.4
Name: fluctifex
Version: 0.1.0
Summary: Differentiable Earth-system modeling in JAX: assimilate raw observations, forecast, and control.
Project-URL: Homepage, https://github.com/avitai/fluctifex
Project-URL: Repository, https://github.com/avitai/fluctifex
Project-URL: Issues, https://github.com/avitai/fluctifex/issues
Author-email: Mahdi Shafiei <m.mahdi.shafiei@gmail.com>
License: MIT
License-File: LICENSE
Keywords: climate,data-assimilation,differentiable-physics,jax,model-predictive-control,neural-operators,scientific-machine-learning,weather
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: beartype>=0.19
Requires-Dist: einops>=0.8
Requires-Dist: flax>=0.12
Requires-Dist: jax>=0.8
Requires-Dist: jaxtyping>=0.2
Requires-Dist: numpy>=1.26
Requires-Dist: optax>=0.2
Provides-Extra: all
Requires-Dist: avitai-artifex>=0.1.2; extra == 'all'
Requires-Dist: calibrax>=0.1.1; extra == 'all'
Requires-Dist: datarax>=0.1.4; extra == 'all'
Requires-Dist: mkdocs-material>=9.5; extra == 'all'
Requires-Dist: mkdocs>=1.6; extra == 'all'
Requires-Dist: mkdocstrings[python]>=0.26; extra == 'all'
Requires-Dist: netcdf4>=1.6; extra == 'all'
Requires-Dist: opifex>=0.2.0; extra == 'all'
Requires-Dist: pre-commit>=3.8; extra == 'all'
Requires-Dist: pyright>=1.1.380; extra == 'all'
Requires-Dist: pytest-cov>=5.0; extra == 'all'
Requires-Dist: pytest>=8.0; extra == 'all'
Requires-Dist: ruff>=0.6; extra == 'all'
Requires-Dist: xarray>=2024.1; extra == 'all'
Requires-Dist: zarr>=2.17; extra == 'all'
Provides-Extra: avitai
Requires-Dist: calibrax>=0.1.1; extra == 'avitai'
Requires-Dist: datarax>=0.1.4; extra == 'avitai'
Requires-Dist: opifex>=0.2.0; extra == 'avitai'
Provides-Extra: data
Requires-Dist: netcdf4>=1.6; extra == 'data'
Requires-Dist: xarray>=2024.1; extra == 'data'
Requires-Dist: zarr>=2.17; extra == 'data'
Provides-Extra: dev
Requires-Dist: pre-commit>=3.8; extra == 'dev'
Requires-Dist: pyright>=1.1.380; extra == 'dev'
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.5; extra == 'docs'
Requires-Dist: mkdocs>=1.6; extra == 'docs'
Requires-Dist: mkdocstrings[python]>=0.26; extra == 'docs'
Provides-Extra: generative
Requires-Dist: avitai-artifex>=0.1.2; extra == 'generative'
Provides-Extra: gpu
Requires-Dist: jax[cuda12]>=0.8; extra == 'gpu'
Description-Content-Type: text/markdown

# 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`](https://github.com/avitai/datarax), whose pipeline stages are themselves Flax NNX
modules.

## Architecture

| Module | Responsibility |
|---|---|
| `fluctifex.data` | Dense grids + sparse station observations, as a differentiable DAG pipeline |
| `fluctifex.models` | Forecast backbones (spherical/planar neural operators) and multi-sensor fusion |
| `fluctifex.assimilation` | Variational (4D-Var) and sequential (Kalman) state estimation |
| `fluctifex.control` | Differentiable model-predictive control with safety constraints |
| `fluctifex.training` | Rollout-through-training (a-posteriori) for long-horizon stability |
| `fluctifex.evaluation` | Latitude-weighted and probabilistic verification (WeatherBench-2 style) |

## Ecosystem

`fluctifex` is the Earth-system domain layer of the [Avitai](https://github.com/avitai) stack. It
composes rather than reimplements:

- **[opifex](https://github.com/avitai/opifex)** — neural operators (SFNO, FNO, …), PINNs, UQ,
  assimilation primitives, differentiable MPC
- **[datarax](https://github.com/avitai/datarax)** — differentiable, DAG-structured, shardable data pipelines
- **[calibrax](https://github.com/avitai/calibrax)** — metrics, calibration, FLOP/roofline profiling
- **[artifex](https://github.com/avitai/artifex)** — generative models (diffusion ensembles)

Its sibling [cellifex](https://github.com/avitai) applies the same substrate to cellular dynamics.

## Install

```bash
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:

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

## Quickstart

```python
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:

```python
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:

```python
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](LICENSE).
