AIFS on Modal

aifs-modal is a Python library for running serverless AIFS weather forecasts on Modal. Only AIFS inference runs on a Modal GPU; initial-conditions ingestion is dispatched to co-located Modal CPU workers and cached on a Modal Volume. Forecast outputs are stored in Icechunk repositories on S3-compatible object storage (or Earthmover ArrayLake), giving them a git-like version history you can open with xarray from anywhere.

Features

• Two forecast modes: deterministic (AIFS-Single), ensemble (AIFS-ENS).

• Versioned array storage via Icechunk: forecast outputs are stored with git-like branching and commit history, so you can reproduce any past run and safely extend ensemble experiments. All storage backends supported by Icechunk are available (Tigris, AWS S3, Cloudflare R2, Google Cloud Storage, Azure Blob Storage, Earthmover ArrayLake).

• Task-based pipeline: ingestion and inference are framed as pipeline steps, skipped when outputs already exist. Re-running a notebook cell never duplicates work or wastes GPU resources.

• Reproducible ensembles: each member is seeded by its index (torch.manual_seed(member_id)), so you can safely extend an ensemble run later by increasing n_members.

  • Flexible initial conditions: ingest from operational IFS HRES via the Earthmover Mar (default, alternatively ECMWF open data without any login needed) or ERA5 (via Google ARCO-ERA5 or Copernicus CDS), or any custom field array following the AIFS variable convention.

Why aifs-modal?

Two publicly accessible AIFS archives already exist: the ECMWF operational archive and dynamical.org. The key aim of aifs-modal is to fill the gaps for cases not covered by the archives:

1. Extended lead times: operational forecasts top out at 10–15 days, whereas aifs-modal runs to any lead time. See the jet-stream free-run notebook for a 105-day forecast to explore NH jet-stream comparison with ERA5 climatology.

2. Retrospective reforecasts from ERA5: ERA5 spans 80+ years of reanalysis, enabling ensemble reforecasts for any past date regardless of operational archive retention. See the heatwave reforecast notebook for a 10-member AIFS-ENS reforecast of the June 2025 European heatwave verified against station observations.

3. Custom initial conditions: ingestion and inference are decoupled, so any field array following the AIFS variable convention can serve as initial conditions. Applications include perturbed or bias-corrected states from climate-model scenarios. See the CMIP6 SST-patch notebook for an example patching ERA5 sea-surface temperatures to mimic a warmer-climate storyline before running the forecast.

4. Custom checkpoints: the built-in AIFS-Single and AIFS-ENS checkpoints can be overridden via the checkpoint parameter of run_forecast, accepting any anemoi-inference-compatible checkpoint — aifs-modal reads the required variable set from the checkpoint and automatically drops any surplus IC variables. Applications include models fine-tuned on regional data or trained with domain-specific variables.

AIFS-Single deterministic 96-hour forecast of 2-meter temperature over Europe.

AIFS-Single 105-day free-run forecast of the NH jet stream compared against ERA5 climatology.

AIFS-ENS 10-member ensemble forecast of 2-meter temperature for Lausanne, verified against MeteoSwiss station data.

AIFS-Single forecast difference (warmer SST − baseline) showing the downstream impact of a CMIP6 SST patch on 2-meter temperature.

Usage

Before running, complete the setup steps (Modal account, storage bucket, IC-source credentials, and pip install aifs-modal). Then a forecast runs in a few lines:

import datetime
import xarray as xr
from aifs_modal import app, run_forecast

date = datetime.datetime(2025, 6, 20, 0, tzinfo=datetime.UTC)

# 1. Run a 96-hour forecast on Modal — ICs are ingested automatically (~$0.05 GPU cost)
with app.run():
    run_forecast.remote(
        date,
        storage_bucket="my-tigris-bucket",
        source_repo="org/ecmwf-ifs-hres-ics-open",  # Brightband IFS dataset
        outputs_prefix="aifs-outputs",
        lead_time=96,
    )

# 2. Open the results locally with xarray — no GPU needed
import icechunk
repo = icechunk.Repository.open(
    icechunk.tigris_storage(bucket="my-tigris-bucket", prefix="aifs-outputs", ...)
)
ds = xr.open_dataset(
    repo.readonly_session("main").store,
    group="2025-06-20/00z",
    engine="zarr",
    zarr_format=3,
)

See the user guide for a full worked example including ensemble forecasts and output visualization.

Requirements

A Modal account, an S3-compatible storage bucket, and credentials for your chosen initial-conditions source. See the setup guide for step-by-step instructions.

Installation

pip install aifs-modal

Development

Install pixi and use the generated environments and tasks:

pixi install -e dev
pixi run -e test test
pixi run -e doc build-doc

Commit the generated pixi.lock file once you have resolved the environments for your project, and update it whenever dependencies change.

Acknowledgements

• This package was created with the martibosch/cookiecutter-pypixi-package project template.