Predictions with Aurora 1.5#
In this example, we download ERA5 data for 1 Jan 2023 at 0.25° resolution and run Aurora 1.5 to produce hourly forecasts out to 10 days (240 hours).
Aurora 1.5 extends Aurora with:
26 surface-level variables (including insolation), up from 4
36 static variables (land-surface and vegetation fields)
Variable lead-time embeddings, enabling sub-6-hour output
7 output-only surface variables (e.g. precipitation, wind gusts) that the model predicts but are not present in the input
Running this notebook requires additional Python packages:
pip install cdsapi matplotlib
Downloading the Data#
Register an account with the Climate Data Store and create $HOME/.cdsapirc:
url: https://cds.climate.copernicus.eu/api
key: <API key>
Accept the terms of use on the ERA5 single-levels and ERA5 pressure-levels dataset pages.
We now download surface-level variables, atmospheric variables, and the additional surface fields used by V1.5.
Please note thate ERA5 data are used here for illustrative example purposes only. The model is fine-tuned on IFS data and is therefore not guaranteed to have the same accuracy when using ERA5.
from pathlib import Path
import cdsapi
# Data will be downloaded here.
download_path = Path("~/downloads/era5_v1p5")
download_path = download_path.expanduser()
download_path.mkdir(parents=True, exist_ok=True)
# ----- Surface-level variables (extended set for Aurora 1.5) -----
if not (download_path / "2023-01-01-surface-level.nc").exists():
c = cdsapi.Client()
c.retrieve(
"reanalysis-era5-single-levels",
{
"product_type": "reanalysis",
"variable": [
"2m_temperature",
"10m_u_component_of_wind",
"10m_v_component_of_wind",
"mean_sea_level_pressure",
"2m_dewpoint_temperature",
"total_column_water_vapour",
"total_cloud_cover",
"100m_u_component_of_wind",
"100m_v_component_of_wind",
"surface_pressure",
"low_cloud_cover",
"medium_cloud_cover",
"high_cloud_cover",
"skin_temperature",
"soil_temperature_level_1",
"volumetric_soil_water_layer_1",
"sea_ice_cover",
"snow_depth",
],
"year": "2023",
"month": "01",
"day": "01",
"time": ["00:00", "06:00", "12:00", "18:00"],
"data_format": "netcdf",
},
str(download_path / "2023-01-01-surface-level.nc"),
)
print("Surface-level variables downloaded!")
# ----- Atmospheric variables -----
if not (download_path / "2023-01-01-atmospheric.nc").exists():
c = cdsapi.Client()
c.retrieve(
"reanalysis-era5-pressure-levels",
{
"product_type": "reanalysis",
"variable": [
"temperature",
"u_component_of_wind",
"v_component_of_wind",
"specific_humidity",
"geopotential",
],
"pressure_level": [
"50",
"100",
"150",
"200",
"250",
"300",
"400",
"500",
"600",
"700",
"850",
"925",
"1000",
],
"year": "2023",
"month": "01",
"day": "01",
"time": ["00:00", "06:00", "12:00", "18:00"],
"data_format": "netcdf",
},
str(download_path / "2023-01-01-atmospheric.nc"),
)
print("Atmospheric variables downloaded!")
Downloading the Static Variables#
Aurora 1.5 uses an extended set of 36 static variables (land-sea mask, orography, vegetation fields, soil types, etc.). These are provided in a pickle file. We assume it has been saved as static_aurora_v1p5.pkl alongside the checkpoint.
# Path to the static variables pickle and model checkpoint.
# Adjust these to match your local paths.
static_path = "aurora-0.25-v1.5-static.pickle"
checkpoint_path = "aurora-0.25-v1.5.ckpt"
Preparing a Batch#
We convert the downloaded data to an aurora.Batch. Aurora 1.5 requires:
18 input surface variables (+ insolation, computed below)
36 static variables (from the pickle file)
5 atmospheric variables at 13 pressure levels
The CDS variable names differ from the canonical short names used by Aurora. The mapping is applied below.
Note that the 7 “output-only” surface variables (wind gusts, boundary layer height, radiation fluxes, precipitation, snowfall) are not needed in the input — the model and rollout handle zero-padding automatically.
import pickle
from datetime import datetime
import numpy as np
import torch
import xarray as xr
from aurora import Batch, Metadata
from aurora.insolation import insolation
surf_vars_ds = xr.open_dataset(download_path / "2023-01-01-surface-level.nc", engine="netcdf4")
atmos_vars_ds = xr.open_dataset(download_path / "2023-01-01-atmospheric.nc", engine="netcdf4")
# CDS NetCDF variable name -> Aurora canonical short name.
# The CDS API returns variables under short GRIB names.
surf_name_map = {
"t2m": "2t",
"u10": "10u",
"v10": "10v",
"msl": "msl",
"d2m": "2d",
"tcwv": "tcwv",
"tcc": "tcc",
"u100": "100u",
"v100": "100v",
"sp": "sp",
"lcc": "lcc",
"mcc": "mcc",
"hcc": "hcc",
"skt": "skt",
"stl1": "stl1",
"swvl1": "swvl1",
"siconc": "ci",
"sd": "scaled_sd",
}
# Build surface variable tensors: select hours 00:00 and 06:00, insert batch dim.
surf_vars = {}
for nc_name, aurora_name in surf_name_map.items():
data = surf_vars_ds[nc_name].values[:2][None] # (1, 2, H, W)
surf_vars[aurora_name] = torch.from_numpy(np.nan_to_num(data, nan=0.0).astype(np.float32))
# Compute and add insolation for the two input time steps.
lat = surf_vars_ds.latitude.values.astype(np.float32)
lon = surf_vars_ds.longitude.values.astype(np.float32)
input_times = [datetime(2023, 1, 1, 0), datetime(2023, 1, 1, 6)]
sol = np.stack(
[insolation([t], lat, lon, enforce_2d=True)[0] for t in input_times],
axis=0,
)
surf_vars["insolation"] = torch.from_numpy(sol[None].astype(np.float32))
# Build atmospheric variable tensors: (1, 2, 13, H, W).
levels = (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)
atmos_vars = {}
for nc_name, aurora_name in [("t", "t"), ("u", "u"), ("v", "v"), ("q", "q"), ("z", "z")]:
data = atmos_vars_ds[nc_name].sel(pressure_level=list(levels)).values[:2][None]
atmos_vars[aurora_name] = torch.from_numpy(np.nan_to_num(data, nan=0.0).astype(np.float32))
# Load static variables from the pickle file.
with open(static_path, "rb") as f:
static_raw = pickle.load(f)
static_vars = {k: torch.from_numpy(v) for k, v in static_raw.items()}
# Longitude: CDS provides values in [0, 360) already for 0.25° global grids.
lon_vals = surf_vars_ds.longitude.values.astype(np.float64)
if lon_vals.min() < 0:
lon_vals = lon_vals % 360
batch = Batch(
surf_vars=surf_vars,
static_vars=static_vars,
atmos_vars=atmos_vars,
metadata=Metadata(
lat=torch.from_numpy(surf_vars_ds.latitude.values),
lon=torch.from_numpy(lon_vals),
# Select element 1 (06:00) as the "current" time.
time=(input_times[1],),
atmos_levels=levels,
),
)
print(
f"Batch ready: {len(batch.surf_vars)} surf vars, "
f"{len(batch.static_vars)} static vars, "
f"{len(batch.atmos_vars)} atmos vars, "
f"shape {batch.spatial_shape}"
)
Loading and Running the Model#
Aurora 1.5 uses a different set of normalisation statistics for the extended surface and static variables. These are built into the model and applied automatically.
We perform a roll-out for 4 main steps (24 hours) with hourly sub-steps (fine_lead_times=[1, 2, 3, 4, 5, 6]), producing 24 hourly predictions.
For a full 10-day (240 h) forecast, set steps=40.
from aurora import AuroraV1p5, rollout
model = AuroraV1p5()
model.load_checkpoint("microsoft/aurora", "aurora-0.25-v1.5.ckpt")
model.eval()
model = model.to("cuda")
# 4 main AR steps × 6 hourly sub-steps = 24 hourly predictions.
# For a full 240 h forecast, use steps=40.
steps = 4
fine_lead_times = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
with torch.inference_mode():
preds = [
pred.to("cpu")
for pred in rollout(model, batch, steps=steps, fine_lead_times=fine_lead_times)
]
print(f"Generated {len(preds)} hourly predictions.")
Visualising the Results#
We plot the 2-metre temperature at selected lead times alongside ERA5 for comparison. Since we produced hourly output, we can inspect predictions at T+6 h, T+12 h, T+18 h, and T+24 h.
import matplotlib.pyplot as plt
# Show predictions at T+6h, T+12h, T+18h, T+24h.
# fine_lead_times has 6 entries per main step, so indices 5, 11, 17, 23 correspond
# to the 6th, 12th, 18th, and 24th hour.
show_indices = [5, 11, 17, 23]
# ERA5 truth at 12:00, 18:00 on Jan 1 and 00:00, 06:00 on Jan 2.
# The downloaded data has 00:00, 06:00, 12:00, 18:00 on Jan 1.
# Index 2 = 12:00 (T+6h from init 06:00), Index 3 = 18:00 (T+12h).
# For T+18h and T+24h we would need Jan 2 data; we only plot what we have.
era5_indices = [2, 3] # 12:00, 18:00
fig, ax = plt.subplots(2, 2, figsize=(12, 6.5))
for row, (pred_idx, era5_idx) in enumerate(zip(show_indices[:2], era5_indices)):
pred = preds[pred_idx]
lead_hours = pred_idx + 1 # 1-indexed hours
ax[row, 0].imshow(pred.surf_vars["2t"][0, 0].numpy() - 273.15, vmin=-50, vmax=50)
ax[row, 0].set_ylabel(f"T+{lead_hours}h\n{pred.metadata.time[0]}")
if row == 0:
ax[row, 0].set_title("Aurora 1.5 Prediction")
ax[row, 0].set_xticks([])
ax[row, 0].set_yticks([])
ax[row, 1].imshow(surf_vars_ds["t2m"][era5_idx].values - 273.15, vmin=-50, vmax=50)
if row == 0:
ax[row, 1].set_title("ERA5")
ax[row, 1].set_xticks([])
ax[row, 1].set_yticks([])
plt.tight_layout()
plt.show()