{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# AIFS Extended Free-Running Forecast: Mid-February to End of May 2026\n", "\n", "This notebook:\n", "\n", "1. Runs [AIFS](https://huggingface.co/ecmwf/aifs-single-1.1) in **free-running mode** from **2026-02-15 00 UTC** to **2026-05-31 00 UTC** (105 days, 2520 h) on Modal.\n", "2. Evaluates the forecast with the **Anomaly Correlation Coefficient (ACC) for Z500** at **+14 days** (2026-03-01 00 UTC) against the ERA5 analysis.\n", "3. Diagnoses the **Northern-Hemisphere jet-stream latitude** from the AIFS 300 hPa zonal wind over the whole free-run period and compares it with the ERA5 **WMO 1991–2020 climatological normals** for the same calendar period.\n", "\n", "Each experiment is isolated in its own icechunk **branch** of the outputs repository, named after the `experiment` variable set in the configuration cell below. Re-running with a different branch name (e.g. a different checkpoint or start date) keeps results completely independent." ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "import datetime as dt\n", "import os\n", "import tempfile\n", "import time\n", "\n", "import cartopy.crs as ccrs\n", "import cartopy.feature as cfeature\n", "import cdsapi\n", "import icechunk\n", "import matplotlib.dates as mdates\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import pandas as pd\n", "import xarray as xr\n", "\n", "from aifs_modal import app, ingest, run_forecast" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "## Configuration" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Experiment : jet-stream-feb26\n", "Start : 2026-02-15 00:00:00+00:00\n", "End : 2026-05-31 00:00:00+00:00\n", "Lead : 2520 h (105 days)\n", "Verification date (T+14d): 2026-03-01 00:00:00+00:00\n" ] } ], "source": [ "# --- experiment label (drives the icechunk branch for outputs) ---\n", "experiment = \"jet-stream-feb26\"\n", "\n", "# --- forecast window ---\n", "start_date = dt.datetime(2026, 2, 15, 0, tzinfo=dt.UTC)\n", "end_date = dt.datetime(2026, 5, 31, 0, tzinfo=dt.UTC)\n", "lead_time = int((end_date - start_date).total_seconds() // 3600) # hours\n", "print(f\"Experiment : {experiment}\")\n", "print(f\"Start : {start_date}\")\n", "print(f\"End : {end_date}\")\n", "print(f\"Lead : {lead_time} h ({lead_time / 24:.0f} days)\")\n", "\n", "# --- skill-score reference ---\n", "SKILL_LEAD_DAYS = 14\n", "verification_date = start_date + dt.timedelta(days=SKILL_LEAD_DAYS)\n", "print(f\"Verification date (T+{SKILL_LEAD_DAYS}d): {verification_date}\")\n", "\n", "# --- storage ---\n", "storage_bucket = \"aifs-modal\"\n", "initial_conditions_prefix = \"aifs-initial-conditions\"\n", "initial_conditions_branch = \"main\"\n", "outputs_prefix = \"aifs-outputs\"\n", "outputs_branch = experiment # one branch per experiment" ] }, { "cell_type": "markdown", "id": "4", "metadata": {}, "source": [ "## 1. Ingesting initial conditions\n", "\n", "AIFS needs two consecutive analysis states (t−6 h and t) as initial conditions." ] }, { "cell_type": "code", "execution_count": null, "id": "5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " \u001b[2m2026-03-06T11:11:08.094571Z\u001b[0m \u001b[33m WARN\u001b[0m \u001b[1;33maws_runtime::env_config::normalize\u001b[0m\u001b[33m: \u001b[33mprofile [plugins] ignored; sections in the AWS config file (other than [default]) must have a prefix i.e. [profile my-profile]\u001b[0m\n", " \u001b[2;3mat\u001b[0m /home/conda/feedstock_root/build_artifacts/icechunk_1771452354651/_build_env/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/aws-runtime-1.5.18/src/env_config/normalize.rs:121\n", "\n", "Initial conditions missing for 2026-02-15T00:00:00+00:00 (group: 2026-02-15/00z); ingesting now\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " " ] }, { "name": "stdout", "output_type": "stream", "text": [ "By downloading data from the ECMWF open data dataset, you agree to the terms: Attribution 4.0 International (CC BY 4.0). Please attribute ECMWF when downloading this data.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " " ] }, { "name": "stdout", "output_type": "stream", "text": [ "Ingested initial conditions for 2026-02-15T00:00:00+00:00\n", "Initial conditions missing for 2026-02-14T18:00:00+00:00 (group: 2026-02-14/18z); ingesting now\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " " ] }, { "name": "stdout", "output_type": "stream", "text": [ "Ingested initial conditions for 2026-02-14T18:00:00+00:00\n" ] } ], "source": [ "initial_conditions_storage = icechunk.tigris_storage(\n", " bucket=storage_bucket,\n", " prefix=initial_conditions_prefix,\n", " region=os.getenv(\"AWS_REGION\", None),\n", " access_key_id=os.environ[\"AWS_ACCESS_KEY_ID\"],\n", " secret_access_key=os.environ[\"AWS_SECRET_ACCESS_KEY\"],\n", ")\n", "initial_conditions_repo = icechunk.Repository.open(initial_conditions_storage)\n", "\n", "for date in [start_date, start_date - dt.timedelta(hours=6)]:\n", " ingest.ensure_date_ingested(date, initial_conditions_repo)" ] }, { "cell_type": "markdown", "id": "6", "metadata": {}, "source": [ "## 2. Running the extended free-running forecast on Modal\n", "\n", "The forecast is written to the branch named after `experiment` (`outputs_branch`). Running this notebook again with a different `experiment` value produces a fully independent branch, leaving previous results untouched." ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [], "source": [ "# with modal.enable_output(): # uncomment for live logs\n", "with app.run():\n", " _start = time.perf_counter()\n", " run_forecast.remote(\n", " start_date,\n", " storage_bucket,\n", " lead_time=lead_time,\n", " outputs_prefix=outputs_prefix,\n", " outputs_branch=outputs_branch,\n", " include_pressure_levels=True, # needed for Z500 and U300\n", " )\n", " print(\n", " f\"Run {lead_time}-hour free-run forecast in {time.perf_counter() - _start:.1f}s\"\n", " )" ] }, { "cell_type": "markdown", "id": "8", "metadata": {}, "source": [ "## 3. Loading forecast outputs" ] }, { "cell_type": "code", "execution_count": null, "id": "9", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " \u001b[2m2026-03-06T13:47:48.204011Z\u001b[0m \u001b[33m WARN\u001b[0m \u001b[1;33maws_runtime::env_config::normalize\u001b[0m\u001b[33m: \u001b[33mprofile [plugins] ignored; sections in the AWS config file (other than [default]) must have a prefix i.e. [profile my-profile]\u001b[0m\n", " \u001b[2;3mat\u001b[0m /home/conda/feedstock_root/build_artifacts/icechunk_1771452354651/_build_env/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/aws-runtime-1.5.18/src/env_config/normalize.rs:121\n", "\n" ] }, { "data": { "text/html": [ "
<xarray.Dataset> Size: 178GB\n",
"Dimensions: (valid_time: 420, lat: 721, lon: 1440, pressure: 13)\n",
"Coordinates:\n",
" * valid_time (valid_time) datetime64[ns] 3kB 2026-02-15T06:00:00 ... 2026-...\n",
" * lat (lat) float64 6kB 90.0 89.75 89.5 89.25 ... -89.5 -89.75 -90.0\n",
" * lon (lon) float64 12kB 0.0 0.25 0.5 0.75 ... 359.0 359.2 359.5 359.8\n",
" * pressure (pressure) int64 104B 50 100 150 200 250 ... 700 850 925 1000\n",
"Data variables: (12/30)\n",
" 100u (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" 100v (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" 10v (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" 2t (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" hcc (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" lcc (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" ... ...\n",
" ro (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" u (valid_time, pressure, lat, lon) float32 23GB dask.array<chunksize=(2, 13, 721, 1440), meta=np.ndarray>\n",
" v (valid_time, pressure, lat, lon) float32 23GB dask.array<chunksize=(2, 13, 721, 1440), meta=np.ndarray>\n",
" z (valid_time, pressure, lat, lon) float32 23GB dask.array<chunksize=(2, 13, 721, 1440), meta=np.ndarray>\n",
" tcc (valid_time, lat, lon) float32 2GB dask.array<chunksize=(32, 721, 1440), meta=np.ndarray>\n",
" w (valid_time, pressure, lat, lon) float32 23GB dask.array<chunksize=(2, 13, 721, 1440), meta=np.ndarray><xarray.DataArray 'z' (latitude: 721, longitude: 1440)> Size: 4MB\n",
"[1038240 values with dtype=float32]\n",
"Coordinates:\n",
" * latitude (latitude) float64 6kB 90.0 89.75 89.5 ... -89.75 -90.0\n",
" * longitude (longitude) float64 12kB 0.0 0.25 0.5 ... 359.2 359.5 359.8\n",
" number int64 8B ...\n",
" valid_time datetime64[ns] 8B 2026-03-01\n",
" pressure_level float64 8B 500.0\n",
" expver <U4 16B ...\n",
"Attributes: (12/31)\n",
" GRIB_paramId: 129\n",
" GRIB_dataType: an\n",
" GRIB_numberOfPoints: 1038240\n",
" GRIB_typeOfLevel: isobaricInhPa\n",
" GRIB_stepUnits: 1\n",
" GRIB_stepType: instant\n",
" ... ...\n",
" GRIB_shortName: z\n",
" GRIB_totalNumber: 0\n",
" GRIB_units: m**2 s**-2\n",
" long_name: Geopotential\n",
" units: m**2 s**-2\n",
" standard_name: geopotential<xarray.DataArray 'z' (latitude: 721, longitude: 1440)> Size: 4MB\n",
"array([[49978.207, 49978.207, 49978.207, ..., 49978.207, 49978.207,\n",
" 49978.207],\n",
" [49971.332, 49971.332, 49971.332, ..., 49971.293, 49971.3 ,\n",
" 49971.31 ],\n",
" [49964.727, 49964.773, 49964.793, ..., 49964.64 , 49964.7 ,\n",
" 49964.707],\n",
" ...,\n",
" [48822.566, 48822.582, 48822.59 , ..., 48822.6 , 48822.55 ,\n",
" 48822.6 ],\n",
" [48815.883, 48815.86 , 48815.867, ..., 48815.875, 48815.875,\n",
" 48815.867],\n",
" [48808.316, 48808.316, 48808.316, ..., 48808.316, 48808.316,\n",
" 48808.316]], shape=(721, 1440), dtype=float32)\n",
"Coordinates:\n",
" * latitude (latitude) float64 6kB 90.0 89.75 89.5 ... -89.75 -90.0\n",
" * longitude (longitude) float64 12kB 0.0 0.25 0.5 ... 359.2 359.5 359.8\n",
" number int64 8B ...\n",
" pressure_level float64 8B 500.0\n",
"Attributes: (12/31)\n",
" GRIB_paramId: 129\n",
" GRIB_dataType: an\n",
" GRIB_numberOfPoints: 1038240\n",
" GRIB_typeOfLevel: isobaricInhPa\n",
" GRIB_stepUnits: 1\n",
" GRIB_stepType: avgua\n",
" ... ...\n",
" GRIB_shortName: z\n",
" GRIB_totalNumber: 0\n",
" GRIB_units: m**2 s**-2\n",
" long_name: Geopotential\n",
" units: m**2 s**-2\n",
" standard_name: geopotential"
],
"text/plain": [
"