Skip to content

StormCast-CONUS Score-Based Data Assimilation

Running StormCast-CONUS with guided diffusion posterior sampling to assimilate observations.

This example demonstrates how to use the StormCast-CONUS generative model with score-based data assimilation (SDA) over the Continental United States. Sparse in-situ surface observations from NOAA Global Historical Climatology Network Hourly (GHCNh) are assimilated at each forecast step using diffusion posterior sampling (DPS) guidance. Two forecasts from the central United States are run to illustrate the impact of data assimilation: one without observations and one with GHCNh surface station data.

In this example you will learn:

  • How to load StormCast-CONUS and configure SDA parameters
  • Fetching HRRR initial conditions and GHCNh surface observations
  • Running the model iteratively with and without observation assimilation
  • Comparing assimilated and non-assimilated forecasts

Set Up

This example requires the following components:

StormCast-CONUS extends the StormCast generative architecture to the full CONUS HRRR domain at 3 km resolution. When an observation DataFrame is passed to each generator step, diffusion posterior sampling (DPS) steers the denoising trajectory toward the observed values.

import os

os.makedirs("outputs", exist_ok=True)
from dotenv import load_dotenv

load_dotenv()  # TODO: make common example prep function

from datetime import timedelta

import numpy as np
import torch
from loguru import logger
from tqdm import tqdm

logger.remove()
logger.add(lambda msg: tqdm.write(msg, end=""), colorize=True)

from earth2studio.data import HRRR, GHCNHourly, fetch_data
from earth2studio.models.px import StormCastCONUS
from earth2studio.utils.coords import map_coords

# Load the default model package
package = StormCastCONUS.load_default_package()

# Configure SDA: sda_std_obs is the assumed normalised observation noise std per
# variable (lower = trust observations more).
# sda_gamma is the DPS step-size scaling (lower = stronger guidance from obs).
#
# By default the example runs on the central-US subdomain to reduce GPU memory
# requirements. For the full CONUS domain, comment out hrrr_lat_lim and
# hrrr_lon_lim below (requires ~200 GB VRAM for SDA).
hrrr_lat_lim = (273, 785)
hrrr_lon_lim = (579, 1219)
model = StormCastCONUS.load_model(
    package,
    hrrr_lat_lim=hrrr_lat_lim,  # comment out for full CONUS domain
    hrrr_lon_lim=hrrr_lon_lim,  # comment out for full CONUS domain
    num_diffusion_steps=18,
    num_sda_diffusion_steps=96,
    sda_std_obs=0.15,
    sda_gamma=1e-3,
)
model = model.to("cuda:0")

hrrr = HRRR()
Console output9 lines
/__w/earth2studio/earth2studio/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:64: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:
    PyTorch 2.10.0+cu128 with CUDA 1208 (you have 2.13.0+cu130)
    Python  3.10.19 (you have 3.13.13)
  Please reinstall xformers (see https://github.com/facebookresearch/xformers#installing-xformers)
  Memory-efficient attention, SwiGLU, sparse and more won't be available.
  Set XFORMERS_MORE_DETAILS=1 for more details
CuPy distance computation test failed with error: cuVS >= 24.12 or pylibraft < 24.12 should be installed to use this feature

Fetch Initial Conditions

Pull HRRR analysis data for 17 April 2026, a date that saw a significant tornado outbreak across the central United States, and align it to the model's coordinate system.

init_time = np.array([np.datetime64("2026-04-17T18:00")])

x, coords = fetch_data(
    hrrr,
    time=init_time,
    variable=model.variables,
    lead_time=np.array([np.timedelta64(0, "h")]),
    device="cuda:0",
)
x, coords = map_coords(x, coords, model.input_coords())
Console output19 lines
Fetching HRRR data:   0%|          | 0/99 [00:00<?, ?it/s]
Fetching HRRR data:   1%|          | 1/99 [00:00<00:46,  2.10it/s]
Fetching HRRR data:   2%|โ–         | 2/99 [00:00<00:25,  3.77it/s]
Fetching HRRR data:   4%|โ–         | 4/99 [00:00<00:12,  7.38it/s]
Fetching HRRR data:   8%|โ–Š         | 8/99 [00:00<00:06, 14.10it/s]
Fetching HRRR data:  12%|โ–ˆโ–        | 12/99 [00:00<00:04, 18.82it/s]
Fetching HRRR data:  15%|โ–ˆโ–Œ        | 15/99 [00:01<00:04, 19.08it/s]
Fetching HRRR data:  18%|โ–ˆโ–Š        | 18/99 [00:01<00:04, 17.05it/s]
Fetching HRRR data:  20%|โ–ˆโ–ˆ        | 20/99 [00:01<00:05, 15.34it/s]
Fetching HRRR data:  22%|โ–ˆโ–ˆโ–       | 22/99 [00:01<00:05, 14.81it/s]
Fetching HRRR data:  24%|โ–ˆโ–ˆโ–       | 24/99 [00:01<00:06, 12.46it/s]
Fetching HRRR data:  27%|โ–ˆโ–ˆโ–‹       | 27/99 [00:02<00:05, 12.87it/s]
Fetching HRRR data:  30%|โ–ˆโ–ˆโ–ˆ       | 30/99 [00:02<00:05, 12.28it/s]
Fetching HRRR data:  37%|โ–ˆโ–ˆโ–ˆโ–‹      | 37/99 [00:02<00:04, 14.37it/s]
Fetching HRRR data:  47%|โ–ˆโ–ˆโ–ˆโ–ˆโ–‹     | 47/99 [00:03<00:03, 16.34it/s]
Fetching HRRR data:  57%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‹    | 56/99 [00:03<00:01, 22.10it/s]
Fetching HRRR data:  70%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‰   | 69/99 [00:03<00:01, 28.53it/s]
Fetching HRRR data:  82%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ– | 81/99 [00:04<00:00, 34.77it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 99/99 [00:04<00:00, 24.07it/s]

Run Without Observations

Step the model forward 6 hours without any observations. This uses the standard EDM diffusion sampler, equivalent to running StormCast-CONUS as a pure generative forecast model. We store only the 10-m zonal wind (u10m) used for comparison.

nsteps = 6
plot_var = "u10m"
plot_vars = ["u10m", "v10m", "t2m"]
var_idx = list(model.variables).index(plot_var)

np.random.seed(42)
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)

no_obs_fields = []
gen = model.create_generator(x.clone(), coords.copy())
x_cur, c_cur = next(gen)  # prime the generator, yields initial state (lead_time = 0 h)

for step in tqdm(range(nsteps), desc="No-obs forecast"):
    logger.info(f"Running no-obs forecast step {step + 1}/{nsteps}")
    x_cur, c_cur = gen.send(None)  # advance one hour without observations
    no_obs_fields.append(x_cur[0, 0, var_idx].cpu().numpy())

gen.close()
no_obs_fields = np.stack(no_obs_fields)  # (nsteps, H, W)
Console output106 lines
No-obs forecast:   0%|          | 0/6 [00:00<?, ?it/s]

2026-08-15 05:39:35.671 | INFO     | __main__:<module>:16 - Running no-obs forecast step 1/6

No-obs forecast:   0%|          | 0/6 [00:00<?, ?it/s]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:10,  2.31it/s]

Fetching GFS data:   8%|โ–Š         | 2/26 [00:00<00:06,  3.51it/s]

Fetching GFS data:  19%|โ–ˆโ–‰        | 5/26 [00:00<00:02,  8.69it/s]

Fetching GFS data:  35%|โ–ˆโ–ˆโ–ˆโ–      | 9/26 [00:00<00:01, 14.70it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 26.32it/s]

No-obs forecast:  17%|โ–ˆโ–‹        | 1/6 [00:05<00:29,  5.88s/it]

2026-08-15 05:39:41.555 | INFO     | __main__:<module>:16 - Running no-obs forecast step 2/6

No-obs forecast:  17%|โ–ˆโ–‹        | 1/6 [00:05<00:29,  5.88s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:05,  4.30it/s]

Fetching GFS data:  19%|โ–ˆโ–‰        | 5/26 [00:00<00:02,  7.74it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 34.76it/s]

No-obs forecast:  33%|โ–ˆโ–ˆโ–ˆโ–Ž      | 2/6 [00:11<00:21,  5.47s/it]

2026-08-15 05:39:46.741 | INFO     | __main__:<module>:16 - Running no-obs forecast step 3/6

No-obs forecast:  33%|โ–ˆโ–ˆโ–ˆโ–Ž      | 2/6 [00:11<00:21,  5.47s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:06,  3.79it/s]

Fetching GFS data:  12%|โ–ˆโ–        | 3/26 [00:00<00:02,  9.08it/s]

Fetching GFS data:  19%|โ–ˆโ–‰        | 5/26 [00:00<00:01, 12.00it/s]

Fetching GFS data:  35%|โ–ˆโ–ˆโ–ˆโ–      | 9/26 [00:00<00:00, 19.08it/s]

Fetching GFS data:  92%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–| 24/26 [00:00<00:00, 55.45it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:01<00:00, 21.04it/s]

No-obs forecast:  50%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ     | 3/6 [00:16<00:16,  5.55s/it]

2026-08-15 05:39:52.389 | INFO     | __main__:<module>:16 - Running no-obs forecast step 4/6

No-obs forecast:  50%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ     | 3/6 [00:16<00:16,  5.55s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:04,  5.20it/s]

Fetching GFS data:  15%|โ–ˆโ–Œ        | 4/26 [00:00<00:01, 15.28it/s]

Fetching GFS data:  27%|โ–ˆโ–ˆโ–‹       | 7/26 [00:00<00:01, 13.86it/s]

Fetching GFS data:  42%|โ–ˆโ–ˆโ–ˆโ–ˆโ–     | 11/26 [00:00<00:00, 19.33it/s]

Fetching GFS data:  92%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–| 24/26 [00:00<00:00, 47.10it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 33.83it/s]

No-obs forecast:  67%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‹   | 4/6 [00:21<00:10,  5.43s/it]

2026-08-15 05:39:57.626 | INFO     | __main__:<module>:16 - Running no-obs forecast step 5/6

No-obs forecast:  67%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‹   | 4/6 [00:21<00:10,  5.43s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:05,  4.66it/s]

Fetching GFS data:  19%|โ–ˆโ–‰        | 5/26 [00:00<00:01, 18.43it/s]

Fetching GFS data:  31%|โ–ˆโ–ˆโ–ˆ       | 8/26 [00:00<00:00, 21.29it/s]

Fetching GFS data:  77%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‹  | 20/26 [00:00<00:00, 53.29it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:01<00:00, 18.59it/s]

No-obs forecast:  83%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–Ž | 5/6 [00:27<00:05,  5.59s/it]

2026-08-15 05:40:03.502 | INFO     | __main__:<module>:16 - Running no-obs forecast step 6/6

No-obs forecast:  83%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–Ž | 5/6 [00:27<00:05,  5.59s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:06,  3.90it/s]

Fetching GFS data:  15%|โ–ˆโ–Œ        | 4/26 [00:00<00:01, 12.17it/s]

Fetching GFS data:  23%|โ–ˆโ–ˆโ–Ž       | 6/26 [00:00<00:01, 13.89it/s]

Fetching GFS data:  69%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‰   | 18/26 [00:00<00:00, 44.20it/s]

Fetching GFS data:  92%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–| 24/26 [00:01<00:00, 17.34it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:01<00:00, 16.62it/s]

No-obs forecast: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 6/6 [00:33<00:00,  5.74s/it]
No-obs forecast: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 6/6 [00:33<00:00,  5.64s/it]

Fetch Observations and Plot Station Locations

Fetch GHCNh surface observations covering the model domain. The bounding box is derived from the model's lat/lon grid so it automatically adjusts when running on a subdomain. We visualise the station network before running the assimilation forecast.

# Model lon is in [0, 360); convert to [-180, 180) for the western-hemisphere
# CONUS domain. GHCNHourly.get_stations_bbox accepts either convention.
lat_min = float(model.lat.min())
lat_max = float(model.lat.max())
lon_min = float(model.lon.min()) - 360.0
lon_max = float(model.lon.max()) - 360.0
stations = GHCNHourly.get_stations_bbox((lat_min, lon_min, lat_max, lon_max))
ghcn = GHCNHourly(
    stations=stations, time_tolerance=timedelta(minutes=15), verbose=False
)

import cartopy
import cartopy.crs as ccrs
import matplotlib.pyplot as plt

# Fetch a sample observation at the initial time to retrieve station positions.
# GHCNh returns lon in [0, 360); shift to [-180, 180) for PlateCarree plots.
sample_df = ghcn(init_time, plot_vars)
station_lats = sample_df["lat"].values
station_lons = sample_df["lon"].values - 360.0

plt.close("all")
fig, ax = plt.subplots(subplot_kw={"projection": ccrs.PlateCarree()}, figsize=(8, 6))
ax.set_extent(
    [lon_min - 2, lon_max + 2, lat_min - 2, lat_max + 2], crs=ccrs.PlateCarree()
)
ax.add_feature(
    cartopy.feature.STATES.with_scale("50m"), linewidth=0.5, edgecolor="black"
)
ax.add_feature(cartopy.feature.LAND, facecolor="lightyellow")
ax.gridlines(draw_labels=True, linewidth=0.3, alpha=0.5)
ax.scatter(
    station_lons,
    station_lats,
    s=20,
    marker="x",
    transform=ccrs.PlateCarree(),
    zorder=3,
)
ax.set_title("GHCNh Station Locations - StormCast-CONUS Domain")
plt.savefig("outputs/01_ghcn_stations.jpg", dpi=150, bbox_inches="tight")

Output from StormCast-CONUS Score-Based Data Assimilation

Run Inference With Streaming Observations

At each forecast step the next valid time is determined from the current generator state, GHCNh observations are fetched for that time, and the observation DataFrame is sent to the generator. Zero-value wind reports (likely anemometer failure) are filtered out before assimilation.

np.random.seed(42)
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)

obs_fields = []
gen = model.create_generator(x.clone(), coords.copy())
x_cur, c_cur = next(gen)  # prime the generator, yields initial state

for step in tqdm(range(nsteps), desc="Obs forecast"):
    # Target valid time is one step ahead of the current generator state
    valid_time = np.array(
        [c_cur["time"][0] + c_cur["lead_time"][0] + np.timedelta64(1, "h")]
    )
    obs_df = ghcn(valid_time, plot_vars)

    # Drop zero-value wind reports (common anemometer failure mode)
    if len(obs_df) > 0:
        is_wind = obs_df["variable"].str.startswith(("u", "v"))
        zero_wind = is_wind & (obs_df["observation"].abs() < 1e-5)
        obs_df = obs_df[~zero_wind]

    obs = obs_df if len(obs_df) > 0 else None
    logger.info(
        f"Step {step + 1}/{nsteps} โ€” {len(obs_df) if obs is not None else 0} obs"
    )
    x_cur, c_cur = gen.send(obs)  # advance one hour with observations
    obs_fields.append(x_cur[0, 0, var_idx].cpu().numpy())

gen.close()
obs_fields = np.stack(obs_fields)  # (nsteps, H, W)
Console output68 lines
Obs forecast:   0%|          | 0/6 [00:00<?, ?it/s]

2026-08-15 05:41:48.751 | INFO     | __main__:<module>:24 - Step 1/6 โ€” 392 obs

Obs forecast:   0%|          | 0/6 [00:14<?, ?it/s]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:11,  2.15it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 54.94it/s]

Obs forecast:  17%|โ–ˆโ–‹        | 1/6 [01:16<06:20, 76.02s/it]

2026-08-15 05:43:04.187 | INFO     | __main__:<module>:24 - Step 2/6 โ€” 392 obs

Obs forecast:  17%|โ–ˆโ–‹        | 1/6 [01:30<06:20, 76.02s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:10,  2.28it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 56.90it/s]

Obs forecast:  33%|โ–ˆโ–ˆโ–ˆโ–Ž      | 2/6 [02:38<05:18, 79.65s/it]

2026-08-15 05:44:26.533 | INFO     | __main__:<module>:24 - Step 3/6 โ€” 392 obs

Obs forecast:  33%|โ–ˆโ–ˆโ–ˆโ–Ž      | 2/6 [02:52<05:18, 79.65s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:11,  2.23it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 57.82it/s]

Obs forecast:  50%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ     | 3/6 [04:09<04:15, 85.02s/it]

2026-08-15 05:45:58.076 | INFO     | __main__:<module>:24 - Step 4/6 โ€” 392 obs

Obs forecast:  50%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ     | 3/6 [04:23<04:15, 85.02s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:12,  2.01it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 51.66it/s]

Obs forecast:  67%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‹   | 4/6 [05:35<02:50, 85.49s/it]

2026-08-15 05:47:23.786 | INFO     | __main__:<module>:24 - Step 5/6 โ€” 392 obs

Obs forecast:  67%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–‹   | 4/6 [05:49<02:50, 85.49s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:12,  1.93it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 49.90it/s]

Obs forecast:  83%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–Ž | 5/6 [06:59<01:24, 84.84s/it]

2026-08-15 05:48:48.630 | INFO     | __main__:<module>:24 - Step 6/6 โ€” 392 obs

Obs forecast:  83%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–Ž | 5/6 [07:14<01:24, 84.84s/it]

Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]

Fetching GFS data:   4%|โ–         | 1/26 [00:00<00:10,  2.33it/s]
Fetching GFS data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 26/26 [00:00<00:00, 58.18it/s]

Obs forecast: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 6/6 [08:32<00:00, 87.52s/it]
Obs forecast: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 6/6 [08:32<00:00, 85.38s/it]

Post Processing

Compare the two forecasts. The top row shows the baseline (no observations), the middle row shows the assimilation forecast with station locations overlaid, and the bottom row shows the signed difference (assimilated โˆ’ baseline).

plt.close("all")

# HRRR Lambert Conformal projection
projection = ccrs.LambertConformal(
    central_longitude=262.5,
    central_latitude=38.5,
    standard_parallels=(38.5, 38.5),
    globe=ccrs.Globe(semimajor_axis=6371229, semiminor_axis=6371229),
)

fig, axes = plt.subplots(
    3,
    nsteps,
    subplot_kw={"projection": projection},
    figsize=(4 * nsteps, 8),
)
fig.subplots_adjust(wspace=0.02, hspace=0.08, left=0.1, right=0.9)

vmin, vmax = -10, 10
for step in range(nsteps):
    lead_hr = step + 1
    no_obs_field = no_obs_fields[step]
    obs_field = obs_fields[step]
    diff_field = obs_field - no_obs_field

    # Row 0: No-obs forecast
    ax = axes[0, step]
    im0 = ax.pcolormesh(
        model.lon,
        model.lat,
        no_obs_field,
        transform=ccrs.PlateCarree(),
        cmap="PRGn",
        vmin=vmin,
        vmax=vmax,
    )
    ax.add_feature(
        cartopy.feature.STATES.with_scale("50m"),
        linewidth=0.5,
        edgecolor="black",
        zorder=2,
    )
    ax.set_title(f"+{lead_hr}h")

    # Row 1: With-obs forecast + station locations
    ax = axes[1, step]
    im1 = ax.pcolormesh(
        model.lon,
        model.lat,
        obs_field,
        transform=ccrs.PlateCarree(),
        cmap="PRGn",
        vmin=vmin,
        vmax=vmax,
    )
    ax.scatter(
        station_lons,
        station_lats,
        s=8,
        facecolors="none",
        edgecolors="black",
        linewidths=0.8,
        transform=ccrs.PlateCarree(),
        zorder=3,
    )
    ax.add_feature(
        cartopy.feature.STATES.with_scale("50m"),
        linewidth=0.5,
        edgecolor="black",
        zorder=2,
    )

    # Row 2: Difference (assimilated โˆ’ baseline)
    ax = axes[2, step]
    im2 = ax.pcolormesh(
        model.lon,
        model.lat,
        diff_field,
        transform=ccrs.PlateCarree(),
        cmap="RdBu_r",
        vmin=-3,
        vmax=3,
    )
    ax.add_feature(
        cartopy.feature.STATES.with_scale("50m"),
        linewidth=0.5,
        edgecolor="black",
        zorder=2,
    )

for row_label, ax_row in zip(["No Obs", "Obs", "Difference"], axes[:, 0]):
    ax_row.text(
        -0.07,
        0.5,
        row_label,
        va="bottom",
        ha="center",
        rotation="vertical",
        rotation_mode="anchor",
        fontsize=12,
        transform=ax_row.transAxes,
    )

plt.colorbar(im0, ax=axes[0, -1], shrink=0.6, label=f"{plot_var} (m/s)")
plt.colorbar(im1, ax=axes[1, -1], shrink=0.6, label=f"{plot_var} (m/s)")
plt.colorbar(im2, ax=axes[2, -1], shrink=0.6, label=f"{plot_var} (m/s)")

plt.tight_layout()
plt.savefig("outputs/01_stormcast_conus_sda_comparison.jpg", dpi=150)

Output from StormCast-CONUS Score-Based Data Assimilation

Ground Truth Comparison

Fetch HRRR analysis at each valid forecast time and compute the absolute error of both the no-obs and assimilation forecasts. This shows whether assimilation improves accuracy relative to the actual analysis.

truth_x, truth_coords = fetch_data(
    hrrr,
    time=init_time,
    variable=np.array([plot_var]),
    lead_time=np.array([np.timedelta64(h + 1, "h") for h in range(nsteps)]),
    device="cpu",
)
truth_x, truth_coords = map_coords(
    truth_x,
    truth_coords,
    {"hrrr_y": model.hrrr_y, "hrrr_x": model.hrrr_x},
)
# truth_x shape: (time=1, lead_time=nsteps, variable=1, H, W)
truth_fields = truth_x[0, :, 0].numpy()  # (nsteps, H, W)

no_obs_err = np.abs(no_obs_fields - truth_fields)
obs_err = np.abs(obs_fields - truth_fields)
Console output23 lines
Fetching HRRR data:   0%|          | 0/1 [00:00<?, ?it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  2.10it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  2.09it/s]

Fetching HRRR data:   0%|          | 0/1 [00:00<?, ?it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  2.11it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  2.10it/s]

Fetching HRRR data:   0%|          | 0/1 [00:00<?, ?it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  4.38it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  4.37it/s]

Fetching HRRR data:   0%|          | 0/1 [00:00<?, ?it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  5.62it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  5.61it/s]

Fetching HRRR data:   0%|          | 0/1 [00:00<?, ?it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  3.60it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  3.60it/s]

Fetching HRRR data:   0%|          | 0/1 [00:00<?, ?it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  5.54it/s]
Fetching HRRR data: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00,  5.53it/s]

Plot absolute errors between the StormCast-CONUS predictions and HRRR analysis ground truth. In later time-steps the assimilated forecast typically shows lower errors near the observation stations.

plt.close("all")
fig, axes = plt.subplots(
    2,
    nsteps,
    subplot_kw={"projection": projection},
    figsize=(4 * nsteps, 6),
)
fig.subplots_adjust(wspace=0.02, hspace=0.08, left=0.1, right=0.9)

err_max = 5
for step in range(nsteps):
    lead_hr = step + 1

    # Row 0: No-obs absolute error
    ax = axes[0, step]
    im0 = ax.pcolormesh(
        model.lon,
        model.lat,
        no_obs_err[step],
        transform=ccrs.PlateCarree(),
        cmap="viridis",
        vmin=0,
        vmax=err_max,
    )
    ax.add_feature(
        cartopy.feature.STATES.with_scale("50m"),
        linewidth=0.5,
        edgecolor="grey",
        zorder=2,
    )
    ax.set_title(f"+{lead_hr}h")

    # Row 1: Obs absolute error
    ax = axes[1, step]
    im1 = ax.pcolormesh(
        model.lon,
        model.lat,
        obs_err[step],
        transform=ccrs.PlateCarree(),
        cmap="viridis",
        vmin=0,
        vmax=err_max,
    )
    ax.add_feature(
        cartopy.feature.STATES.with_scale("50m"),
        linewidth=0.5,
        edgecolor="grey",
        zorder=2,
    )

for row_label, ax_row in zip(["|No Obs โˆ’ Truth|", "|Obs โˆ’ Truth|"], axes[:, 0]):
    ax_row.text(
        -0.07,
        0.5,
        row_label,
        va="bottom",
        ha="center",
        rotation="vertical",
        rotation_mode="anchor",
        fontsize=12,
        transform=ax_row.transAxes,
    )

plt.colorbar(im0, ax=axes[0, -1], shrink=0.6, label=f"|ฮ”{plot_var}| (m/s)")
plt.colorbar(im1, ax=axes[1, -1], shrink=0.6, label=f"|ฮ”{plot_var}| (m/s)")
plt.tight_layout()
plt.savefig("outputs/01_stormcast_conus_sda_gt_comparison.jpg", dpi=150)

Output from StormCast-CONUS Score-Based Data Assimilation

Download Python sourceDownload Jupyter notebookexecuted in 712.8s
Execution profile

Runtime telemetry

Total runtime11m 53s

Execution environment

CPUAMD EPYC 7313P 16-Core Processor
GPUNVIDIA H100 PCIe ยท 79.6 GiB
System RAM58.5 GiB
PlatformLinux 6.8.0-136-generic
Python3.13.13
GPU driver / CUDADriver 595.84 ยท CUDA support 13.2