HealDA Global Data Assimilation¶

Producing a global weather analysis from satellite and in-situ observations.

This example demonstrates how to use the HealDA data assimilation model to produce a global weather analysis on a HEALPix grid from sparse in-situ (conventional) and satellite radiance observations sourced from the NOAA UFS replay archive. Three runs are compared: conventional observations only, satellite observations only, and both combined to illustrate the impact of each observation type.

In this example you will learn:

How to load and initialise the HealDA data assimilation model
Fetching UFS conventional and satellite observation DataFrames
Running the model with different observation combinations
Comparing the assimilated global fields against ERA5 data

Set Up¶

This example requires the following components:

Assimilation Model: HealDA earth2studio.models.da.HealDA.
Datasource (conv): UFS conventional observations earth2studio.data.UFSObsConv.
Datasource (sat): UFS satellite observations earth2studio.data.UFSObsSat.

HealDA is a stateless neural-network-based data assimilation model that ingests conventional (radiosonde, surface station, GPS-RO, etc.) and satellite radiance observations and produces a single global weather analysis on a HEALPix level-6 grid.

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 NCAR_ERA5, UFSObsConv, UFSObsSat, fetch_dataframe
from earth2studio.models.da import HealDA

# Load the default model package (downloads checkpoint from HuggingFace)
# Setting lat_lon=True regrids the native HEALPix output to a regular lat-lon grid.
package = HealDA.load_default_package()
model = HealDA.load_model(package, lat_lon=True)
model = model.to("cuda:0")

Console output119 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]

Downloading config.json: 0%|          | 0.00/25.0 [00:00<?, ?B/s]
Downloading config.json: 100%|██████████| 25.0/25.0 [00:00<00:00, 204B/s]
Downloading config.json: 100%|██████████| 25.0/25.0 [00:00<00:00, 201B/s]

Downloading healda_ufs_era5.mdlus: 0%|          | 0.00/1.22G [00:00<?, ?B/s]
Downloading healda_ufs_era5.mdlus: 1%|          | 10.0M/1.22G [00:01<02:37, 8.26MB/s]
Downloading healda_ufs_era5.mdlus: 2%|▏         | 30.0M/1.22G [00:01<00:48, 26.1MB/s]
Downloading healda_ufs_era5.mdlus: 3%|▎         | 40.0M/1.22G [00:01<00:41, 30.7MB/s]
Downloading healda_ufs_era5.mdlus: 4%|▍         | 50.0M/1.22G [00:01<00:31, 39.7MB/s]
Downloading healda_ufs_era5.mdlus: 5%|▍         | 60.0M/1.22G [00:01<00:26, 46.9MB/s]
Downloading healda_ufs_era5.mdlus: 6%|▌         | 70.0M/1.22G [00:02<00:22, 54.5MB/s]
Downloading healda_ufs_era5.mdlus: 6%|▋         | 80.0M/1.22G [00:02<00:19, 63.5MB/s]
Downloading healda_ufs_era5.mdlus: 8%|▊         | 100M/1.22G [00:02<00:17, 68.6MB/s] 
Downloading healda_ufs_era5.mdlus: 10%|▉         | 120M/1.22G [00:02<00:14, 79.5MB/s]
Downloading healda_ufs_era5.mdlus: 10%|█         | 130M/1.22G [00:02<00:14, 78.4MB/s]
Downloading healda_ufs_era5.mdlus: 11%|█         | 140M/1.22G [00:02<00:14, 82.0MB/s]
Downloading healda_ufs_era5.mdlus: 13%|█▎        | 160M/1.22G [00:03<00:12, 90.7MB/s]
Downloading healda_ufs_era5.mdlus: 14%|█▎        | 170M/1.22G [00:03<00:14, 78.8MB/s]
Downloading healda_ufs_era5.mdlus: 14%|█▍        | 180M/1.22G [00:03<00:13, 83.0MB/s]
Downloading healda_ufs_era5.mdlus: 16%|█▌        | 200M/1.22G [00:03<00:14, 76.9MB/s]
Downloading healda_ufs_era5.mdlus: 17%|█▋        | 210M/1.22G [00:03<00:17, 62.6MB/s]
Downloading healda_ufs_era5.mdlus: 18%|█▊        | 220M/1.22G [00:04<00:20, 51.9MB/s]
Downloading healda_ufs_era5.mdlus: 18%|█▊        | 230M/1.22G [00:04<00:22, 48.1MB/s]
Downloading healda_ufs_era5.mdlus: 20%|██        | 250M/1.22G [00:04<00:16, 63.1MB/s]
Downloading healda_ufs_era5.mdlus: 21%|██        | 260M/1.22G [00:04<00:17, 58.0MB/s]
Downloading healda_ufs_era5.mdlus: 22%|██▏       | 270M/1.22G [00:05<00:15, 65.1MB/s]
Downloading healda_ufs_era5.mdlus: 22%|██▏       | 280M/1.22G [00:05<00:15, 64.9MB/s]
Downloading healda_ufs_era5.mdlus: 24%|██▍       | 300M/1.22G [00:05<00:13, 72.6MB/s]
Downloading healda_ufs_era5.mdlus: 25%|██▍       | 310M/1.22G [00:05<00:13, 70.7MB/s]
Downloading healda_ufs_era5.mdlus: 26%|██▌       | 320M/1.22G [00:05<00:15, 62.0MB/s]
Downloading healda_ufs_era5.mdlus: 27%|██▋       | 340M/1.22G [00:06<00:12, 75.8MB/s]
Downloading healda_ufs_era5.mdlus: 28%|██▊       | 350M/1.22G [00:06<00:11, 80.0MB/s]
Downloading healda_ufs_era5.mdlus: 30%|██▉       | 370M/1.22G [00:06<00:10, 88.6MB/s]
Downloading healda_ufs_era5.mdlus: 30%|███       | 380M/1.22G [00:06<00:10, 91.0MB/s]
Downloading healda_ufs_era5.mdlus: 31%|███       | 390M/1.22G [00:06<00:11, 77.7MB/s]
Downloading healda_ufs_era5.mdlus: 33%|███▎      | 410M/1.22G [00:07<00:12, 68.0MB/s]
Downloading healda_ufs_era5.mdlus: 34%|███▎      | 420M/1.22G [00:07<00:12, 68.6MB/s]
Downloading healda_ufs_era5.mdlus: 34%|███▍      | 430M/1.22G [00:07<00:11, 73.7MB/s]
Downloading healda_ufs_era5.mdlus: 35%|███▌      | 440M/1.22G [00:07<00:11, 74.2MB/s]
Downloading healda_ufs_era5.mdlus: 36%|███▌      | 450M/1.22G [00:07<00:16, 52.3MB/s]
Downloading healda_ufs_era5.mdlus: 38%|███▊      | 470M/1.22G [00:08<00:12, 63.5MB/s]
Downloading healda_ufs_era5.mdlus: 39%|███▉      | 490M/1.22G [00:08<00:10, 73.5MB/s]
Downloading healda_ufs_era5.mdlus: 41%|████      | 510M/1.22G [00:08<00:10, 72.0MB/s]
Downloading healda_ufs_era5.mdlus: 42%|████▏     | 520M/1.22G [00:08<00:11, 68.9MB/s]
Downloading healda_ufs_era5.mdlus: 42%|████▏     | 530M/1.22G [00:08<00:12, 60.1MB/s]
Downloading healda_ufs_era5.mdlus: 44%|████▍     | 550M/1.22G [00:09<00:11, 63.1MB/s]
Downloading healda_ufs_era5.mdlus: 46%|████▌     | 570M/1.22G [00:09<00:11, 60.0MB/s]
Downloading healda_ufs_era5.mdlus: 46%|████▋     | 580M/1.22G [00:09<00:13, 53.9MB/s]
Downloading healda_ufs_era5.mdlus: 47%|████▋     | 590M/1.22G [00:10<00:11, 60.3MB/s]
Downloading healda_ufs_era5.mdlus: 49%|████▉     | 610M/1.22G [00:10<00:09, 72.9MB/s]
Downloading healda_ufs_era5.mdlus: 50%|█████     | 630M/1.22G [00:10<00:07, 82.4MB/s]
Downloading healda_ufs_era5.mdlus: 52%|█████▏    | 650M/1.22G [00:10<00:07, 89.6MB/s]
Downloading healda_ufs_era5.mdlus: 54%|█████▎    | 670M/1.22G [00:10<00:06, 94.5MB/s]
Downloading healda_ufs_era5.mdlus: 55%|█████▌    | 690M/1.22G [00:10<00:05, 100MB/s] 
Downloading healda_ufs_era5.mdlus: 57%|█████▋    | 710M/1.22G [00:11<00:05, 101MB/s]
Downloading healda_ufs_era5.mdlus: 58%|█████▊    | 730M/1.22G [00:11<00:05, 103MB/s]
Downloading healda_ufs_era5.mdlus: 60%|██████    | 750M/1.22G [00:11<00:05, 104MB/s]
Downloading healda_ufs_era5.mdlus: 62%|██████▏   | 770M/1.22G [00:11<00:04, 107MB/s]
Downloading healda_ufs_era5.mdlus: 63%|██████▎   | 790M/1.22G [00:11<00:04, 104MB/s]
Downloading healda_ufs_era5.mdlus: 65%|██████▍   | 810M/1.22G [00:12<00:04, 106MB/s]
Downloading healda_ufs_era5.mdlus: 66%|██████▋   | 830M/1.22G [00:12<00:04, 107MB/s]
Downloading healda_ufs_era5.mdlus: 68%|██████▊   | 850M/1.22G [00:12<00:03, 108MB/s]
Downloading healda_ufs_era5.mdlus: 70%|██████▉   | 870M/1.22G [00:12<00:03, 108MB/s]
Downloading healda_ufs_era5.mdlus: 71%|███████▏  | 890M/1.22G [00:12<00:03, 109MB/s]
Downloading healda_ufs_era5.mdlus: 73%|███████▎  | 910M/1.22G [00:13<00:03, 103MB/s]
Downloading healda_ufs_era5.mdlus: 74%|███████▎  | 920M/1.22G [00:13<00:04, 85.0MB/s]
Downloading healda_ufs_era5.mdlus: 75%|███████▌  | 940M/1.22G [00:13<00:03, 90.8MB/s]
Downloading healda_ufs_era5.mdlus: 77%|███████▋  | 960M/1.22G [00:13<00:03, 88.9MB/s]
Downloading healda_ufs_era5.mdlus: 79%|███████▊  | 980M/1.22G [00:14<00:03, 93.3MB/s]
Downloading healda_ufs_era5.mdlus: 80%|████████  | 0.98G/1.22G [00:14<00:02, 91.9MB/s]
Downloading healda_ufs_era5.mdlus: 82%|████████▏ | 1.00G/1.22G [00:14<00:02, 96.0MB/s]
Downloading healda_ufs_era5.mdlus: 83%|████████▎ | 1.01G/1.22G [00:14<00:02, 95.2MB/s]
Downloading healda_ufs_era5.mdlus: 84%|████████▍ | 1.03G/1.22G [00:14<00:02, 100MB/s] 
Downloading healda_ufs_era5.mdlus: 85%|████████▍ | 1.04G/1.22G [00:14<00:01, 99.8MB/s]
Downloading healda_ufs_era5.mdlus: 87%|████████▋ | 1.05G/1.22G [00:15<00:01, 104MB/s] 
Downloading healda_ufs_era5.mdlus: 87%|████████▋ | 1.06G/1.22G [00:15<00:01, 96.5MB/s]
Downloading healda_ufs_era5.mdlus: 88%|████████▊ | 1.07G/1.22G [00:15<00:01, 96.8MB/s]
Downloading healda_ufs_era5.mdlus: 90%|████████▉ | 1.09G/1.22G [00:15<00:01, 99.1MB/s]
Downloading healda_ufs_era5.mdlus: 91%|█████████ | 1.10G/1.22G [00:15<00:01, 99.9MB/s]
Downloading healda_ufs_era5.mdlus: 91%|█████████▏| 1.11G/1.22G [00:15<00:01, 99.0MB/s]
Downloading healda_ufs_era5.mdlus: 93%|█████████▎| 1.13G/1.22G [00:15<00:00, 99.5MB/s]
Downloading healda_ufs_era5.mdlus: 95%|█████████▍| 1.15G/1.22G [00:16<00:00, 102MB/s] 
Downloading healda_ufs_era5.mdlus: 96%|█████████▌| 1.17G/1.22G [00:16<00:00, 105MB/s]
Downloading healda_ufs_era5.mdlus: 97%|█████████▋| 1.18G/1.22G [00:16<00:00, 103MB/s]
Downloading healda_ufs_era5.mdlus: 98%|█████████▊| 1.19G/1.22G [00:16<00:00, 96.1MB/s]
Downloading healda_ufs_era5.mdlus: 99%|█████████▊| 1.20G/1.22G [00:16<00:00, 94.7MB/s]
Downloading healda_ufs_era5.mdlus: 99%|█████████▉| 1.21G/1.22G [00:16<00:00, 92.6MB/s]
Downloading healda_ufs_era5.mdlus: 100%|██████████| 1.22G/1.22G [00:16<00:00, 77.4MB/s]

Downloading condition_hpx6_padxy.npy: 0%|          | 0.00/384k [00:00<?, ?B/s]
Downloading condition_hpx6_padxy.npy: 100%|██████████| 384k/384k [00:00<00:00, 508kB/s]
Downloading condition_hpx6_padxy.npy: 100%|██████████| 384k/384k [00:00<00:00, 505kB/s]

Downloading atms_normalizations.csv: 0%|          | 0.00/2.78k [00:00<?, ?B/s]
Downloading atms_normalizations.csv: 100%|██████████| 2.78k/2.78k [00:00<00:00, 26.1kB/s]
Downloading atms_normalizations.csv: 100%|██████████| 2.78k/2.78k [00:00<00:00, 25.4kB/s]

Downloading mhs_normalizations.csv: 0%|          | 0.00/1.28k [00:00<?, ?B/s]
Downloading mhs_normalizations.csv: 100%|██████████| 1.28k/1.28k [00:00<00:00, 11.7kB/s]
Downloading mhs_normalizations.csv: 100%|██████████| 1.28k/1.28k [00:00<00:00, 11.4kB/s]

Downloading amsua_normalizations.csv: 0%|          | 0.00/4.67k [00:00<?, ?B/s]
Downloading amsua_normalizations.csv: 100%|██████████| 4.67k/4.67k [00:00<00:00, 44.5kB/s]
Downloading amsua_normalizations.csv: 100%|██████████| 4.67k/4.67k [00:00<00:00, 43.4kB/s]

Downloading amsub_normalizations.csv: 0%|          | 0.00/675 [00:00<?, ?B/s]
Downloading amsub_normalizations.csv: 100%|██████████| 675/675 [00:00<00:00, 5.78kB/s]
Downloading amsub_normalizations.csv: 100%|██████████| 675/675 [00:00<00:00, 5.65kB/s]

Downloading conv_normalizations.csv: 0%|          | 0.00/632 [00:00<?, ?B/s]
Downloading conv_normalizations.csv: 100%|██████████| 632/632 [00:00<00:00, 5.75kB/s]
Downloading conv_normalizations.csv: 100%|██████████| 632/632 [00:00<00:00, 5.62kB/s]

Downloading era5_13_levels_stats.csv: 0%|          | 0.00/3.36k [00:00<?, ?B/s]
Downloading era5_13_levels_stats.csv: 100%|██████████| 3.36k/3.36k [00:00<00:00, 30.7kB/s]
Downloading era5_13_levels_stats.csv: 100%|██████████| 3.36k/3.36k [00:00<00:00, 30.0kB/s]

Fetch Observations¶

Pull conventional and satellite observation DataFrames for the analysis time. The UFS data sources return pandas DataFrames that match the schemas expected by HealDA.input_coords. We use earth2studio.data.fetch_dataframe which attaches request_time metadata required by the model. The time_tolerance parameter defines a time window around the analysis time so that observations will be retrieved for.

analysis_time = np.array([np.datetime64("2024-01-01T00:00")])

conv_source = UFSObsConv(time_tolerance=(timedelta(hours=-21), timedelta(hours=3)))
conv_schema, sat_schema = model.input_coords()
conv_df = fetch_dataframe(
    conv_source,
    time=analysis_time,
    variable=np.array(conv_schema["variable"]),
    fields=np.array(list(conv_schema.keys())),
)
logger.info(f"Fetched {len(conv_df)} conventional observations")

sat_source = UFSObsSat(time_tolerance=(timedelta(hours=-21), timedelta(hours=3)))
sat_df = fetch_dataframe(
    sat_source,
    time=analysis_time,
    variable=np.array(sat_schema["variable"]),
    fields=np.array(list(sat_schema.keys())),
)
logger.info(f"Fetched {len(sat_df)} satellite observations")

Console output235 lines

Fetching GSI files:   0%|          | 0/25 [00:00<?, ?it/s]
Fetching GSI files:   4%|▍         | 1/25 [00:00<00:16,  1.46it/s]
Fetching GSI files:  20%|██        | 5/25 [00:00<00:02,  7.85it/s]
Fetching GSI files:  40%|████      | 10/25 [00:00<00:00, 15.47it/s]
Fetching GSI files:  56%|█████▌    | 14/25 [00:01<00:00, 19.57it/s]
Fetching GSI files:  72%|███████▏  | 18/25 [00:01<00:00, 19.74it/s]
Fetching GSI files:  88%|████████▊ | 22/25 [00:01<00:00, 16.03it/s]
Fetching GSI files: 100%|██████████| 25/25 [00:02<00:00, 11.65it/s]
Fetching GSI files: 100%|██████████| 25/25 [00:02<00:00, 12.29it/s]
2026-08-15 05:51:25.895 | INFO     | __main__:<module>:11 - Fetched 8515112 conventional observations

Fetching GSI files:   0%|          | 0/90 [00:00<?, ?it/s]

2026-08-15 05:51:26.220 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsua_metop-a_ges.2023123118_control.nc4 not found

Fetching GSI files:   0%|          | 0/90 [00:00<?, ?it/s]
Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.223 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsub_n17_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.224 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsub_n17_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.227 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_mhs_metop-a_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.228 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsua_n17_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.230 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_mhs_metop-a_ges.2023123118_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.231 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsub_n15_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.233 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_mhs_n18_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.234 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsub_n15_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.235 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_mhs_n18_ges.2023123118_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.238 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_mhs_n18_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.246 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsua_n16_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.248 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_mhs_n18_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.252 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsub_n16_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.253 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsub_n16_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.258 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsua_n17_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.260 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsub_n15_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.262 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsua_n17_ges.2023123118_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.264 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsua_n16_ges.2023123118_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.265 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsua_n17_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.266 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsub_n15_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.273 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsub_n17_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.274 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsub_n17_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.275 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_mhs_metop-a_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.281 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsub_n16_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.285 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsua_metop-a_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.286 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsub_n15_ges.2023123118_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.287 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsua_n17_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.287 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsua_metop-a_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.289 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsub_n17_ges.2023123118_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.289 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_mhs_n18_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.290 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsua_n16_ges.2023123100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.291 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsua_metop-a_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.292 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_mhs_metop-a_ges.2023123106_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.294 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsub_n16_ges.2023123118_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.296 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsua_n16_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.299 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsub_n16_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.299 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsua_metop-a_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.301 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_mhs_metop-a_ges.2023123112_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]

2026-08-15 05:51:26.311 | WARNING  | earth2studio.data.ufs:_handle_missing_file:196 - File s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsua_n16_ges.2024010100_control.nc4 not found

Fetching GSI files:   1%|          | 1/90 [00:00<00:26,  3.31it/s]
Fetching GSI files:  46%|████▌     | 41/90 [00:01<00:01, 40.10it/s]
Fetching GSI files:  50%|█████     | 45/90 [00:01<00:01, 30.92it/s]
Fetching GSI files:  53%|█████▎    | 48/90 [00:01<00:01, 25.35it/s]
Fetching GSI files:  57%|█████▋    | 51/90 [00:01<00:01, 25.07it/s]
Fetching GSI files:  60%|██████    | 54/90 [00:02<00:01, 23.53it/s]
Fetching GSI files:  63%|██████▎   | 57/90 [00:02<00:01, 23.84it/s]
Fetching GSI files:  68%|██████▊   | 61/90 [00:02<00:01, 24.64it/s]
Fetching GSI files:  72%|███████▏  | 65/90 [00:02<00:00, 25.66it/s]
Fetching GSI files:  76%|███████▌  | 68/90 [00:02<00:01, 19.40it/s]
Fetching GSI files:  79%|███████▉  | 71/90 [00:02<00:00, 19.44it/s]
Fetching GSI files:  82%|████████▏ | 74/90 [00:03<00:01, 15.40it/s]
Fetching GSI files:  86%|████████▌ | 77/90 [00:03<00:00, 15.47it/s]
Fetching GSI files:  88%|████████▊ | 79/90 [00:03<00:00, 14.31it/s]
Fetching GSI files:  90%|█████████ | 81/90 [00:03<00:00, 12.61it/s]
Fetching GSI files:  92%|█████████▏| 83/90 [00:03<00:00, 13.32it/s]
Fetching GSI files:  94%|█████████▍| 85/90 [00:04<00:00,  6.41it/s]
Fetching GSI files:  97%|█████████▋| 87/90 [00:05<00:00,  3.77it/s]
Fetching GSI files:  98%|█████████▊| 88/90 [00:06<00:00,  3.20it/s]
Fetching GSI files:  99%|█████████▉| 89/90 [00:08<00:00,  1.70it/s]
Fetching GSI files: 100%|██████████| 90/90 [00:08<00:00,  1.71it/s]
Fetching GSI files: 100%|██████████| 90/90 [00:08<00:00, 10.34it/s]
2026-08-15 05:51:37.460 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_mhs_metop-a_ges.2023123100_control.nc4
2026-08-15 05:51:37.461 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_mhs_metop-a_ges.2023123106_control.nc4
2026-08-15 05:51:37.461 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_mhs_metop-a_ges.2023123112_control.nc4
2026-08-15 05:51:37.461 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_mhs_metop-a_ges.2023123118_control.nc4
2026-08-15 05:51:37.461 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_mhs_metop-a_ges.2024010100_control.nc4
2026-08-15 05:51:39.469 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_mhs_n18_ges.2023123100_control.nc4
2026-08-15 05:51:39.469 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_mhs_n18_ges.2023123106_control.nc4
2026-08-15 05:51:39.469 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_mhs_n18_ges.2023123112_control.nc4
2026-08-15 05:51:39.469 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_mhs_n18_ges.2023123118_control.nc4
2026-08-15 05:51:39.470 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_mhs_n18_ges.2024010100_control.nc4
2026-08-15 05:51:40.443 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsua_metop-a_ges.2023123100_control.nc4
2026-08-15 05:51:40.443 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsua_metop-a_ges.2023123106_control.nc4
2026-08-15 05:51:40.444 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsua_metop-a_ges.2023123112_control.nc4
2026-08-15 05:51:40.444 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsua_metop-a_ges.2023123118_control.nc4
2026-08-15 05:51:40.444 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsua_metop-a_ges.2024010100_control.nc4
2026-08-15 05:51:43.686 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsua_n16_ges.2023123100_control.nc4
2026-08-15 05:51:43.686 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsua_n16_ges.2023123106_control.nc4
2026-08-15 05:51:43.686 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsua_n16_ges.2023123112_control.nc4
2026-08-15 05:51:43.686 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsua_n16_ges.2023123118_control.nc4
2026-08-15 05:51:43.686 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsua_n16_ges.2024010100_control.nc4
2026-08-15 05:51:43.687 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsua_n17_ges.2023123100_control.nc4
2026-08-15 05:51:43.687 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsua_n17_ges.2023123106_control.nc4
2026-08-15 05:51:43.687 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsua_n17_ges.2023123112_control.nc4
2026-08-15 05:51:43.687 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsua_n17_ges.2023123118_control.nc4
2026-08-15 05:51:43.687 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsua_n17_ges.2024010100_control.nc4
2026-08-15 05:51:45.923 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsub_n15_ges.2023123100_control.nc4
2026-08-15 05:51:45.923 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsub_n15_ges.2023123106_control.nc4
2026-08-15 05:51:45.923 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsub_n15_ges.2023123112_control.nc4
2026-08-15 05:51:45.923 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsub_n15_ges.2023123118_control.nc4
2026-08-15 05:51:45.924 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsub_n15_ges.2024010100_control.nc4
2026-08-15 05:51:45.924 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsub_n16_ges.2023123100_control.nc4
2026-08-15 05:51:45.924 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsub_n16_ges.2023123106_control.nc4
2026-08-15 05:51:45.924 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsub_n16_ges.2023123112_control.nc4
2026-08-15 05:51:45.924 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsub_n16_ges.2023123118_control.nc4
2026-08-15 05:51:45.924 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsub_n16_ges.2024010100_control.nc4
2026-08-15 05:51:45.925 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123100/gsi/diag_amsub_n17_ges.2023123100_control.nc4
2026-08-15 05:51:45.925 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123106/gsi/diag_amsub_n17_ges.2023123106_control.nc4
2026-08-15 05:51:45.925 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123112/gsi/diag_amsub_n17_ges.2023123112_control.nc4
2026-08-15 05:51:45.925 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2023/12/2023123118/gsi/diag_amsub_n17_ges.2023123118_control.nc4
2026-08-15 05:51:45.925 | WARNING  | earth2studio.data.ufs:_compile_dataframe:223 - Cached file missing for s3://noaa-ufs-gefsv13replay-pds/2024/01/2024010100/gsi/diag_amsub_n17_ges.2024010100_control.nc4
2026-08-15 05:51:46.094 | INFO     | __main__:<module>:20 - Fetched 4378286 satellite observations

Observation Locations¶

Plot the spatial distribution of conventional and satellite observations to visualise their coverage before running the assimilation. There are 12-14 million observations typically for the model's 24-hour time window.

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

plt.close("all")
fig, axes = plt.subplots(
    1,
    2,
    subplot_kw={"projection": ccrs.Robinson()},
    figsize=(16, 4),
)

# Conventional observations
ax = axes[0]
ax.set_global()
ax.coastlines(linewidth=0.5)
ax.gridlines(linewidth=0.3, alpha=0.5)
ax.scatter(
    conv_df["lon"].values[::10],
    conv_df["lat"].values[::10],
    s=0.1,
    alpha=0.3,
    c="tab:blue",
    transform=ccrs.PlateCarree(),
)
ax.set_title(f"Conventional obs (n={len(conv_df):,})", fontsize=13)

# Satellite observations
ax = axes[1]
ax.set_global()
ax.coastlines(linewidth=0.5)
ax.gridlines(linewidth=0.3, alpha=0.5)
ax.scatter(
    sat_df["lon"].values[::10],
    sat_df["lat"].values[::10],
    s=0.1,
    alpha=0.3,
    c="tab:orange",
    transform=ccrs.PlateCarree(),
)
ax.set_title(f"Satellite obs (n={len(sat_df):,})", fontsize=13)
fig.suptitle(
    f"Observation Locations {str(analysis_time[0])[:16]} UTC",
    fontsize=15,
)
plt.tight_layout()
plt.savefig("outputs/22_healda_obs_locations.jpg", dpi=150)

Output from HealDA Global Data Assimilation

DA models can be called directly for stateless inference or via create_generator for stateful (iterative) assimilation workflows. Here we use the direct call API to invoke the model.

HealDA is designed to work with the (-21, 3) hour observation window from the UFS replay archive using both conventional and satellite observations. However, the DA model interface is flexible enough to accept different time windows and observation sources. Below we test three configurations - conventional only, satellite only, and proper combined - to illustrate the impact each observation type has on the analysis.

torch.manual_seed(42)
result_both = model(conv_obs=conv_df, sat_obs=sat_df)
logger.info(f"Combined analysis shape: {result_both.shape}")

Console output1 line

2026-08-15 05:52:01.899 | INFO     | __main__:<module>:3 - Combined analysis shape: (1, 74, 181, 360)

torch.manual_seed(42)
result_sat = model(sat_obs=sat_df)
logger.info(f"Sat-only analysis shape: {result_sat.shape}")

Console output1 line

2026-08-15 05:52:05.725 | INFO     | __main__:<module>:3 - Sat-only analysis shape: (1, 74, 181, 360)

torch.manual_seed(42)
result_conv = model(conv_obs=conv_df)
logger.info(f"Conv-only analysis shape: {result_conv.shape}")

Console output1 line

2026-08-15 05:52:10.782 | INFO     | __main__:<module>:3 - Conv-only analysis shape: (1, 74, 181, 360)

Post Processing¶

Because we loaded the model with lat_lon=True the output is already on a regular equiangular lat-lon grid, so no manual regridding is needed. Compare the three runs for surface temperature (t2m) and geopotential 500 hPa (z500). Each row shows a different observation configuration.

plt.close("all")
plot_vars = ["t2m", "z500"]
titles = ["Conv + Sat", "Sat only", "Conv only"]
results = [result_both, result_sat, result_conv]
projection = ccrs.Robinson()

fig, axes = plt.subplots(
    len(results),
    len(plot_vars),
    subplot_kw={"projection": projection},
    figsize=(14, 8),
)
fig.subplots_adjust(wspace=0.02, hspace=0.08, left=0.1, right=0.9)

lat = results[0].coords["lat"].values
lon = results[0].coords["lon"].values
cmaps = ["Spectral_r", "PRGn"]


def to_numpy(arr):
    """CuPy / Numpy helper function"""
    return arr.get() if hasattr(arr, "get") else arr


for row, (title, da) in enumerate(zip(titles, results)):
    for col, var in enumerate(plot_vars):
        ax = axes[row, col]
        field = to_numpy(da.sel(variable=var).data[0])
        im = ax.pcolormesh(
            lon,
            lat,
            field,
            transform=ccrs.PlateCarree(),
            cmap=cmaps[col],
        )
        ax.coastlines(linewidth=0.5)
        ax.gridlines(linewidth=0.3, alpha=0.5)
        fig.colorbar(im, ax=ax, shrink=0.6)
        if row == 0:
            ax.set_title(var, fontsize=14)
        if col == 0:
            ax.text(
                -0.05,
                0.5,
                title,
                fontsize=12,
                va="bottom",
                ha="center",
                rotation="vertical",
                rotation_mode="anchor",
                transform=ax.transAxes,
            )

fig.suptitle(f"HealDA Analysis {str(analysis_time[0])[:16]} UTC", fontsize=18, y=0.97)
plt.tight_layout()
plt.savefig("outputs/22_healda_analysis.jpg", dpi=150)

Output from HealDA Global Data Assimilation

HealDA vs ERA5¶

Next fetch ERA5 reanalysis at 0.25° resolution from the NCAR archive to compare against the assimilated fields. HealDA outputs standard Earth2Studio variable names so we can query ERA5 with the same identifiers. We expect the runs that are missing an observation source to show larger errors, while the combined run yields the most accurate global prediction.

era5_ds = NCAR_ERA5()
era5_da = era5_ds(analysis_time, plot_vars)
era5_interp = era5_da.interp(lat=lat, lon=lon, method="nearest")

Console output3 lines

Fetching NCAR ERA5 data:   0%|          | 0/2 [00:00<?, ?it/s]
Fetching NCAR ERA5 data:  50%|█████     | 1/2 [00:05<00:05,  5.92s/it]
Fetching NCAR ERA5 data: 100%|██████████| 2/2 [00:05<00:00,  2.99s/it]

diff_titles = ["Conv+Sat - ERA5", "Sat - ERA5", "Conv - ERA5"]
diff_results = [result_both, result_sat, result_conv]
for title, da_pred in zip(diff_titles, diff_results):
    for var in plot_vars:
        field_pred = to_numpy(da_pred.sel(variable=var).data[0])
        field_era5 = to_numpy(era5_interp.sel(variable=var).data[0])
        mae = float(np.abs(field_pred - field_era5).mean())
        logger.info(f"{title} | {var} MAE: {mae:.4f}")

Console output6 lines

2026-08-15 05:52:52.178 | INFO     | __main__:<module>:8 - Conv+Sat - ERA5 | t2m MAE: 0.7994
2026-08-15 05:52:52.179 | INFO     | __main__:<module>:8 - Conv+Sat - ERA5 | z500 MAE: 51.1967
2026-08-15 05:52:52.180 | INFO     | __main__:<module>:8 - Sat - ERA5 | t2m MAE: 0.8496
2026-08-15 05:52:52.181 | INFO     | __main__:<module>:8 - Sat - ERA5 | z500 MAE: 164.9191
2026-08-15 05:52:52.182 | INFO     | __main__:<module>:8 - Conv - ERA5 | t2m MAE: 4.3319
2026-08-15 05:52:52.182 | INFO     | __main__:<module>:8 - Conv - ERA5 | z500 MAE: 570.9412

plt.close("all")

diff_ranges = {"t2m": (-10, 10), "z500": (-500, 500)}
fig, axes = plt.subplots(
    len(diff_results),
    len(plot_vars),
    subplot_kw={"projection": projection},
    figsize=(14, 8),
)
fig.subplots_adjust(wspace=0.02, hspace=0.08, left=0.1, right=0.9)

for row, (title, da_pred) in enumerate(zip(diff_titles, diff_results)):
    for col, var in enumerate(plot_vars):
        ax = axes[row, col]
        field_pred = to_numpy(da_pred.sel(variable=var).data[0])
        field_era5 = to_numpy(era5_interp.sel(variable=var).data[0])
        diff = field_pred - field_era5
        im = ax.pcolormesh(
            lon,
            lat,
            diff,
            transform=ccrs.PlateCarree(),
            cmap="RdBu_r",
            vmin=diff_ranges[var][0],
            vmax=diff_ranges[var][1],
        )
        ax.coastlines(linewidth=0.5)
        ax.gridlines(linewidth=0.3, alpha=0.5)
        fig.colorbar(im, ax=ax, shrink=0.6)
        if row == 0:
            ax.set_title(var, fontsize=14)
        if col == 0:
            bbox = ax.get_position()
            ax.text(
                -0.05,
                0.5,
                title,
                fontsize=12,
                va="bottom",
                ha="center",
                rotation="vertical",
                rotation_mode="anchor",
                transform=ax.transAxes,
            )

fig.suptitle(
    f"HealDA Analysis Error {str(analysis_time[0])[:16]} UTC",
    fontsize=18,
    y=0.97,
)
plt.savefig("outputs/22_healda_differences.jpg", dpi=150, bbox_inches="tight")

Output from HealDA Global Data Assimilation

Download Python source Download Jupyter notebookexecuted in 167.8s

Execution profile

Runtime telemetry

Total runtime2m 48s

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