Note
Go to the end to download the full example code.
Tropical Cyclone Tracking#
Tropical cyclone tracking with tracker diagnostic models.
This example will demonstrate how to use the tropical cyclone (TC) tracker diagnostic models for creating TC paths. The diagnostics used here can be combined with other AI weather models and ensemble methods to create complex inference workflow that enable downstream analysis.
In this example you will learn:
How to instantiate a TC tracker diagnostic
How to apply the TC tracker to data
How to couple the TC tracker to a prognostic model
Post-processing results
# /// script
# dependencies = [
# "torch==2.11.0", # Match lock file to avoid torch-harmonics issue
# "earth2studio[cyclone,sfno] @ git+https://github.com/NVIDIA/earth2studio.git",
# "cartopy",
# ]
# ///
Set Up#
This example will look at tracking cyclones during August 2009, a moment in time when
multiple tropical cyclones where impacting East Asia.
Earth2Studio provides multiple variations of TC trackers such as earth2studio.models.dx.TCTrackerVitart
and earth2studio.models.dx.TCTrackerWuDuan.
The difference being the underlying algorithm used to identify the center.
This example needs the following:
Diagostic Model: Use the TC tracker
earth2studio.models.dx.TCTrackerWuDuan.Datasource: Pull data from the WB2 ERA5 data api
earth2studio.data.WB2ERA5.Prognostic Model: Use the built in FourCastNet Model
earth2studio.models.px.FCN.
import os
os.makedirs("outputs", exist_ok=True)
from dotenv import load_dotenv
load_dotenv() # TODO: make common example prep function
from datetime import datetime, timedelta
import torch
from earth2studio.data import ARCO
from earth2studio.models.dx import TCTrackerWuDuan
from earth2studio.models.px import SFNO
from earth2studio.utils.time import to_time_array
# Create tropical cyclone tracker
tracker = TCTrackerWuDuan()
# Load the default model package which downloads the check point from NGC
package = SFNO.load_default_package()
prognostic = SFNO.load_model(package)
# Create the data source
data = ARCO()
nsteps = 16 # Number of steps to run the tracker for into future
start_time = datetime(2009, 8, 5) # Start date for inference
Tracking Analysis Data#
Before coupling the TC tracker with a prognostic model, we will first apply it to analysis data. We can fetch a small time range from the data source and provide it to our model.
For the forecast we will predict for two days (these will get executed as a batch) for 20 forecast steps which is 5 days.
from earth2studio.data import fetch_data, prep_data_array
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
tracker = tracker.to(device)
# Land fall occured August 25th 2017
times = [start_time + timedelta(hours=6 * i) for i in range(nsteps + 1)]
for step, time in enumerate(times):
da = data(time, tracker.input_coords()["variable"])
x, coords = prep_data_array(da, device=device)
output, output_coords = tracker(x, coords)
print(f"Step {step}: ARCO tracks output shape {output.shape}")
era5_tracks = output.cpu()
torch.save(era5_tracks, "outputs/13_era5_paths.pt")
2026-05-18 02:17:18.193 | DEBUG | earth2studio.data.arco:_async_init:125 - Using Multi-Storage Client for ARCO data access
2026-05-18 02:17:19.125 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/.zattrs to local cache
2026-05-18 02:17:19.304 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/.zgroup to local cache
2026-05-18 02:17:19.308 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/.zmetadata to local cache
2026-05-18 02:17:19.474 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/level/0 to local cache
2026-05-18 02:17:19.638 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/model-level-1h-0p25deg.zarr-v1/.zattrs to local cache
2026-05-18 02:17:19.764 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/model-level-1h-0p25deg.zarr-v1/.zgroup to local cache
2026-05-18 02:17:19.781 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/model-level-1h-0p25deg.zarr-v1/.zmetadata to local cache
2026-05-18 02:17:19.913 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/model-level-1h-0p25deg.zarr-v1/hybrid/0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:17:20.038 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960672.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:17:20.043 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960672.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:17:20.061 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960672.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:17:20.090 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960672.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:17:20.095 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960672.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:02, 1.84it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:02<00:00, 1.33it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:03<00:00, 1.44it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:03<00:00, 1.43it/s]
Step 0: ARCO tracks output shape torch.Size([1, 7, 1, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:03.924 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960678.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:03.928 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960678.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:03.930 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960678.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:03.931 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960678.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:03.932 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960678.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.87it/s]
Fetching ARCO data: 60%|██████ | 3/5 [00:00<00:00, 5.89it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:03<00:01, 1.10s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.07it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.24it/s]
Step 1: ARCO tracks output shape torch.Size([1, 7, 2, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:08.180 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960684.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:08.181 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960684.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:08.185 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960684.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:08.186 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960684.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:08.187 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960684.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.09it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:04<00:01, 1.07s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.00s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.01it/s]
Step 2: ARCO tracks output shape torch.Size([1, 7, 3, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:13.315 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960690.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:13.318 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960690.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:13.320 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960690.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:13.323 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960690.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:13.332 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960690.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.66it/s]
Fetching ARCO data: 40%|████ | 2/5 [00:00<00:00, 4.08it/s]
Fetching ARCO data: 60%|██████ | 3/5 [00:00<00:00, 5.02it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:03<00:01, 1.05s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.34s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.02it/s]
Step 3: ARCO tracks output shape torch.Size([1, 8, 4, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:18.414 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960696.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:18.422 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960696.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:18.425 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960696.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:18.425 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960696.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:18.427 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960696.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.81it/s]
Fetching ARCO data: 40%|████ | 2/5 [00:00<00:00, 4.48it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:02<00:00, 1.19it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:05<00:00, 1.32s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:05<00:00, 1.06s/it]
Step 4: ARCO tracks output shape torch.Size([1, 9, 5, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:23.938 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960702.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:23.942 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960702.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:23.944 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960702.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:23.947 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960702.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:23.962 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960702.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 3.50it/s]
Fetching ARCO data: 40%|████ | 2/5 [00:00<00:00, 3.75it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:02<00:00, 1.17it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.07s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.11it/s]
Step 5: ARCO tracks output shape torch.Size([1, 10, 6, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:33.986 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960708.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:33.989 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960708.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:33.992 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960708.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:33.997 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960708.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:34.011 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960708.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.03it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:03<00:00, 1.29it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.09s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.02it/s]
Step 6: ARCO tracks output shape torch.Size([1, 11, 7, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:39.082 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960714.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:39.085 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960714.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:39.085 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960714.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:39.086 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960714.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:39.089 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960714.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 3.23it/s]
Fetching ARCO data: 40%|████ | 2/5 [00:00<00:00, 5.16it/s]
Fetching ARCO data: 60%|██████ | 3/5 [00:00<00:00, 6.46it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:03<00:01, 1.18s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.25s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.08it/s]
Step 7: ARCO tracks output shape torch.Size([1, 12, 8, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:43.937 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960720.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:43.941 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960720.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:43.943 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960720.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:43.950 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960720.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:43.950 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960720.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.73it/s]
Fetching ARCO data: 40%|████ | 2/5 [00:00<00:00, 4.71it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:03<00:00, 1.12it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.05s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.12it/s]
Step 8: ARCO tracks output shape torch.Size([1, 13, 9, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:48.617 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960726.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:48.620 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960726.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:48.621 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960726.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:48.623 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960726.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:48.636 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960726.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.79it/s]
Fetching ARCO data: 60%|██████ | 3/5 [00:00<00:00, 6.56it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:03<00:01, 1.01s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:05<00:00, 1.32s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:05<00:00, 1.02s/it]
Step 9: ARCO tracks output shape torch.Size([1, 14, 10, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:53.953 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960732.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:53.955 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960732.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:53.957 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960732.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:53.959 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960732.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:53.963 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960732.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.82it/s]
Fetching ARCO data: 60%|██████ | 3/5 [00:00<00:00, 6.63it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:04<00:01, 1.53s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:05<00:00, 1.34s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:05<00:00, 1.12s/it]
Step 10: ARCO tracks output shape torch.Size([1, 15, 11, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:59.741 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960738.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:59.747 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960738.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:59.750 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960738.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:59.755 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960738.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:18:59.793 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960738.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.03it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:02<00:00, 1.37it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.10it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.19it/s]
Step 11: ARCO tracks output shape torch.Size([1, 16, 12, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:04.167 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960744.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:04.168 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960744.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:04.171 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960744.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:04.186 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960744.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:04.196 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960744.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:02, 1.94it/s]
Fetching ARCO data: 60%|██████ | 3/5 [00:00<00:00, 5.35it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:03<00:00, 1.06it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.01it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.22it/s]
Step 12: ARCO tracks output shape torch.Size([1, 17, 13, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:11.091 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960750.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:11.094 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960750.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:11.094 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960750.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:11.097 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960750.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:11.099 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960750.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.26it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:02<00:00, 1.39it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.02it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.12it/s]
Step 13: ARCO tracks output shape torch.Size([1, 18, 14, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:15.751 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960756.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:15.756 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960756.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:15.760 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960756.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:15.766 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960756.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:15.771 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960756.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.05it/s]
Fetching ARCO data: 60%|██████ | 3/5 [00:00<00:00, 6.11it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.14s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.06it/s]
Step 14: ARCO tracks output shape torch.Size([1, 19, 15, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:20.687 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960762.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:20.688 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960762.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:20.694 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960762.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:20.695 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960762.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:20.703 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960762.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.04it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:04<00:01, 1.06s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.12it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.11it/s]
Step 15: ARCO tracks output shape torch.Size([1, 20, 16, 4])
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:25.425 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/u_component_of_wind/960768.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:25.425 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_u_component_of_wind/960768.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:25.427 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/10m_v_component_of_wind/960768.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:25.429 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/mean_sea_level_pressure/960768.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
2026-05-18 02:19:25.430 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/v_component_of_wind/960768.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/5 [00:00<?, ?it/s]
Fetching ARCO data: 20%|██ | 1/5 [00:00<00:01, 2.71it/s]
Fetching ARCO data: 40%|████ | 2/5 [00:00<00:00, 4.68it/s]
Fetching ARCO data: 80%|████████ | 4/5 [00:02<00:00, 1.30it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.02s/it]
Fetching ARCO data: 100%|██████████| 5/5 [00:04<00:00, 1.18it/s]
Step 16: ARCO tracks output shape torch.Size([1, 21, 17, 4])
Notice that the output tensor grows as iterations are performed. This is because the tracker builds tracks based on previous forward passes returning a tensor with the dimensions [batch, path, step, variable]. Not all paths are guaranteed to be the same length or have the same start / stop time so any missing data is populated with a nan value.
Up next lets also repeat the same process using the prognostic AI model. One could use one of the build in workflows but here we will manually implement the inference loop.
from tqdm import tqdm
from earth2studio.utils.coords import map_coords
prognostic = prognostic.to(device)
# Reset the internal path buffer of tracker
tracker.reset_path_buffer()
# Load the initial state
x, coords = fetch_data(
source=data,
time=to_time_array([start_time]),
variable=prognostic.input_coords()["variable"],
lead_time=prognostic.input_coords()["lead_time"],
device=device,
)
# Create prognostic iterator
model = prognostic.create_iterator(x, coords)
with tqdm(total=nsteps + 1, desc="Running inference") as pbar:
for step, (x, coords) in enumerate(model):
# Run tracker
x, coords = map_coords(x, coords, tracker.input_coords())
output, output_coords = tracker(x, coords)
# lets remove the lead time dim
output = output[:, 0]
print(f"Step {step}: SFNO tracks output shape {output.shape}")
pbar.update(1)
if step == nsteps:
break
sfno_tracks = output.cpu()
torch.save(sfno_tracks, "outputs/13_sfno_paths.pt")
Fetching ARCO data: 0%| | 0/13 [00:00<?, ?it/s]
2026-05-18 02:19:30.332 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/geopotential/960672.0.0.0 to local cache
Fetching ARCO data: 0%| | 0/13 [00:00<?, ?it/s]
2026-05-18 02:19:30.332 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/2m_temperature/960672.0.0 to local cache
Fetching ARCO data: 0%| | 0/13 [00:00<?, ?it/s]
2026-05-18 02:19:30.337 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/total_column_water_vapour/960672.0.0 to local cache
Fetching ARCO data: 0%| | 0/13 [00:00<?, ?it/s]
Fetching ARCO data: 8%|▊ | 1/13 [00:00<00:01, 7.15it/s]
2026-05-18 02:19:30.379 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/surface_pressure/960672.0.0 to local cache
Fetching ARCO data: 8%|▊ | 1/13 [00:00<00:01, 7.15it/s]
2026-05-18 02:19:30.386 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/specific_humidity/960672.0.0.0 to local cache
Fetching ARCO data: 8%|▊ | 1/13 [00:00<00:01, 7.15it/s]
2026-05-18 02:19:30.400 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/100m_u_component_of_wind/960672.0.0 to local cache
Fetching ARCO data: 8%|▊ | 1/13 [00:00<00:01, 7.15it/s]
2026-05-18 02:19:30.407 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/temperature/960672.0.0.0 to local cache
Fetching ARCO data: 8%|▊ | 1/13 [00:00<00:01, 7.15it/s]
2026-05-18 02:19:30.525 | DEBUG | earth2studio.data.utils:_make_local_details:1066 - Copying /gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3/100m_v_component_of_wind/960672.0.0 to local cache
Fetching ARCO data: 8%|▊ | 1/13 [00:00<00:01, 7.15it/s]
Fetching ARCO data: 31%|███ | 4/13 [00:00<00:00, 13.07it/s]
Fetching ARCO data: 46%|████▌ | 6/13 [00:00<00:00, 15.10it/s]
Fetching ARCO data: 62%|██████▏ | 8/13 [00:00<00:00, 10.80it/s]
Fetching ARCO data: 77%|███████▋ | 10/13 [00:00<00:00, 11.37it/s]
Fetching ARCO data: 92%|█████████▏| 12/13 [00:02<00:00, 2.52it/s]
Fetching ARCO data: 100%|██████████| 13/13 [00:04<00:00, 1.44it/s]
Fetching ARCO data: 100%|██████████| 13/13 [00:04<00:00, 2.70it/s]
Running inference: 0%| | 0/17 [00:00<?, ?it/s]Step 0: SFNO tracks output shape torch.Size([1, 7, 1, 4])
Running inference: 6%|▌ | 1/17 [00:00<00:03, 4.63it/s]Step 1: SFNO tracks output shape torch.Size([1, 8, 2, 4])
Running inference: 12%|█▏ | 2/17 [00:00<00:04, 3.49it/s]Step 2: SFNO tracks output shape torch.Size([1, 8, 3, 4])
Running inference: 18%|█▊ | 3/17 [00:00<00:04, 3.32it/s]Step 3: SFNO tracks output shape torch.Size([1, 9, 4, 4])
Running inference: 24%|██▎ | 4/17 [00:01<00:03, 3.29it/s]Step 4: SFNO tracks output shape torch.Size([1, 10, 5, 4])
Running inference: 29%|██▉ | 5/17 [00:01<00:03, 3.25it/s]Step 5: SFNO tracks output shape torch.Size([1, 10, 6, 4])
Running inference: 35%|███▌ | 6/17 [00:01<00:03, 3.16it/s]Step 6: SFNO tracks output shape torch.Size([1, 10, 7, 4])
Running inference: 41%|████ | 7/17 [00:02<00:03, 3.25it/s]Step 7: SFNO tracks output shape torch.Size([1, 11, 8, 4])
Running inference: 47%|████▋ | 8/17 [00:02<00:02, 3.18it/s]Step 8: SFNO tracks output shape torch.Size([1, 12, 9, 4])
Running inference: 53%|█████▎ | 9/17 [00:02<00:02, 3.14it/s]Step 9: SFNO tracks output shape torch.Size([1, 13, 10, 4])
Running inference: 59%|█████▉ | 10/17 [00:03<00:02, 3.21it/s]Step 10: SFNO tracks output shape torch.Size([1, 13, 11, 4])
Running inference: 65%|██████▍ | 11/17 [00:03<00:01, 3.33it/s]Step 11: SFNO tracks output shape torch.Size([1, 14, 12, 4])
Running inference: 71%|███████ | 12/17 [00:03<00:01, 3.34it/s]Step 12: SFNO tracks output shape torch.Size([1, 15, 13, 4])
Running inference: 76%|███████▋ | 13/17 [00:03<00:01, 3.26it/s]Step 13: SFNO tracks output shape torch.Size([1, 16, 14, 4])
Running inference: 82%|████████▏ | 14/17 [00:04<00:00, 3.22it/s]Step 14: SFNO tracks output shape torch.Size([1, 17, 15, 4])
Running inference: 88%|████████▊ | 15/17 [00:04<00:00, 3.30it/s]Step 15: SFNO tracks output shape torch.Size([1, 18, 16, 4])
Running inference: 94%|█████████▍| 16/17 [00:04<00:00, 3.26it/s]Step 16: SFNO tracks output shape torch.Size([1, 19, 17, 4])
Running inference: 100%|██████████| 17/17 [00:05<00:00, 3.18it/s]
Running inference: 100%|██████████| 17/17 [00:05<00:00, 3.26it/s]
Note that before the inference loop of the AI model, the path buffer of the tracker was reset which refreshes the tracker’s path state starting from scratch. Otherwise, it would attempt to append to the existing tracks from the data source loop in the previous section.
Finally we can plot the results to compare the track ground truths from ERA5 with
those produced by SFNO.
Recall the outputs of earth2studio.models.dx.TCTrackerWuDuan has the path
ID in the second dimension, thus that is what will determine the number of lines.
The lat/lon coords are the first two variables in the last dimension.
Lastly we just need to be mindful of the NaN filler values which can get easily
masked out and any “path” that isnt over 2 points long
from datetime import datetime, timedelta
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
# Convert tracks from tensors to numpy arrays
era5_paths = era5_tracks.numpy()
sfno_paths = sfno_tracks.numpy()
# Calculate end date
end_time = start_time + timedelta(hours=6 * nsteps)
# Create figure with cartopy projection
plt.figure(figsize=(10, 8))
projection = ccrs.LambertConformal(
central_longitude=130.0, central_latitude=30.0, standard_parallels=(20.0, 40.0)
)
ax = plt.axes(projection=projection)
# Add map features
ax.add_feature(cfeature.COASTLINE)
ax.add_feature(cfeature.LAND, alpha=0.1)
ax.gridlines(draw_labels=True, alpha=0.6)
ax.set_extent([90, 170, 0, 50], crs=ccrs.PlateCarree())
era5_cmap = plt.cm.autumn
sfno_cmap = plt.cm.winter
for path in range(era5_paths.shape[1]):
# Get lat/lon coordinates, filtering out nans
lats = era5_paths[0, path, :, 0]
lons = era5_paths[0, path, :, 1]
mask = ~np.isnan(lats) & ~np.isnan(lons)
if mask.any() and len(lons[mask]) > 2:
color = era5_cmap(path / era5_paths.shape[1])
ax.plot(
lons[mask],
lats[mask],
color=color,
linestyle="-.",
marker="x",
label="ERA5" if path == 0 else "",
transform=ccrs.PlateCarree(),
)
for path in range(sfno_paths.shape[1]):
# Get lat/lon coordinates, filtering out nans
lats = sfno_paths[0, path, :, 0]
lons = sfno_paths[0, path, :, 1]
mask = ~np.isnan(lats) & ~np.isnan(lons)
if mask.any() and len(lons[mask]) > 2:
color = sfno_cmap(path / sfno_paths.shape[1])
ax.plot(
lons[mask],
lats[mask],
color=color,
linestyle="-",
label="SFNO" if path == 0 else "",
transform=ccrs.PlateCarree(),
)
era5_patch = mpatches.Rectangle(
(0, 0), 1, 1, fc=era5_cmap(0.3), alpha=0.9, label="ERA5"
)
sfno_patch = mpatches.Rectangle(
(0, 0), 1, 1, fc=sfno_cmap(0.3), alpha=0.9, label="SFNO"
)
ax.legend(handles=[era5_patch, sfno_patch], loc="upper right", title="Cyclone Tracks")
plt.title(
f'Tropical Cyclone Tracks\n{start_time.strftime("%Y-%m-%d")} to {end_time.strftime("%Y-%m-%d")}'
)
plt.savefig(f"outputs/13_{start_time}_cyclone_tracks.jpg", bbox_inches="tight", dpi=300)
In addition to filtering out the NaN values, users may want to apply other post processing steps on the paths which may be enforcing path lengths are above a certain threshold or other geography based filters.
No cyclone tracker is perfect, we encourage users to experiment and tune the tracker as needed.
Total running time of the script: (2 minutes 33.262 seconds)