Add Aurora v1.5 deterministic and ensemble prognostic model wrappers#965
Open
swbg wants to merge 4 commits into
Open
Add Aurora v1.5 deterministic and ensemble prognostic model wrappers#965swbg wants to merge 4 commits into
swbg wants to merge 4 commits into
Conversation
Adds AuroraV1p5 and AuroraV1p5Ensemble wrappers for the Microsoft Aurora v1.5 0.25-degree global weather model. The wrappers use an hourly rollout: the 6-hour auto-regressive base step is queried at lead times t+1h through t+6h before advancing the AR state, yielding 90-variable outputs (65 atmos + 18 surface + 7 output-only diagnostics). Bit-perfect agreement with the official aurora.rollout() path verified across 12 hourly steps (2 full AR cycles). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Collaborator
Author
Reference Comparison ValidationModel: Initial condition: ERA5 at 2025-01-01T00:00 UTC (ARCO GCS zarr store, 0.25°). Results Summary
Full per-step table (all 12 steps): Key findings:
Reference Plots
Validation ScriptsVanilla reference script"""Vanilla Aurora v1.5 inference — NO Earth2Studio imports.
Uses the official aurora.rollout() function and fetches initial conditions
directly from the ARCO ERA5 zarr store on GCS (no Earth2Studio data utilities).
Runs 12 hourly steps (2 full AR cycles of 6 hours each).
Saves output tensors to reference_aurora_v1p5_vanilla.pt.
Requires: pip install microsoft-aurora>=2.0.0 gcsfs zarr
"""
import os
import pickle
from datetime import datetime, timezone
import gcsfs
import numpy as np
import torch
import zarr
from aurora import AuroraV1p5 as AuroraModel
from aurora import Batch, Metadata
from aurora.insolation import insolation as aurora_insolation
from aurora.normalisation import log_untransform as aurora_log_untransform
from aurora.rollout import rollout
DEVICE = "cuda:0"
T0 = np.datetime64("2025-01-01T00:00")
T_MINUS_6 = T0 - np.timedelta64(6, "h")
ATMOS_LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50]
ARCO_PATH = "/gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
ARCO_EPOCH = np.datetime64("1900-01-01T00:00")
# Populate via: python -c "from aurora import AuroraV1p5; AuroraV1p5.load_model(AuroraV1p5.load_default_package())"
# or run the E2S script first so the cache is warm.
CACHE_DIR = os.path.expanduser("~/.cache/earth2studio/aurora-v1p5")
CKPT_PATH = os.path.join(CACHE_DIR, "aurora-0.25-v1.5.ckpt")
STATIC_PATH = os.path.join(CACHE_DIR, "aurora-0.25-v1.5-static.pickle")
_AURORA_TO_ARCO_SURF = {
"msl": "mean_sea_level_pressure",
"10u": "10m_u_component_of_wind",
"10v": "10m_v_component_of_wind",
"2t": "2m_temperature",
"2d": "2m_dewpoint_temperature",
"tcwv": "total_column_water_vapour",
"tcc": "total_cloud_cover",
"100u": "100m_u_component_of_wind",
"100v": "100m_v_component_of_wind",
"sp": "surface_pressure",
"lcc": "low_cloud_cover",
"mcc": "medium_cloud_cover",
"hcc": "high_cloud_cover",
"skt": "skin_temperature",
"stl1": "soil_temperature_level_1",
"swvl1": "volumetric_soil_water_layer_1",
"ci": "sea_ice_cover",
"scaled_sd": "snow_depth",
}
_AURORA_TO_ARCO_ATMOS = {
"z": "geopotential",
"q": "specific_humidity",
"t": "temperature",
"u": "u_component_of_wind",
"v": "v_component_of_wind",
}
_OUTPUT_ONLY_SURF_VARS = [
("i10fg", "i10fg", False),
("blh", "blh", False),
("uvb1h", "uvb_1h", False),
("ssrd1h", "ssrd_1h", False),
("ttr1h", "ttr_1h", False),
("tp1h", "scaled_tp_1h", True),
("sf1h", "scaled_sf_1h", True),
]
def ts_to_datetime(ts: np.datetime64) -> datetime:
epoch = np.datetime64("1970-01-01T00:00:00")
seconds = float((ts.astype("datetime64[s]") - epoch) / np.timedelta64(1, "s"))
return datetime.fromtimestamp(seconds, tz=timezone.utc)
def arco_time_index(ts: np.datetime64) -> int:
return int((ts - ARCO_EPOCH) / np.timedelta64(1, "h"))
print("Loading Aurora v1.5 weights from local cache...")
model = AuroraModel()
model.load_checkpoint_local(CKPT_PATH)
model.eval().to(DEVICE)
print("Loading static vars from pickle...")
with open(STATIC_PATH, "rb") as f:
static_raw = pickle.load(f) # noqa: S301
static_vars = {k: torch.from_numpy(np.asarray(v))[:720, :] for k, v in static_raw.items()}
LAT = np.linspace(90, -90, 720, endpoint=False)
LON = np.linspace(0, 359.75, 1440)
print("Opening ARCO zarr store...")
fs = gcsfs.GCSFileSystem(token="anon", asynchronous=False)
zstore = zarr.storage.FsspecStore(fs, path=ARCO_PATH)
zg = zarr.open(store=zstore, mode="r")
all_levels = zg["level"][:]
level_indices = [int(np.searchsorted(all_levels, lv)) for lv in ATMOS_LEVELS]
ti_minus6 = arco_time_index(T_MINUS_6)
ti_0 = arco_time_index(T0)
atmos_data = {}
for aurora_name, arco_name in _AURORA_TO_ARCO_ATMOS.items():
arr = zg[arco_name]
chunk = np.stack(
[arr[ti_minus6, level_indices, :720, :], arr[ti_0, level_indices, :720, :]],
axis=0,
).astype(np.float32)
atmos_data[aurora_name] = torch.from_numpy(chunk).unsqueeze(0)
surf_data = {}
for aurora_name, arco_name in _AURORA_TO_ARCO_SURF.items():
arr = zg[arco_name]
chunk = np.stack(
[arr[ti_minus6, :720, :], arr[ti_0, :720, :]],
axis=0,
).astype(np.float32)
surf_data[aurora_name] = torch.from_numpy(chunk).unsqueeze(0)
for k in surf_data:
surf_data[k] = torch.nan_to_num(surf_data[k], nan=0.0)
dt0 = ts_to_datetime(T_MINUS_6)
dt1 = ts_to_datetime(T0)
insol = aurora_insolation((dt0, dt1), LAT, LON, enforce_2d=True)
surf_data["insolation"] = torch.from_numpy(np.asarray(insol)).float().unsqueeze(0)
batch = Batch(
surf_vars=surf_data,
static_vars=static_vars,
atmos_vars=atmos_data,
metadata=Metadata(
lat=torch.from_numpy(LAT).float(),
lon=torch.from_numpy(LON).float(),
time=(dt1,),
atmos_levels=tuple(int(lv) for lv in ATMOS_LEVELS),
rollout_step=0,
),
)
print("Running official aurora.rollout() — 2 AR cycles, fine_lead_times=[1..6]...")
sub_outputs = {}
with torch.inference_mode():
for h, pred in enumerate(
rollout(model, batch, steps=2, fine_lead_times=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]),
start=1,
):
atmos_out = torch.cat(
[pred.atmos_vars["z"], pred.atmos_vars["q"], pred.atmos_vars["t"],
pred.atmos_vars["u"], pred.atmos_vars["v"]],
dim=2,
)
surf_keys_ordered = list(_AURORA_TO_ARCO_SURF.keys())
surf_out = torch.cat(
[pred.surf_vars[k].unsqueeze(2) for k in surf_keys_ordered],
dim=2,
)
diag_tensors = []
for _e2s_name, aurora_name, log_scaled in _OUTPUT_ONLY_SURF_VARS:
v = pred.surf_vars[aurora_name]
if log_scaled:
v = aurora_log_untransform(v)
diag_tensors.append(v.unsqueeze(2))
diag_out = torch.cat(diag_tensors, dim=2)
out_tensor = torch.cat([atmos_out, surf_out, diag_out], dim=2).cpu()
sub_outputs[h] = out_tensor
torch.save(sub_outputs, "reference_aurora_v1p5_vanilla.pt")
print("Saved reference_aurora_v1p5_vanilla.pt")Earth2Studio reference script"""Earth2Studio AuroraV1p5 wrapper inference.
Fetches the same ERA5 initial conditions from ARCO as the vanilla script,
runs 12 hourly steps (2 full AR cycles) using the E2S AuroraV1p5 wrapper via
create_iterator, and saves the output tensors to reference_aurora_v1p5_e2s.pt.
Requires: pip install earth2studio[aurora-v1p5]
"""
import numpy as np
import torch
from earth2studio.data import ARCO
from earth2studio.data.utils import fetch_data
from earth2studio.models.px.aurora_v1p5 import VARIABLES, AuroraV1p5
from earth2studio.utils.coords import map_coords
DEVICE = "cuda:0"
TIME = np.array([np.datetime64("2025-01-01T00:00")])
N_STEPS = 12
print("Loading AuroraV1p5 via E2S...")
package = AuroraV1p5.load_default_package()
model = AuroraV1p5.load_model(package).to(DEVICE)
model.eval()
print("Fetching ARCO ERA5 data...")
data = ARCO()
lead_time = model.input_coords()["lead_time"]
variable = np.array(VARIABLES)
x, coords = fetch_data(data, TIME, variable, lead_time, device=DEVICE)
x, coords = map_coords(x, coords, model.input_coords())
print(f"Running {N_STEPS} E2S steps via create_iterator...")
sub_outputs = {}
with torch.inference_mode():
p_iter = model.create_iterator(x, coords)
next(p_iter) # skip IC yield (h+0 with NaN output-only channels)
for h, (out, _) in enumerate(p_iter, start=1):
result = out.cpu().view(1, 1, 90, 720, 1440)
sub_outputs[h] = result
print(f" h+{h}: shape={result.shape} max={result.abs().max():.4f}")
if h >= N_STEPS:
break
torch.save(sub_outputs, "reference_aurora_v1p5_e2s.pt")
print("Saved reference_aurora_v1p5_e2s.pt")Comparison script"""Compare vanilla Aurora v1.5 outputs against Earth2Studio wrapper outputs.
Loads reference_aurora_v1p5_vanilla.pt and reference_aurora_v1p5_e2s.pt and
reports max absolute difference, relative error, and correlation for each of
the 12 hourly steps (2 full AR cycles). Also generates a comparison plot
(t2m at h+1, h+6, h+12; rows: E2S / Vanilla / |diff|).
Relative error denominator: max(abs(vanilla), 1e-6).
"""
import numpy as np
import torch
from earth2studio.models.px.aurora_v1p5 import VARIABLES
vanilla = torch.load("reference_aurora_v1p5_vanilla.pt", map_location="cpu")
e2s = torch.load("reference_aurora_v1p5_e2s.pt", map_location="cpu")
N_STEPS = len(vanilla)
T2M_IDX = list(VARIABLES).index("t2m")
print(f"{'Hour':>4} {'MaxAbsErr':>12} {'RelErr':>12} {'Corr':>10} {'NaN(v)':>8} {'NaN(e)':>8}")
print("-" * 65)
corrs = []
for h in range(1, N_STEPS + 1):
v = vanilla[h].float() # (1, 1, 90, 720, 1440)
e = e2s[h].float() # (1, 1, 90, 720, 1440)
diff = (v - e).abs()
max_err = diff.max().item()
denom = v.abs().clamp(min=1e-6)
rel_err = (diff / denom).max().item()
v_flat = v.flatten()
e_flat = e.flatten()
corr = torch.corrcoef(torch.stack([v_flat, e_flat]))[0, 1].item()
corrs.append(corr)
nan_v = v.isnan().sum().item()
nan_e = e.isnan().sum().item()
print(f" h+{h:02d} {max_err:>12.6f} {rel_err:>12.6f} {corr:>10.8f} {nan_v:>8} {nan_e:>8}")
print()
print(f"Mean correlation across {N_STEPS} steps: {np.mean(corrs):.8f}")
print(f"Min correlation across {N_STEPS} steps: {np.min(corrs):.8f}")
try:
import matplotlib.pyplot as plt
steps = [1, 6, N_STEPS]
fig, axes = plt.subplots(3, len(steps), figsize=(15, 9))
fig.suptitle("AuroraV1p5: E2S vs Vanilla t2m (K)", fontsize=13)
vmin, vmax = 220, 320
for col, h in enumerate(steps):
v_t2m = vanilla[h][0, 0, T2M_IDX].numpy()
e_t2m = e2s[h][0, 0, T2M_IDX].numpy()
err = np.abs(v_t2m - e_t2m)
axes[0, col].imshow(e_t2m, vmin=vmin, vmax=vmax, origin="upper", aspect="auto")
axes[0, col].set_title(f"E2S h+{h}")
axes[0, col].axis("off")
axes[1, col].imshow(v_t2m, vmin=vmin, vmax=vmax, origin="upper", aspect="auto")
axes[1, col].set_title(f"Vanilla h+{h}")
axes[1, col].axis("off")
im = axes[2, col].imshow(err, origin="upper", aspect="auto", cmap="Reds")
axes[2, col].set_title(f"|E2S - Vanilla| h+{h}")
axes[2, col].axis("off")
plt.colorbar(im, ax=axes[2, col], fraction=0.03)
plt.tight_layout()
plt.savefig("reference_aurora_v1p5_compare.png", dpi=120, bbox_inches="tight")
print("Saved reference_aurora_v1p5_compare.png")
plt.close()
except ImportError:
print("matplotlib not available — skipping plot") |
Contributor
Greptile SummaryThis PR adds Aurora v1.5 prognostic model support. The main changes are:
|
| Filename | Overview |
|---|---|
| earth2studio/models/px/aurora_v1p5.py | Adds the Aurora v1.5 wrappers, hourly rollout logic, AR clipping, output conversion, and ensemble noise reset; rollout-step state persists across independent iterators. |
| earth2studio/models/px/init.py | Exports the new Aurora v1.5 model classes from the prognostic model package. |
| pyproject.toml | Adds the aurora-v1p5 extra and includes it in the all extra. |
| test/models/px/test_aurora_v1p5.py | Adds unit and package tests for first-use call and iterator behavior. |
| earth2studio/lexicon/base.py | Adds vocabulary entries for Aurora v1.5 diagnostic variables. |
| earth2studio/lexicon/ncar.py | Adds NCAR ERA5 mappings for additional Aurora v1.5 surface variables. |
| docs/modules/models_px.rst | Adds the new Aurora v1.5 wrappers to the model documentation index. |
| docs/userguide/about/install.md | Documents installation for the new Aurora v1.5 extra. |
| CHANGELOG.md | Records the new Aurora v1.5 deterministic and ensemble wrappers. |
Reviews (1): Last reviewed commit: "Add Aurora v1.5 deterministic and ensemb..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Reset preds_idx to 0 at the start of _default_generator so each create_iterator() call starts Aurora's rollout_step counter fresh. Without this, a second forecast on the same instance carries the previous cycle count forward, causing positive-variable clamping to activate at cycle 1 instead of cycle 2 and producing wrong predictions. Add test_aurora_v1p5_iter_repeated to guard against regressions. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Collaborator
Author
|
Re: Greptile P1 — Shared Rollout Step State Fixed in commit 0830e27. |
Collaborator
Author
|
Caveat: Aurora 1.5 is pretrained on ERA5/IFS including sea ice are fraction, which is not available from the IFS public data archive. Users can either run the model on ERA5 (for historical studies, not recommended because the model is finetuned on IFS) or bring their own IFS data. |
NickGeneva
reviewed
Jul 15, 2026
NickGeneva
reviewed
Jul 15, 2026
NickGeneva
reviewed
Jul 15, 2026
| aurora = [ | ||
| "microsoft-aurora>=1.6.0", | ||
| ] | ||
| aurora-v1p5 = [ |
Collaborator
There was a problem hiding this comment.
instead of adding a new group, can we bump aurora groups requirement to 2.0? Is aurora 1 still runnable with v 2.0?
Collaborator
Author
There was a problem hiding this comment.
Seems possible based on the code and I have run a quick sanity check with the original Aurora:
All 6 steps (h+6 through h+36) pass. Summary:
| Variable | Range across all steps | Status |
|---|---|---|
| t2m | [214.8, 316.0] K | OK — realistic global surface temps |
| u10m | [−20.7, +25.4] m/s | OK — moderate surface winds |
| z500 | [46768, 58064] m²/s² | OK — typical mid-latitude range |
Zero NaN/Inf throughout.
NickGeneva
reviewed
Jul 15, 2026
| @@ -153,6 +153,9 @@ atlas = [ | |||
| aurora = [ | |||
NickGeneva
reviewed
Jul 15, 2026
NickGeneva
reviewed
Jul 15, 2026
| self, | ||
| x: torch.Tensor, | ||
| coords: CoordSystem, | ||
| lead_time_hours: float = 1.0, |
Collaborator
There was a problem hiding this comment.
float or should this be an int
Collaborator
Author
There was a problem hiding this comment.
I have made this an int now and pytorch cast it to float ahead of passing it to the model.
NickGeneva
reviewed
Jul 15, 2026
| @check_optional_dependencies() | ||
| def load_model( | ||
| cls, | ||
| package: Package, |
Collaborator
There was a problem hiding this comment.
typically for ens models we need a seed parameter of sorts to make it reproducable, can you see if this is possible?
E.g. http://github.com/NVIDIA/earth2studio/blob/main/earth2studio/models/px/aifs2ens.py#L433
Collaborator
There was a problem hiding this comment.
Thanks for adding the sanity check scripts, this makes be very confident in the determistic model. Can you perhaps produce similar scripts for the ENS model since youre adding both there.
Good test for getting the random state controlled
Collaborator
Author
There was a problem hiding this comment.
I made it a bit of a mix of aifs2ens and fcn3 now, calling torch.manual_seed and exposing a set_rng method (to have a similar interface to fcn3). The Aurora repo simply calls torch.rand at some point without a much better way to set a seed I think. There is some difference between running the Earth2Studio wrapper vs. the original repo (because of different GPU ops like cat etc. after calling torch.manual_seed) but runs with the same seed are reproducible within Earth2Studio.
NickGeneva
reviewed
Jul 15, 2026
| grid (720, 1440) with 5 atmospheric variables across 13 pressure levels and | ||
| 18 surface variables plus 7 output-only surface variables. | ||
|
|
||
| This wrapper uses an hourly rollout by default: the underlying 6-hour |
Collaborator
There was a problem hiding this comment.
any reason for it to be 6 hour batched? Just curious, also seem kinda implmentation detail, not sure if needed for doc string
Collaborator
Author
There was a problem hiding this comment.
The same batch is used to generate 6 consecutive hours (t-6, t+0) > (t+1, t+2, t+3, t+4, t+5, t+6). The rollout advances in groups of 6 so it doesn't have to prepare the same batch over and over again. I've set the default interval to 1 hour, in theory one could advance 2 or 3 hours each step. I have adjusted the docstrings though, seemed a bit too detailed in some parts.
NickGeneva
reviewed
Jul 15, 2026
NickGeneva
left a comment
NickGeneva
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for using the skill Stefan, left a few initial comments. Looking good
Collaborator
Author
Reference Comparison ValidationModels: Initial condition: ERA5 at 2025-01-01T00:00 UTC (ARCO GCS zarr store, 0.25°). Aurora1p5 (Deterministic) Results
Full 12-step table (all steps: bit-perfect): Key findings:
Aurora1p5Ensemble (Stochastic, seed=42) ResultsExact bit-level agreement between a stochastic E2S wrapper and the vanilla Per-variable MaxAbsErr (seed=42, 1 member): Reproducibility check — two Key findings:
Reference Plots
Validation ScriptsAurora1p5 — vanilla reference script"""Vanilla Aurora v1.5 inference — NO Earth2Studio imports.
Uses the official aurora.rollout() function and fetches initial conditions
directly from the ARCO ERA5 zarr store on GCS (no Earth2Studio data utilities).
Runs 12 hourly steps (2 full AR cycles of 6 hours each).
Saves output tensors to reference_aurora_v1p5_vanilla.pt.
Run with:
uv run python reference_aurora_v1p5_vanilla.py
"""
import os
import pickle
from datetime import datetime, timezone
import gcsfs
import numpy as np
import torch
import zarr
from aurora import AuroraV1p5 as AuroraModel
from aurora import Batch, Metadata
from aurora.insolation import insolation as aurora_insolation
from aurora.normalisation import log_untransform as aurora_log_untransform
from aurora.rollout import rollout
DEVICE = "cuda:0"
T0 = np.datetime64("2025-01-01T00:00")
T_MINUS_6 = T0 - np.timedelta64(6, "h")
ATMOS_LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50]
ARCO_PATH = "/gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
ARCO_EPOCH = np.datetime64("1900-01-01T00:00")
CACHE_DIR = os.path.expanduser("~/.cache/earth2studio/aurora1p5")
CKPT_PATH = os.path.join(CACHE_DIR, "aurora-0.25-v1.5.ckpt")
STATIC_PATH = os.path.join(CACHE_DIR, "aurora-0.25-v1.5-static.pickle")
_AURORA_TO_ARCO_SURF = {
"msl": "mean_sea_level_pressure",
"10u": "10m_u_component_of_wind",
"10v": "10m_v_component_of_wind",
"2t": "2m_temperature",
"2d": "2m_dewpoint_temperature",
"tcwv": "total_column_water_vapour",
"tcc": "total_cloud_cover",
"100u": "100m_u_component_of_wind",
"100v": "100m_v_component_of_wind",
"sp": "surface_pressure",
"lcc": "low_cloud_cover",
"mcc": "medium_cloud_cover",
"hcc": "high_cloud_cover",
"skt": "skin_temperature",
"stl1": "soil_temperature_level_1",
"swvl1": "volumetric_soil_water_layer_1",
"ci": "sea_ice_cover",
"scaled_sd": "snow_depth",
}
_AURORA_TO_ARCO_ATMOS = {
"z": "geopotential",
"q": "specific_humidity",
"t": "temperature",
"u": "u_component_of_wind",
"v": "v_component_of_wind",
}
_OUTPUT_ONLY_SURF_VARS = [
("i10fg", "i10fg", False),
("blh", "blh", False),
("uvb1h", "uvb_1h", False),
("ssrd1h", "ssrd_1h", False),
("ttr1h", "ttr_1h", False),
("tp1h", "scaled_tp_1h", True),
("sf1h", "scaled_sf_1h", True),
]
def ts_to_datetime(ts: np.datetime64) -> datetime:
epoch = np.datetime64("1970-01-01T00:00:00")
seconds = float((ts.astype("datetime64[s]") - epoch) / np.timedelta64(1, "s"))
return datetime.fromtimestamp(seconds, tz=timezone.utc)
def arco_time_index(ts: np.datetime64) -> int:
return int((ts - ARCO_EPOCH) / np.timedelta64(1, "h"))
print("Loading Aurora v1.5 weights from local cache...")
model = AuroraModel()
model.load_checkpoint_local(CKPT_PATH)
model.eval().to(DEVICE)
print("Loading static vars from pickle...")
with open(STATIC_PATH, "rb") as f:
static_raw = pickle.load(f) # noqa: S301
static_vars = {k: torch.from_numpy(np.asarray(v))[:720, :] for k, v in static_raw.items()}
LAT = np.linspace(90, -90, 720, endpoint=False)
LON = np.linspace(0, 359.75, 1440)
print("Opening ARCO zarr store...")
fs = gcsfs.GCSFileSystem(token="anon", asynchronous=False)
zstore = zarr.storage.FsspecStore(fs, path=ARCO_PATH)
zg = zarr.open(store=zstore, mode="r")
all_levels = zg["level"][:]
level_indices = [int(np.searchsorted(all_levels, lv)) for lv in ATMOS_LEVELS]
ti_minus6 = arco_time_index(T_MINUS_6)
ti_0 = arco_time_index(T0)
atmos_data = {}
for aurora_name, arco_name in _AURORA_TO_ARCO_ATMOS.items():
arr = zg[arco_name]
chunk = np.stack(
[arr[ti_minus6, level_indices, :720, :], arr[ti_0, level_indices, :720, :]],
axis=0,
).astype(np.float32)
atmos_data[aurora_name] = torch.from_numpy(chunk).unsqueeze(0)
surf_data = {}
for aurora_name, arco_name in _AURORA_TO_ARCO_SURF.items():
arr = zg[arco_name]
chunk = np.stack(
[arr[ti_minus6, :720, :], arr[ti_0, :720, :]],
axis=0,
).astype(np.float32)
surf_data[aurora_name] = torch.from_numpy(chunk).unsqueeze(0)
for k in surf_data:
surf_data[k] = torch.nan_to_num(surf_data[k], nan=0.0)
dt0 = ts_to_datetime(T_MINUS_6)
dt1 = ts_to_datetime(T0)
insol = aurora_insolation((dt0, dt1), LAT, LON, enforce_2d=True)
surf_data["insolation"] = torch.from_numpy(np.asarray(insol)).float().unsqueeze(0)
batch = Batch(
surf_vars=surf_data,
static_vars=static_vars,
atmos_vars=atmos_data,
metadata=Metadata(
lat=torch.from_numpy(LAT).float(),
lon=torch.from_numpy(LON).float(),
time=(dt1,),
atmos_levels=tuple(int(lv) for lv in ATMOS_LEVELS),
rollout_step=0,
),
)
print("Running official aurora.rollout() — 2 AR cycles, fine_lead_times=[1..6]...")
sub_outputs = {}
with torch.inference_mode():
for h, pred in enumerate(
rollout(model, batch, steps=2, fine_lead_times=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]),
start=1,
):
atmos_out = torch.cat(
[pred.atmos_vars["z"], pred.atmos_vars["q"], pred.atmos_vars["t"],
pred.atmos_vars["u"], pred.atmos_vars["v"]],
dim=2,
)
surf_keys_ordered = list(_AURORA_TO_ARCO_SURF.keys())
surf_out = torch.cat(
[pred.surf_vars[k].unsqueeze(2) for k in surf_keys_ordered],
dim=2,
)
diag_tensors = []
for _, aurora_name, log_scaled in _OUTPUT_ONLY_SURF_VARS:
v = pred.surf_vars[aurora_name]
if log_scaled:
v = aurora_log_untransform(v)
diag_tensors.append(v.unsqueeze(2))
diag_out = torch.cat(diag_tensors, dim=2)
out_tensor = torch.cat([atmos_out, surf_out, diag_out], dim=2).cpu()
sub_outputs[h] = out_tensor
torch.save(sub_outputs, "reference_aurora_v1p5_vanilla.pt")
print("Saved reference_aurora_v1p5_vanilla.pt")Aurora1p5 — Earth2Studio reference script"""Earth2Studio Aurora1p5 wrapper inference.
Fetches the same ERA5 initial conditions from ARCO as the vanilla script,
runs 12 hourly steps (2 full AR cycles) using the E2S Aurora1p5 wrapper via
create_iterator, and saves the output tensors to reference_aurora_v1p5_e2s.pt.
Run with:
uv run python reference_aurora_v1p5_e2s.py
"""
import numpy as np
import torch
from earth2studio.data import ARCO
from earth2studio.data.utils import fetch_data
from earth2studio.models.px.aurora1p5 import VARIABLES, Aurora1p5
from earth2studio.utils.coords import map_coords
DEVICE = "cuda:0"
TIME = np.array([np.datetime64("2025-01-01T00:00")])
N_STEPS = 12
print("Loading Aurora1p5 via E2S...")
package = Aurora1p5.load_default_package()
model = Aurora1p5.load_model(package).to(DEVICE)
model.eval()
print("Fetching ARCO ERA5 data...")
data = ARCO()
lead_time = model.input_coords()["lead_time"]
variable = np.array(VARIABLES)
x, coords = fetch_data(data, TIME, variable, lead_time, device=DEVICE)
x, coords = map_coords(x, coords, model.input_coords())
print(f"Running {N_STEPS} E2S steps via create_iterator...")
sub_outputs = {}
with torch.inference_mode():
p_iter = model.create_iterator(x, coords)
next(p_iter) # skip IC yield (h+0 with NaN output-only channels)
for h, (out, _) in enumerate(p_iter, start=1):
result = out.cpu().view(1, 1, 90, 720, 1440)
sub_outputs[h] = result
if h >= N_STEPS:
break
torch.save(sub_outputs, "reference_aurora_v1p5_e2s.pt")
print("Saved reference_aurora_v1p5_e2s.pt")Aurora1p5 — comparison script"""Compare vanilla Aurora v1.5 outputs against Earth2Studio wrapper outputs.
Loads reference_aurora_v1p5_vanilla.pt and reference_aurora_v1p5_e2s.pt and
reports max absolute difference, relative error, and correlation for each of
the 12 hourly steps. Also generates a comparison plot (t2m at h+1, h+6, h+12;
rows: E2S / Vanilla / |diff|).
Relative error denominator: max(abs(vanilla), 1e-6).
Run with:
uv run python reference_aurora_v1p5_compare.py
"""
import numpy as np
import torch
from earth2studio.models.px.aurora1p5 import OUTPUT_VARIABLES as VARIABLES
vanilla = torch.load("reference_aurora_v1p5_vanilla.pt", map_location="cpu", weights_only=False)
e2s = torch.load("reference_aurora_v1p5_e2s.pt", map_location="cpu", weights_only=False)
N_STEPS = len(vanilla)
T2M_IDX = list(VARIABLES).index("t2m")
print(f"{'Hour':>4} {'MaxAbsErr':>12} {'RelErr':>12} {'Corr':>10} {'NaN(v)':>8} {'NaN(e)':>8}")
print("-" * 65)
corrs = []
for h in range(1, N_STEPS + 1):
v = vanilla[h].float()
e = e2s[h].float()
diff = (v - e).abs()
max_err = diff.max().item()
denom = v.abs().clamp(min=1e-6)
rel_err = (diff / denom).max().item()
corr = torch.corrcoef(torch.stack([v.flatten(), e.flatten()]))[0, 1].item()
corrs.append(corr)
nan_v = v.isnan().sum().item()
nan_e = e.isnan().sum().item()
print(f" h+{h:02d} {max_err:>12.6f} {rel_err:>12.6f} {corr:>10.8f} {nan_v:>8} {nan_e:>8}")
print()
print(f"Mean correlation across {N_STEPS} steps: {np.mean(corrs):.8f}")
print(f"Min correlation across {N_STEPS} steps: {np.min(corrs):.8f}")
try:
import matplotlib.pyplot as plt
steps = [1, 6, N_STEPS]
fig, axes = plt.subplots(3, len(steps), figsize=(15, 9))
fig.suptitle("Aurora1p5: E2S vs Vanilla t2m (K)", fontsize=13)
vmin, vmax = 220, 320
for col, h in enumerate(steps):
v_t2m = vanilla[h][0, 0, T2M_IDX].numpy()
e_t2m = e2s[h][0, 0, T2M_IDX].numpy()
err = np.abs(v_t2m - e_t2m)
axes[0, col].imshow(e_t2m, vmin=vmin, vmax=vmax, origin="upper", aspect="auto")
axes[0, col].set_title(f"E2S h+{h}")
axes[0, col].axis("off")
axes[1, col].imshow(v_t2m, vmin=vmin, vmax=vmax, origin="upper", aspect="auto")
axes[1, col].set_title(f"Vanilla h+{h}")
axes[1, col].axis("off")
im = axes[2, col].imshow(err, origin="upper", aspect="auto", cmap="Reds")
axes[2, col].set_title(f"|E2S - Vanilla| h+{h}")
axes[2, col].axis("off")
plt.colorbar(im, ax=axes[2, col], fraction=0.03)
plt.tight_layout()
plt.savefig("reference_aurora_v1p5_compare.png", dpi=120, bbox_inches="tight")
plt.close()
except ImportError:
print("matplotlib not available — skipping plot")Aurora1p5 — sanity check script"""Sanity-check plots for Aurora1p5 using ARCO ERA5 initial conditions.
Runs 12 hourly forecast steps (2 AR cycles) and plots t2m, msl, and z500
at h+6 and h+12. Checks for NaN/Inf in all outputs.
Run with:
uv run python reference_aurora_v1p5_sanity.py
"""
import numpy as np
import torch
from earth2studio.data import ARCO
from earth2studio.data.utils import fetch_data
from earth2studio.models.px.aurora1p5 import VARIABLES, Aurora1p5
from earth2studio.utils.coords import map_coords
DEVICE = "cuda:0"
TIME = np.array([np.datetime64("2025-01-01T00:00")])
N_STEPS = 12
PLOT_VARS = {
"t2m": {"idx": list(VARIABLES).index("t2m"), "vmin": 220, "vmax": 320, "label": "t2m (K)"},
"msl": {"idx": list(VARIABLES).index("msl"), "vmin": 97000, "vmax": 103000, "label": "msl (Pa)"},
"z500": {"idx": list(VARIABLES).index("z500"), "vmin": 49000, "vmax": 59000, "label": "z500 (m²/s²)"},
}
package = Aurora1p5.load_default_package()
model = Aurora1p5.load_model(package).to(DEVICE)
model.eval()
data = ARCO()
lead_time = model.input_coords()["lead_time"]
variable = np.array(VARIABLES)
x, coords = fetch_data(data, TIME, variable, lead_time, device=DEVICE)
x, coords = map_coords(x, coords, model.input_coords())
outputs = {}
iterator = model.create_iterator(x, coords)
next(iterator)
with torch.inference_mode():
for out, out_coords in iterator:
h = int(out_coords["lead_time"][0] / np.timedelta64(1, "h"))
outputs[h] = out[0, 0].cpu()
if h >= N_STEPS:
break
try:
import matplotlib.pyplot as plt
plot_steps = [6, 12]
fig, axes = plt.subplots(len(PLOT_VARS), len(plot_steps), figsize=(12, len(PLOT_VARS) * 3))
fig.suptitle("Aurora1p5 sanity check — IC: 2025-01-01T00Z", fontsize=12)
for row, (var_name, meta) in enumerate(PLOT_VARS.items()):
for col, h in enumerate(plot_steps):
field = outputs[h][meta["idx"]].numpy()
im = axes[row, col].imshow(
field, vmin=meta["vmin"], vmax=meta["vmax"], origin="upper", aspect="auto"
)
axes[row, col].set_title(f"{meta['label']} h+{h}")
axes[row, col].axis("off")
plt.colorbar(im, ax=axes[row, col], fraction=0.03)
plt.tight_layout()
plt.savefig("reference_aurora_v1p5_sanity.png", dpi=120, bbox_inches="tight")
plt.close()
except ImportError:
print("matplotlib not available — skipping plot")
total_nan = sum(t.isnan().sum().item() for t in outputs.values())
total_inf = sum(t.isinf().sum().item() for t in outputs.values())
print(f"Total NaN: {total_nan} Total Inf: {total_inf}")
t2m_idx = list(VARIABLES).index("t2m")
t2m_range = [outputs[h][t2m_idx] for h in sorted(outputs)]
print(f"t2m range h+1…h+{N_STEPS}: [{min(t.min() for t in t2m_range):.1f}, {max(t.max() for t in t2m_range):.1f}] K")Aurora1p5Ensemble — vanilla reference script"""Vanilla Aurora v1.5 ensemble inference — NO Earth2Studio imports.
Uses aurora.rollout() with AuroraV1p5Ensemble and fetches initial conditions
directly from the ARCO ERA5 zarr store on GCS. Runs 12 hourly steps (2 AR
cycles). Seeds with torch.manual_seed + reset_noise to match the E2S
wrapper's set_rng behaviour.
Saves output tensors to reference_aurora_v1p5_ens_vanilla.pt.
Run with:
uv run python reference_aurora_v1p5_ens_vanilla.py
"""
import os
import pickle
from datetime import datetime, timezone
import gcsfs
import numpy as np
import torch
import zarr
from aurora import AuroraV1p5Ensemble as AuroraEnsembleModel
from aurora import Batch, Metadata
from aurora.insolation import insolation as aurora_insolation
from aurora.normalisation import log_untransform as aurora_log_untransform
from aurora.rollout import rollout
DEVICE = "cuda:0"
SEED = 42
T0 = np.datetime64("2025-01-01T00:00")
T_MINUS_6 = T0 - np.timedelta64(6, "h")
ATMOS_LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50]
ARCO_PATH = "/gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
ARCO_EPOCH = np.datetime64("1900-01-01T00:00")
CACHE_DIR = os.path.expanduser("~/.cache/earth2studio/aurora1p5")
CKPT_PATH = os.path.join(CACHE_DIR, "aurora-0.25-v1.5-ensemble.ckpt")
STATIC_PATH = os.path.join(CACHE_DIR, "aurora-0.25-v1.5-static.pickle")
_AURORA_TO_ARCO_SURF = {
"msl": "mean_sea_level_pressure", "10u": "10m_u_component_of_wind",
"10v": "10m_v_component_of_wind", "2t": "2m_temperature",
"2d": "2m_dewpoint_temperature", "tcwv": "total_column_water_vapour",
"tcc": "total_cloud_cover", "100u": "100m_u_component_of_wind",
"100v": "100m_v_component_of_wind", "sp": "surface_pressure",
"lcc": "low_cloud_cover", "mcc": "medium_cloud_cover",
"hcc": "high_cloud_cover", "skt": "skin_temperature",
"stl1": "soil_temperature_level_1", "swvl1": "volumetric_soil_water_layer_1",
"ci": "sea_ice_cover", "scaled_sd": "snow_depth",
}
_AURORA_TO_ARCO_ATMOS = {
"z": "geopotential", "q": "specific_humidity", "t": "temperature",
"u": "u_component_of_wind", "v": "v_component_of_wind",
}
_OUTPUT_ONLY_SURF_VARS = [
("i10fg", "i10fg", False), ("blh", "blh", False),
("uvb1h", "uvb_1h", False), ("ssrd1h", "ssrd_1h", False),
("ttr1h", "ttr_1h", False), ("tp1h", "scaled_tp_1h", True),
("sf1h", "scaled_sf_1h", True),
]
def ts_to_datetime(ts):
epoch = np.datetime64("1970-01-01T00:00:00")
seconds = float((ts.astype("datetime64[s]") - epoch) / np.timedelta64(1, "s"))
return datetime.fromtimestamp(seconds, tz=timezone.utc)
def arco_time_index(ts):
return int((ts - ARCO_EPOCH) / np.timedelta64(1, "h"))
model = AuroraEnsembleModel()
model.load_checkpoint_local(CKPT_PATH)
model.eval().to(DEVICE)
with open(STATIC_PATH, "rb") as f:
static_raw = pickle.load(f) # noqa: S301
static_vars = {k: torch.from_numpy(np.asarray(v))[:720, :] for k, v in static_raw.items()}
LAT = np.linspace(90, -90, 720, endpoint=False)
LON = np.linspace(0, 359.75, 1440)
fs = gcsfs.GCSFileSystem(token="anon", asynchronous=False)
zstore = zarr.storage.FsspecStore(fs, path=ARCO_PATH)
zg = zarr.open(store=zstore, mode="r")
all_levels = zg["level"][:]
level_indices = [int(np.searchsorted(all_levels, lv)) for lv in ATMOS_LEVELS]
ti_minus6 = arco_time_index(T_MINUS_6)
ti_0 = arco_time_index(T0)
atmos_data = {}
for aurora_name, arco_name in _AURORA_TO_ARCO_ATMOS.items():
arr = zg[arco_name]
chunk = np.stack(
[arr[ti_minus6, level_indices, :720, :], arr[ti_0, level_indices, :720, :]],
axis=0,
).astype(np.float32)
atmos_data[aurora_name] = torch.from_numpy(chunk).unsqueeze(0)
surf_data = {}
for aurora_name, arco_name in _AURORA_TO_ARCO_SURF.items():
arr = zg[arco_name]
chunk = np.stack(
[arr[ti_minus6, :720, :], arr[ti_0, :720, :]], axis=0,
).astype(np.float32)
surf_data[aurora_name] = torch.from_numpy(chunk).unsqueeze(0)
for k in surf_data:
surf_data[k] = torch.nan_to_num(surf_data[k], nan=0.0)
dt0 = ts_to_datetime(T_MINUS_6)
dt1 = ts_to_datetime(T0)
insol = aurora_insolation((dt0, dt1), LAT, LON, enforce_2d=True)
surf_data["insolation"] = torch.from_numpy(np.asarray(insol)).float().unsqueeze(0)
batch = Batch(
surf_vars=surf_data, static_vars=static_vars, atmos_vars=atmos_data,
metadata=Metadata(
lat=torch.from_numpy(LAT).float(), lon=torch.from_numpy(LON).float(),
time=(dt1,), atmos_levels=tuple(int(lv) for lv in ATMOS_LEVELS), rollout_step=0,
),
)
torch.manual_seed(SEED)
model.reset_noise()
sub_outputs = {}
with torch.inference_mode():
for h, pred in enumerate(
rollout(model, batch, steps=2, fine_lead_times=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]),
start=1,
):
atmos_out = torch.cat(
[pred.atmos_vars["z"], pred.atmos_vars["q"], pred.atmos_vars["t"],
pred.atmos_vars["u"], pred.atmos_vars["v"]], dim=2,
)
surf_out = torch.cat(
[pred.surf_vars[k].unsqueeze(2) for k in _AURORA_TO_ARCO_SURF], dim=2,
)
diag_tensors = []
for _, aurora_name, log_scaled in _OUTPUT_ONLY_SURF_VARS:
v = pred.surf_vars[aurora_name]
if log_scaled:
v = aurora_log_untransform(v)
diag_tensors.append(v.unsqueeze(2))
diag_out = torch.cat(diag_tensors, dim=2)
sub_outputs[h] = torch.cat([atmos_out, surf_out, diag_out], dim=2).cpu()
torch.save(sub_outputs, "reference_aurora_v1p5_ens_vanilla.pt")
print("Saved reference_aurora_v1p5_ens_vanilla.pt")Aurora1p5Ensemble — Earth2Studio reference script"""Earth2Studio Aurora1p5Ensemble wrapper inference.
Fetches the same ERA5 initial conditions from ARCO as the vanilla ensemble
script, runs 12 hourly steps (2 full AR cycles) using the E2S
Aurora1p5Ensemble wrapper via create_iterator with seed=42.
Saves output tensors to reference_aurora_v1p5_ens_e2s.pt.
Run with:
uv run python reference_aurora_v1p5_ens_e2s.py
"""
import numpy as np
import torch
from earth2studio.data import ARCO
from earth2studio.data.utils import fetch_data
from earth2studio.models.px.aurora1p5 import VARIABLES, Aurora1p5Ensemble
from earth2studio.utils.coords import map_coords
DEVICE = "cuda:0"
TIME = np.array([np.datetime64("2025-01-01T00:00")])
N_STEPS = 12
SEED = 42
package = Aurora1p5Ensemble.load_default_package()
model = Aurora1p5Ensemble.load_model(package).to(DEVICE)
model.seed = SEED
model.eval()
data = ARCO()
lead_time = model.input_coords()["lead_time"]
variable = np.array(VARIABLES)
x, coords = fetch_data(data, TIME, variable, lead_time, device=DEVICE)
x, coords = map_coords(x, coords, model.input_coords())
sub_outputs = {}
with torch.inference_mode():
p_iter = model.create_iterator(x, coords)
next(p_iter)
for h, (out, _) in enumerate(p_iter, start=1):
sub_outputs[h] = out.cpu().view(1, 1, 90, 720, 1440)
if h >= N_STEPS:
break
torch.save(sub_outputs, "reference_aurora_v1p5_ens_e2s.pt")
print("Saved reference_aurora_v1p5_ens_e2s.pt")Aurora1p5Ensemble — comparison script"""Compare vanilla Aurora v1.5 ensemble outputs against Earth2Studio wrapper.
Loads reference_aurora_v1p5_ens_vanilla.pt and reference_aurora_v1p5_ens_e2s.pt
and reports per-variable MaxAbsErr (t2m, u10m, z500) and overall correlation.
Generates a comparison plot (t2m at h+1, h+6, h+12; rows: E2S / Vanilla / |diff|).
Relative error denominator: max(abs(vanilla), 1e-6).
Both runs must use SEED=42.
Run with:
uv run python reference_aurora_v1p5_ens_compare.py
"""
import numpy as np
import torch
from earth2studio.models.px.aurora1p5 import OUTPUT_VARIABLES
vanilla = torch.load("reference_aurora_v1p5_ens_vanilla.pt", map_location="cpu", weights_only=False)
e2s = torch.load("reference_aurora_v1p5_ens_e2s.pt", map_location="cpu", weights_only=False)
N_STEPS = len(vanilla)
VARS = list(OUTPUT_VARIABLES)
T2M_IDX = VARS.index("t2m")
U10M_IDX = VARS.index("u10m")
Z500_IDX = VARS.index("z500")
header = f"{'Hour':>4} {'t2m MaxAbsErr':>14} {'u10m MaxAbsErr':>15} {'z500 MaxAbsErr':>15} {'Corr (all)':>11}"
print(header)
print("-" * len(header))
corrs = []
for h in range(1, N_STEPS + 1):
v = vanilla[h].float()
e = e2s[h].float()
diff = (v - e).abs()
t2m_err = diff[0, 0, T2M_IDX].max().item()
u10m_err = diff[0, 0, U10M_IDX].max().item()
z500_err = diff[0, 0, Z500_IDX].max().item()
corr = torch.corrcoef(torch.stack([v.flatten(), e.flatten()]))[0, 1].item()
corrs.append(corr)
print(f" h+{h:02d} {t2m_err:>14.4f} {u10m_err:>15.4f} {z500_err:>15.2f} {corr:>11.8f}")
print(f"\nMean correlation: {np.mean(corrs):.8f} Min: {np.min(corrs):.8f}")
try:
import matplotlib.pyplot as plt
steps = [1, 6, N_STEPS]
fig, axes = plt.subplots(3, len(steps), figsize=(15, 9))
fig.suptitle("Aurora1p5Ensemble: E2S vs Vanilla t2m (K) [seed=42]", fontsize=13)
vmin, vmax = 220, 320
for col, h in enumerate(steps):
v_t2m = vanilla[h][0, 0, T2M_IDX].numpy()
e_t2m = e2s[h][0, 0, T2M_IDX].numpy()
err = np.abs(v_t2m - e_t2m)
axes[0, col].imshow(e_t2m, vmin=vmin, vmax=vmax, origin="upper", aspect="auto")
axes[0, col].set_title(f"E2S h+{h}"); axes[0, col].axis("off")
axes[1, col].imshow(v_t2m, vmin=vmin, vmax=vmax, origin="upper", aspect="auto")
axes[1, col].set_title(f"Vanilla h+{h}"); axes[1, col].axis("off")
im = axes[2, col].imshow(err, origin="upper", aspect="auto", cmap="Reds")
axes[2, col].set_title(f"|E2S - Vanilla| h+{h}"); axes[2, col].axis("off")
plt.colorbar(im, ax=axes[2, col], fraction=0.03)
plt.tight_layout()
plt.savefig("reference_aurora_v1p5_ens_compare.png", dpi=120, bbox_inches="tight")
plt.close()
except ImportError:
print("matplotlib not available — skipping plot")Aurora1p5Ensemble — sanity check and reproducibility script"""Sanity-check and reproducibility test for Aurora1p5Ensemble.
Runs 12 hourly forecast steps with seed=42, then runs again with the same
seed and asserts bit-perfect agreement. Plots t2m, msl, z500 at h+6 and h+12.
Run with:
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True uv run python reference_aurora_v1p5_ensemble_sanity.py
"""
import numpy as np
import torch
from earth2studio.data import ARCO
from earth2studio.data.utils import fetch_data
from earth2studio.models.px.aurora1p5 import VARIABLES, Aurora1p5Ensemble
from earth2studio.utils.coords import map_coords
DEVICE = "cuda:0"
TIME = np.array([np.datetime64("2025-01-01T00:00")])
N_STEPS = 12
N_MEMBERS = 1
SEED = 42
PLOT_VARS = {
"t2m": {"idx": list(VARIABLES).index("t2m"), "vmin": 220, "vmax": 320, "label": "t2m (K)"},
"msl": {"idx": list(VARIABLES).index("msl"), "vmin": 97000, "vmax": 103000, "label": "msl (Pa)"},
"z500": {"idx": list(VARIABLES).index("z500"), "vmin": 49000, "vmax": 59000, "label": "z500 (m²/s²)"},
}
package = Aurora1p5Ensemble.load_default_package()
model = Aurora1p5Ensemble.load_model(package).to(DEVICE)
model.eval()
data = ARCO()
lead_time = model.input_coords()["lead_time"]
variable = np.array(VARIABLES)
x, coords = fetch_data(data, TIME, variable, lead_time, device=DEVICE)
x, coords = map_coords(x, coords, model.input_coords())
x_ens = x.unsqueeze(0).repeat(N_MEMBERS, 1, 1, 1, 1, 1)
coords_ens = coords.copy()
coords_ens["ensemble"] = np.arange(N_MEMBERS)
coords_ens.move_to_end("ensemble", last=False)
def run_iterator(seed):
model.seed = seed
outputs = {}
iterator = model.create_iterator(x_ens, coords_ens)
next(iterator)
with torch.inference_mode():
for out, out_coords in iterator:
h = int(out_coords["lead_time"][0] / np.timedelta64(1, "h"))
outputs[h] = out[:, 0, 0].cpu()
if h >= N_STEPS:
break
return outputs
print(f"Run 1 (seed={SEED})...")
outputs = run_iterator(SEED)
t2m_idx = list(VARIABLES).index("t2m")
for h in sorted(outputs):
t2m = outputs[h][:, t2m_idx]
print(f" h+{h:2d}: t2m [{t2m.min():.1f}, {t2m.max():.1f}] K"
f" NaN={outputs[h].isnan().sum().item()} Inf={outputs[h].isinf().sum().item()}")
print(f"\nRun 2 (seed={SEED}) — reproducibility check...")
outputs2 = run_iterator(SEED)
all_match = all(torch.equal(outputs[h], outputs2[h]) for h in sorted(outputs))
for h in sorted(outputs):
print(f" h+{h:2d}: {'MATCH' if torch.equal(outputs[h], outputs2[h]) else 'MISMATCH'}")
print("✓ Bit-perfect agreement." if all_match else "✗ Reproducibility FAILED.")
try:
import matplotlib.pyplot as plt
plot_steps = [6, 12]
fig, axes = plt.subplots(len(PLOT_VARS), len(plot_steps), figsize=(10, len(PLOT_VARS) * 3))
fig.suptitle(f"Aurora1p5Ensemble sanity — IC: 2025-01-01T00Z (seed={SEED})", fontsize=11)
for row, (var_name, meta) in enumerate(PLOT_VARS.items()):
for col, h in enumerate(plot_steps):
field = outputs[h][0, meta["idx"]].numpy()
im = axes[row, col].imshow(
field, vmin=meta["vmin"], vmax=meta["vmax"], origin="upper", aspect="auto"
)
axes[row, col].set_title(f"{meta['label']} h+{h}")
axes[row, col].axis("off")
plt.colorbar(im, ax=axes[row, col], fraction=0.03)
plt.tight_layout()
plt.savefig("reference_aurora_v1p5_ensemble_sanity.png", dpi=100, bbox_inches="tight")
plt.close()
except ImportError:
print("matplotlib not available — skipping plot")
total_nan = sum(t.isnan().sum().item() for t in outputs.values())
total_inf = sum(t.isinf().sum().item() for t in outputs.values())
print(f"\nTotal NaN: {total_nan} Total Inf: {total_inf}") |
- Rename AuroraV1p5/AuroraV1p5Ensemble → Aurora1p5/Aurora1p5Ensemble (matching FCN1/FCN3 naming convention) - Rename aurora-v1p5 extra → aurora; bump microsoft-aurora requirement to >=2.0.0 and drop the separate aurora extra (backwards-compatible) - Switch HuggingFace checkpoint source from ikwessel/aurora-1.5 to official microsoft/aurora repo at commit c171214 - Add Aurora1p5Ensemble.set_rng(seed) for reproducible noise seeding - Add product badges: precip, land, ocean, solar - Fix year badge to 2026 - Change lead_time_hours to list[int] in _forward_sub_steps - Consolidate _forward/_forward_sub_steps into a single method - Trim docstrings; fix stale wording - Run uv lock Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Add
AuroraV1p5andAuroraV1p5Ensembleprognostic model wrappers for the Microsoft Aurora v1.5 0.25-degree global weather model.Model details
ikwessel/aurora-1.5Dependencies added
microsoft-aurora>=2.0.0Reference comparison
Reference comparison validated against the official
aurora.rollout()path using ERA5 initial conditions (2025-01-01T00:00 UTC) fetched directly from the ARCO GCS zarr store, with deterministic settings (torch.inference_mode()).Cycle 1 (h+1 through h+6):
Cycle 2 (h+7 through h+12):
A key fix was discovered during validation:
_AR_CLIP_BOUNDSwas missing three variables (swvl1,sic,sd) that Aurora'sapply_rollout_input_clippingapplies at the cycle boundary. Once added, both AR cycles achieve full bit-level agreement.Validation
See the reference-comparison validation comment below. Plot placeholders are
included for manual image upload by the PR author.
Checklist