Metadata-Version: 2.5
Name: fluctifex
Version: 0.1.1
Summary: Differentiable Earth-system modeling in JAX: assimilate raw observations, forecast, and control.
Project-URL: Homepage, https://github.com/avitai/fluctifex
Project-URL: Issues, https://github.com/avitai/fluctifex/issues
Project-URL: Source, https://github.com/avitai/fluctifex
Author: Mahdi Shafiei
License: MIT License
        
        Copyright (c) 2026 Mahdi Shafiei
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
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: Operating System :: MacOS
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
Classifier: Typing :: Typed
Requires-Python: <3.14,>=3.12
Requires-Dist: calibrax>=0.1.5
Requires-Dist: datarax>=0.1.6
Requires-Dist: flax>=0.12.9
Requires-Dist: jax>=0.11.1
Requires-Dist: jaxlib>=0.11.1
Requires-Dist: numpy>=1.26
Requires-Dist: opifex>=0.2.2
Requires-Dist: optax>=0.2
Requires-Dist: substrax>=0.1.4
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: gpu
Requires-Dist: jax[cuda12]>=0.11.1; (sys_platform == 'linux') and extra == 'gpu'
Requires-Dist: jaxlib>=0.11.1; (sys_platform == 'linux') and extra == 'gpu'
Provides-Extra: metal
Requires-Dist: jax-metal>=0.1.0; (sys_platform == 'darwin') and extra == 'metal'
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.

```mermaid
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`](https://github.com/avitai/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

| Module | Responsibility |
|---|---|
| `fluctifex.data` | Leakage-safe temporal splits; datarax sources for windowed grids and masked sparse station observations |
| `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
- **[substrax](https://github.com/avitai/substrax)** — device meshes, batch placement, checkpoints

Its sibling [cellifex](https://github.com/avitai) 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):

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

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

## Quickstart

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

```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).
