# How it works

This page describes the architecture of `aifs_modal`, focusing on how AIFS forecasts are run within computational pipelines, highlighting what is executed locally versus on Modal and how data inputs and outputs are stored at each step.

Note that this set up is not intended for running an operational AIFS forecasting platform but rather as a playground to run AIFS forecasts for scientific experiments, where the pre-processing, forecasts and post-processing can be run interactively within the same local notebook, with only the forecasts running on a *serverless* GPU environment and with a task-based approach that prevents unnecessary duplication of work when re-running cells or iterating on the analysis (see the {ref}`re-running-existing-forecasts-and-reproducibility` section below). Idempotency operates at two levels: `run_forecast` checks for existing outputs on the Modal side before ingesting ICs or starting inference; `forecast_exists` provides the same check locally, so the ephemeral app is never spun up when the forecast is already complete. See the {doc}`example notebooks <user-guide/index>` for practical example applications of this set up.

## Key concepts

- **AIFS** (Artificial Intelligence Forecasting System): ECMWF's machine-learning weather model. Two checkpoints are available: [AIFS-Single](https://huggingface.co/ecmwf/aifs-single-1.1) (deterministic) and [AIFS-ENS](https://huggingface.co/ecmwf/aifs-ens-1.0) (ensemble, stochastic perturbations via random seeds).
- **Custom checkpoints**: any [anemoi-inference](https://anemoi-inference.readthedocs.io)-compatible checkpoint can override the built-in defaults by passing a `checkpoint` dict to `run_forecast` or `run_inference` (e.g. `{"huggingface": "org/fine-tuned-model"}` or a path on a Modal Volume). `aifs-modal` reads `variable_to_input_tensor_index` directly from the checkpoint and silently drops any IC variables the model does not expect, so the IC source does not need to be trimmed manually. Requirements: (1) the checkpoint must be in anemoi-inference format (a `.ckpt` file produced by anemoi-training); (2) it must operate on the **N320 reduced Gaussian grid** — the output regridder is hardcoded to N320 → 0.25° lat/lon; (3) the IC source must cover at minimum the full set of variables the checkpoint requires.
- **Modal**: serverless GPU platform. Forecast inference runs on Modal containers with GPU access. Ingestion and orchestration also run on Modal (CPU), with workers co-located with IC sources for in-region bandwidth.
- **Icechunk**: versioned array storage engine backed by S3-compatible object storage. Used for forecast outputs, providing git-like branching and commit history.
- **Initial conditions**: analysis fields ingested from an IC source (Brightband, ARCO-ERA5, or ECMWF open data), regridded to the N320 reduced Gaussian grid and written to a shared Modal IC Volume. AIFS requires two consecutive 6-hourly analyses (t−6h and t) as input. The Volume is a temporary per-run cache; ICs are deleted after successful inference unless `keep_ics=True`.
- **IC sources**: four sources are supported — `ifs-arraylake` (default, Brightband ECMWF IFS on Earthmover ArrayLake, co-located in `us-east`), `era5-arco` (ARCO-ERA5 on GCS, co-located in `us-central1`), `ifs-ekd` (ECMWF open-data S3), and `era5-cds` (Copernicus CDS).

The diagrams below use a consistent four-zone color scheme:

| Zone           | Color  | Description                                    |
| -------------- | ------ | ---------------------------------------------- |
| Local          | grey   | Notebook or script on your machine             |
| Modal CPU      | blue   | Orchestration and IC ingestion on Modal        |
| Modal GPU      | orange | AIFS inference on Modal                        |
| Object storage | green  | IC source and Icechunk output repository on S3 |

## Deterministic forecast (`run_forecast`)

A single AIFS-Single run on one GPU. `run_forecast` is a Modal CPU function that acts as the orchestrator: it checks whether ICs are already on the IC Volume, dispatches ingestion to a co-located Modal CPU worker if they are missing, and then dispatches inference to a GPU container.

```{mermaid}
flowchart LR
    subgraph local ["Local"]
        A["configure\nrun_forecast.remote()"]
        J[read outputs]
    end
    subgraph modal_cpu ["Modal — CPU"]
        B["orchestrate\n+ ingest ICs → Volume"]
    end
    subgraph modal_gpu ["Modal — GPU"]
        G["AIFS-Single\ninference"]
    end
    subgraph storage ["Object storage"]
        IC[("IC source")]
        OUT[("Icechunk outputs")]
    end
    A --> B
    IC --> B
    B --> G
    G --> OUT
    OUT --> J

    style local fill:#555,stroke:#333,color:#fff
    style A fill:#777,stroke:#555,color:#fff
    style J fill:#777,stroke:#555,color:#fff
    style modal_cpu fill:#2c6fad,stroke:#1a4a7a,color:#fff
    style B fill:#4a90d9,stroke:#2c6fad,color:#fff
    style modal_gpu fill:#c44d1a,stroke:#8c3210,color:#fff
    style G fill:#ff6b35,stroke:#c44d1a,color:#fff
    style storage fill:#1a5c3a,stroke:#0d3d26,color:#fff
    style IC fill:#2e7d52,stroke:#1a5c3a,color:#fff
    style OUT fill:#2e7d52,stroke:#1a5c3a,color:#fff
```

**Flow:**

1. **Orchestrate** (Modal CPU, `run_forecast`): check whether the forecast output already exists; if so, return immediately without ingesting ICs or starting inference. Otherwise, check the IC Volume for the target date. If ICs are missing, dispatch `ingest_ifs_arraylake` (Modal `us-east`) or `ingest_era5_arco` (Modal `us-central1`) and wait for completion; inline sources (`ifs-ekd`, `era5-cds`) run directly inside the orchestrator.
2. **Ingest** (Modal CPU, co-located): download and regrid the IC fields to N320; write them to the IC Volume.
3. **Inference** (Modal GPU, `run_inference`): load initial conditions from the IC Volume, load the checkpoint into `SimpleRunner`, drop any IC variables not required by the checkpoint, run the AIFS model, regrid output fields from N320 to 0.25° lat/lon, and stream each 6-hourly step to the outputs Icechunk repository.
4. **Post-process** (local, CPU): open the outputs repository with xarray and analyze them.

## Ensemble forecast (`run_forecast` with `n_members`)

Runs multiple AIFS-ENS members one after another on a single GPU. Each member uses a different random seed for stochastic perturbations. Ingestion and orchestration are identical to the deterministic case; the difference is that `run_inference` loops over members on the same GPU.

```{mermaid}
flowchart LR
    subgraph local ["Local"]
        A["configure\nrun_forecast.remote(n_members=k)"]
        J[read outputs]
    end
    subgraph modal_cpu ["Modal — CPU"]
        B["orchestrate\n+ ingest ICs → Volume"]
    end
    subgraph modal_gpu ["Modal — GPU"]
        G["AIFS-ENS\nmember 0 → 1 → … → k−1"]
    end
    subgraph storage ["Object storage"]
        IC[("IC source")]
        OUT[("Icechunk outputs")]
    end
    A --> B
    IC --> B
    B --> G
    G --> OUT
    OUT --> J

    style local fill:#555,stroke:#333,color:#fff
    style A fill:#777,stroke:#555,color:#fff
    style J fill:#777,stroke:#555,color:#fff
    style modal_cpu fill:#2c6fad,stroke:#1a4a7a,color:#fff
    style B fill:#4a90d9,stroke:#2c6fad,color:#fff
    style modal_gpu fill:#c44d1a,stroke:#8c3210,color:#fff
    style G fill:#ff6b35,stroke:#c44d1a,color:#fff
    style storage fill:#1a5c3a,stroke:#0d3d26,color:#fff
    style IC fill:#2e7d52,stroke:#1a5c3a,color:#fff
    style OUT fill:#2e7d52,stroke:#1a5c3a,color:#fff
```

Each member's output is appended along the `ensemble_member` dimension. This mode is simpler and cheaper (one GPU) but slower for large ensembles.

(re-running-existing-forecasts-and-reproducibility)=

## Re-running existing forecasts and reproducibility

The set up of `aifs-modal` is designed to be used in computational pipelines. Accordingly a key feature is that every step checks whether its output already exists before doing any work, so interrupted or repeated runs pick up where they left off without duplicating effort or GPU cost.

Idempotency operates at two levels. The **local check** (`forecast_exists`) runs in the notebook before `app.run()` — if it returns `True`, the ephemeral app is never created. The **Modal-side check** (`run_forecast`) runs inside the orchestrator — if the output already exists, ingestion and inference are skipped without consuming a GPU.

| Step                           | Where | Skip condition                                      | Override         |
| ------------------------------ | ----- | --------------------------------------------------- | ---------------- |
| `forecast_exists`              | local | Output group exists and has enough members          | —                |
| `run_forecast` output check    | Modal | Same as above — prevents IC ingestion and inference | `overwrite=True` |
| IC ingestion in `run_forecast` | Modal | Both IC dates already present on the IC Volume      | —                |

**Ensemble reproducibility.** Each ensemble member fixes the PyTorch random seed to its member index (`torch.manual_seed(member_id)`) before running AIFS-ENS. This means that member *k* always produces the same trajectory regardless of how many members are in the ensemble or how they are distributed across containers, making it straightforward to extend an existing ensemble run by increasing `n_members` (existing members are skipped, only the new ones are computed).