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

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

# 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

2025-04-18 19:04:02.954 | WARNING  | earth2studio.data.arco:__init__:112 - Using Zarr 2.0 method for ARCO, this can be extremely slow with caching!

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, "era5.pt")

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:16.439 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:16.445 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:16.450 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:16.454 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:16.535 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-05T00:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00, 41.06it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 30.36it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 30.31it/s]
Step 0: ARCO tracks output shape torch.Size([1, 7, 1, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.154 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-05T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.157 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-05T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.160 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-05T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.162 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-05T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.219 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-05T06:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00, 59.47it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.38it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.34it/s]
Step 1: ARCO tracks output shape torch.Size([1, 7, 2, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.478 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-05T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.482 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-05T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.485 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-05T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.487 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-05T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.543 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-05T12:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00, 60.43it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.15it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.12it/s]
Step 2: ARCO tracks output shape torch.Size([1, 7, 3, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.808 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-05T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.811 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-05T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.814 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-05T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.816 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-05T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:46.870 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-05T18:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00, 63.19it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.73it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.69it/s]
Step 3: ARCO tracks output shape torch.Size([1, 8, 4, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.126 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-06T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.129 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-06T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.132 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-06T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.135 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-06T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.192 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-06T00:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00, 59.34it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 27.90it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 27.86it/s]
Step 4: ARCO tracks output shape torch.Size([1, 9, 5, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.470 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-06T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.473 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-06T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.476 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-06T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.478 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-06T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:47.535 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-06T06:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00, 60.12it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.07it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 28.04it/s]
Step 5: ARCO tracks output shape torch.Size([1, 10, 6, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.603 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-06T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.606 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-06T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.609 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-06T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.612 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-06T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.67it/s]

2025-04-18 19:04:51.721 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-06T12:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.67it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.22it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.08it/s]
Step 6: ARCO tracks output shape torch.Size([1, 11, 7, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.973 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-06T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.976 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-06T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.979 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-06T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:51.982 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-06T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.49it/s]

2025-04-18 19:04:52.094 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-06T18:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.49it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.49it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.21it/s]
Step 7: ARCO tracks output shape torch.Size([1, 12, 8, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.347 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-07T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.350 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-07T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.353 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-07T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.355 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-07T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.40it/s]

2025-04-18 19:04:52.469 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-07T00:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.40it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.51it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.20it/s]
Step 8: ARCO tracks output shape torch.Size([1, 13, 9, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.742 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-07T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.745 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-07T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.748 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-07T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:52.751 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-07T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.54it/s]

2025-04-18 19:04:52.862 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-07T06:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.54it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.54it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.27it/s]
Step 9: ARCO tracks output shape torch.Size([1, 14, 10, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.126 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-07T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.130 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-07T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.133 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-07T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.135 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-07T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.48it/s]

2025-04-18 19:04:53.247 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-07T12:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.48it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.56it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.26it/s]
Step 10: ARCO tracks output shape torch.Size([1, 15, 11, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.497 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-07T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.500 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-07T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.503 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-07T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.505 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-07T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.48it/s]

2025-04-18 19:04:53.617 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-07T18:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.48it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.55it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.26it/s]
Step 11: ARCO tracks output shape torch.Size([1, 16, 12, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.871 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-08T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.874 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-08T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.877 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-08T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:53.879 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-08T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.57it/s]

2025-04-18 19:04:53.990 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-08T00:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.57it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.41it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.18it/s]
Step 12: ARCO tracks output shape torch.Size([1, 17, 13, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.108 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-08T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.111 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-08T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.114 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-08T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.117 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-08T06:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.55it/s]

2025-04-18 19:04:56.228 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-08T06:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.55it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.55it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.29it/s]
Step 13: ARCO tracks output shape torch.Size([1, 18, 14, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.481 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-08T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.484 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-08T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.487 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-08T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.489 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-08T12:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.63it/s]

2025-04-18 19:04:56.599 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-08T12:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.63it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.78it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.49it/s]
Step 14: ARCO tracks output shape torch.Size([1, 19, 15, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.843 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-08T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.846 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-08T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.849 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-08T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:56.851 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-08T18:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.58it/s]

2025-04-18 19:04:56.962 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-08T18:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.58it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.73it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.44it/s]
Step 15: ARCO tracks output shape torch.Size([1, 20, 16, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:57.217 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-09T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:57.220 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-09T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:57.223 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-09T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]

2025-04-18 19:04:57.226 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-09T00:00:00

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.53it/s]

2025-04-18 19:04:57.337 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-09T00:00:00

Fetching ARCO data:  80%|████████  | 4/5 [00:00<00:00,  8.53it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 23.68it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:00<00:00, 21.37it/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 garenteed 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=[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, "sfno.pt")

Fetching ARCO data:   0%|          | 0/73 [00:00<?, ?it/s]

2025-04-18 19:04:57.933 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u10m at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/73 [00:00<?, ?it/s]

2025-04-18 19:04:57.936 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v10m at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/73 [00:00<?, ?it/s]

2025-04-18 19:04:57.940 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u100m at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/73 [00:00<?, ?it/s]

2025-04-18 19:04:57.942 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v100m at 2009-08-05T00:00:00

Fetching ARCO data:   0%|          | 0/73 [00:00<?, ?it/s]

2025-04-18 19:04:57.950 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t2m at 2009-08-05T00:00:00

Fetching ARCO data:   5%|▌         | 4/73 [00:00<00:00, 219.21it/s]

2025-04-18 19:04:57.954 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: sp at 2009-08-05T00:00:00

Fetching ARCO data:   5%|▌         | 4/73 [00:00<00:00, 187.19it/s]

2025-04-18 19:04:57.956 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: msl at 2009-08-05T00:00:00

Fetching ARCO data:   5%|▌         | 4/73 [00:00<00:00, 165.74it/s]

2025-04-18 19:04:57.959 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: tcwv at 2009-08-05T00:00:00

Fetching ARCO data:   5%|▌         | 4/73 [00:00<00:00, 149.12it/s]

2025-04-18 19:04:57.966 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u50 at 2009-08-05T00:00:00

Fetching ARCO data:  11%|█         | 8/73 [00:00<00:00, 232.84it/s]

2025-04-18 19:04:58.020 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u100 at 2009-08-05T00:00:00

Fetching ARCO data:  11%|█         | 8/73 [00:00<00:00, 90.89it/s]

2025-04-18 19:04:58.078 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u150 at 2009-08-05T00:00:00

Fetching ARCO data:  11%|█         | 8/73 [00:00<00:01, 54.90it/s]

2025-04-18 19:04:58.136 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u200 at 2009-08-05T00:00:00

Fetching ARCO data:  11%|█         | 8/73 [00:00<00:01, 39.30it/s]
Fetching ARCO data:  12%|█▏        | 9/73 [00:00<00:01, 34.37it/s]

2025-04-18 19:04:58.199 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u250 at 2009-08-05T00:00:00

Fetching ARCO data:  16%|█▋        | 12/73 [00:00<00:01, 34.37it/s]

2025-04-18 19:04:58.257 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u300 at 2009-08-05T00:00:00

Fetching ARCO data:  16%|█▋        | 12/73 [00:00<00:01, 34.37it/s]

2025-04-18 19:04:58.315 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u400 at 2009-08-05T00:00:00

Fetching ARCO data:  16%|█▋        | 12/73 [00:00<00:01, 34.37it/s]

2025-04-18 19:04:58.372 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u500 at 2009-08-05T00:00:00

Fetching ARCO data:  16%|█▋        | 12/73 [00:00<00:01, 34.37it/s]
Fetching ARCO data:  18%|█▊        | 13/73 [00:00<00:02, 24.57it/s]

2025-04-18 19:04:58.435 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u600 at 2009-08-05T00:00:00

Fetching ARCO data:  22%|██▏       | 16/73 [00:00<00:02, 24.57it/s]

2025-04-18 19:04:58.492 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u700 at 2009-08-05T00:00:00

Fetching ARCO data:  22%|██▏       | 16/73 [00:00<00:02, 24.57it/s]

2025-04-18 19:04:58.548 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u850 at 2009-08-05T00:00:00

Fetching ARCO data:  22%|██▏       | 16/73 [00:00<00:02, 24.57it/s]

2025-04-18 19:04:58.605 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u925 at 2009-08-05T00:00:00

Fetching ARCO data:  22%|██▏       | 16/73 [00:00<00:02, 24.57it/s]
Fetching ARCO data:  23%|██▎       | 17/73 [00:00<00:02, 21.37it/s]

2025-04-18 19:04:58.666 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: u1000 at 2009-08-05T00:00:00

Fetching ARCO data:  27%|██▋       | 20/73 [00:00<00:02, 21.37it/s]

2025-04-18 19:04:58.723 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v50 at 2009-08-05T00:00:00

Fetching ARCO data:  27%|██▋       | 20/73 [00:00<00:02, 21.37it/s]

2025-04-18 19:04:58.781 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v100 at 2009-08-05T00:00:00

Fetching ARCO data:  27%|██▋       | 20/73 [00:00<00:02, 21.37it/s]

2025-04-18 19:04:58.839 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v150 at 2009-08-05T00:00:00

Fetching ARCO data:  27%|██▋       | 20/73 [00:00<00:02, 21.37it/s]
Fetching ARCO data:  29%|██▉       | 21/73 [00:00<00:02, 19.64it/s]

2025-04-18 19:04:58.902 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v200 at 2009-08-05T00:00:00

Fetching ARCO data:  33%|███▎      | 24/73 [00:00<00:02, 19.64it/s]

2025-04-18 19:04:58.961 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v250 at 2009-08-05T00:00:00

Fetching ARCO data:  33%|███▎      | 24/73 [00:01<00:02, 19.64it/s]

2025-04-18 19:04:59.019 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v300 at 2009-08-05T00:00:00

Fetching ARCO data:  33%|███▎      | 24/73 [00:01<00:02, 19.64it/s]

2025-04-18 19:04:59.077 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v400 at 2009-08-05T00:00:00

Fetching ARCO data:  33%|███▎      | 24/73 [00:01<00:02, 19.64it/s]
Fetching ARCO data:  34%|███▍      | 25/73 [00:01<00:02, 18.63it/s]

2025-04-18 19:04:59.140 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v500 at 2009-08-05T00:00:00

Fetching ARCO data:  38%|███▊      | 28/73 [00:01<00:02, 18.63it/s]

2025-04-18 19:04:59.198 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v600 at 2009-08-05T00:00:00

Fetching ARCO data:  38%|███▊      | 28/73 [00:01<00:02, 18.63it/s]

2025-04-18 19:04:59.256 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v700 at 2009-08-05T00:00:00

Fetching ARCO data:  38%|███▊      | 28/73 [00:01<00:02, 18.63it/s]

2025-04-18 19:04:59.314 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v850 at 2009-08-05T00:00:00

Fetching ARCO data:  38%|███▊      | 28/73 [00:01<00:02, 18.63it/s]
Fetching ARCO data:  40%|███▉      | 29/73 [00:01<00:02, 18.04it/s]

2025-04-18 19:04:59.377 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v925 at 2009-08-05T00:00:00

Fetching ARCO data:  44%|████▍     | 32/73 [00:01<00:02, 18.04it/s]

2025-04-18 19:04:59.434 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: v1000 at 2009-08-05T00:00:00

Fetching ARCO data:  44%|████▍     | 32/73 [00:01<00:02, 18.04it/s]

2025-04-18 19:04:59.491 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z50 at 2009-08-05T00:00:00

Fetching ARCO data:  44%|████▍     | 32/73 [00:01<00:02, 18.04it/s]

2025-04-18 19:04:59.544 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z100 at 2009-08-05T00:00:00

Fetching ARCO data:  44%|████▍     | 32/73 [00:01<00:02, 18.04it/s]
Fetching ARCO data:  45%|████▌     | 33/73 [00:01<00:02, 17.95it/s]

2025-04-18 19:04:59.603 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z150 at 2009-08-05T00:00:00

Fetching ARCO data:  49%|████▉     | 36/73 [00:01<00:02, 17.95it/s]

2025-04-18 19:04:59.655 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z200 at 2009-08-05T00:00:00

Fetching ARCO data:  49%|████▉     | 36/73 [00:01<00:02, 17.95it/s]

2025-04-18 19:04:59.706 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z250 at 2009-08-05T00:00:00

Fetching ARCO data:  49%|████▉     | 36/73 [00:01<00:02, 17.95it/s]

2025-04-18 19:04:59.759 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z300 at 2009-08-05T00:00:00

Fetching ARCO data:  49%|████▉     | 36/73 [00:01<00:02, 17.95it/s]
Fetching ARCO data:  51%|█████     | 37/73 [00:01<00:01, 18.19it/s]

2025-04-18 19:04:59.815 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z400 at 2009-08-05T00:00:00

Fetching ARCO data:  55%|█████▍    | 40/73 [00:01<00:01, 18.19it/s]

2025-04-18 19:04:59.867 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z500 at 2009-08-05T00:00:00

Fetching ARCO data:  55%|█████▍    | 40/73 [00:01<00:01, 18.19it/s]

2025-04-18 19:04:59.919 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z600 at 2009-08-05T00:00:00

Fetching ARCO data:  55%|█████▍    | 40/73 [00:01<00:01, 18.19it/s]

2025-04-18 19:04:59.970 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z700 at 2009-08-05T00:00:00

Fetching ARCO data:  55%|█████▍    | 40/73 [00:02<00:01, 18.19it/s]
Fetching ARCO data:  56%|█████▌    | 41/73 [00:02<00:01, 18.40it/s]

2025-04-18 19:05:00.027 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z850 at 2009-08-05T00:00:00

Fetching ARCO data:  60%|██████    | 44/73 [00:02<00:01, 18.40it/s]

2025-04-18 19:05:00.079 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z925 at 2009-08-05T00:00:00

Fetching ARCO data:  60%|██████    | 44/73 [00:02<00:01, 18.40it/s]

2025-04-18 19:05:00.131 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: z1000 at 2009-08-05T00:00:00

Fetching ARCO data:  60%|██████    | 44/73 [00:02<00:01, 18.40it/s]

2025-04-18 19:05:00.183 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t50 at 2009-08-05T00:00:00

Fetching ARCO data:  60%|██████    | 44/73 [00:02<00:01, 18.40it/s]
Fetching ARCO data:  62%|██████▏   | 45/73 [00:02<00:01, 18.43it/s]

2025-04-18 19:05:00.243 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t100 at 2009-08-05T00:00:00

Fetching ARCO data:  66%|██████▌   | 48/73 [00:02<00:01, 18.43it/s]

2025-04-18 19:05:00.297 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t150 at 2009-08-05T00:00:00

Fetching ARCO data:  66%|██████▌   | 48/73 [00:02<00:01, 18.43it/s]

2025-04-18 19:05:00.351 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t200 at 2009-08-05T00:00:00

Fetching ARCO data:  66%|██████▌   | 48/73 [00:02<00:01, 18.43it/s]

2025-04-18 19:05:00.406 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t250 at 2009-08-05T00:00:00

Fetching ARCO data:  66%|██████▌   | 48/73 [00:02<00:01, 18.43it/s]
Fetching ARCO data:  67%|██████▋   | 49/73 [00:02<00:01, 18.31it/s]

2025-04-18 19:05:00.465 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t300 at 2009-08-05T00:00:00

Fetching ARCO data:  71%|███████   | 52/73 [00:02<00:01, 18.31it/s]

2025-04-18 19:05:00.518 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t400 at 2009-08-05T00:00:00

Fetching ARCO data:  71%|███████   | 52/73 [00:02<00:01, 18.31it/s]

2025-04-18 19:05:00.571 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t500 at 2009-08-05T00:00:00

Fetching ARCO data:  71%|███████   | 52/73 [00:02<00:01, 18.31it/s]

2025-04-18 19:05:00.625 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t600 at 2009-08-05T00:00:00

Fetching ARCO data:  71%|███████   | 52/73 [00:02<00:01, 18.31it/s]
Fetching ARCO data:  73%|███████▎  | 53/73 [00:02<00:01, 18.32it/s]

2025-04-18 19:05:00.683 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t700 at 2009-08-05T00:00:00

Fetching ARCO data:  77%|███████▋  | 56/73 [00:02<00:00, 18.32it/s]

2025-04-18 19:05:00.736 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t850 at 2009-08-05T00:00:00

Fetching ARCO data:  77%|███████▋  | 56/73 [00:02<00:00, 18.32it/s]

2025-04-18 19:05:00.790 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t925 at 2009-08-05T00:00:00

Fetching ARCO data:  77%|███████▋  | 56/73 [00:02<00:00, 18.32it/s]

2025-04-18 19:05:00.843 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: t1000 at 2009-08-05T00:00:00

Fetching ARCO data:  77%|███████▋  | 56/73 [00:02<00:00, 18.32it/s]
Fetching ARCO data:  78%|███████▊  | 57/73 [00:02<00:00, 18.32it/s]

2025-04-18 19:05:00.902 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q50 at 2009-08-05T00:00:00

Fetching ARCO data:  82%|████████▏ | 60/73 [00:02<00:00, 18.32it/s]

2025-04-18 19:05:00.957 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q100 at 2009-08-05T00:00:00

Fetching ARCO data:  82%|████████▏ | 60/73 [00:03<00:00, 18.32it/s]

2025-04-18 19:05:01.014 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q150 at 2009-08-05T00:00:00

Fetching ARCO data:  82%|████████▏ | 60/73 [00:03<00:00, 18.32it/s]

2025-04-18 19:05:01.070 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q200 at 2009-08-05T00:00:00

Fetching ARCO data:  82%|████████▏ | 60/73 [00:03<00:00, 18.32it/s]
Fetching ARCO data:  84%|████████▎ | 61/73 [00:03<00:00, 18.03it/s]

2025-04-18 19:05:01.132 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q250 at 2009-08-05T00:00:00

Fetching ARCO data:  88%|████████▊ | 64/73 [00:03<00:00, 18.03it/s]

2025-04-18 19:05:01.188 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q300 at 2009-08-05T00:00:00

Fetching ARCO data:  88%|████████▊ | 64/73 [00:03<00:00, 18.03it/s]

2025-04-18 19:05:01.245 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q400 at 2009-08-05T00:00:00

Fetching ARCO data:  88%|████████▊ | 64/73 [00:03<00:00, 18.03it/s]

2025-04-18 19:05:01.301 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q500 at 2009-08-05T00:00:00

Fetching ARCO data:  88%|████████▊ | 64/73 [00:03<00:00, 18.03it/s]
Fetching ARCO data:  89%|████████▉ | 65/73 [00:03<00:00, 17.82it/s]

2025-04-18 19:05:01.362 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q600 at 2009-08-05T00:00:00

Fetching ARCO data:  93%|█████████▎| 68/73 [00:03<00:00, 17.82it/s]

2025-04-18 19:05:01.419 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q700 at 2009-08-05T00:00:00

Fetching ARCO data:  93%|█████████▎| 68/73 [00:03<00:00, 17.82it/s]

2025-04-18 19:05:01.475 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q850 at 2009-08-05T00:00:00

Fetching ARCO data:  93%|█████████▎| 68/73 [00:03<00:00, 17.82it/s]

2025-04-18 19:05:01.531 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q925 at 2009-08-05T00:00:00

Fetching ARCO data:  93%|█████████▎| 68/73 [00:03<00:00, 17.82it/s]
Fetching ARCO data:  95%|█████████▍| 69/73 [00:03<00:00, 17.67it/s]

2025-04-18 19:05:01.593 | DEBUG    | earth2studio.data.arco:fetch_array:253 - Fetching ARCO zarr array for variable: q1000 at 2009-08-05T00:00:00

Fetching ARCO data:  99%|█████████▊| 72/73 [00:03<00:00, 17.67it/s]
Fetching ARCO data: 100%|██████████| 73/73 [00:03<00:00, 19.64it/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:02,  6.31it/s]Step 1: SFNO tracks output shape torch.Size([1, 8, 2, 4])

Running inference:  12%|█▏        | 2/17 [00:00<00:05,  2.80it/s]Step 2: SFNO tracks output shape torch.Size([1, 8, 3, 4])

Running inference:  18%|█▊        | 3/17 [00:01<00:05,  2.40it/s]Step 3: SFNO tracks output shape torch.Size([1, 9, 4, 4])

Running inference:  24%|██▎       | 4/17 [00:01<00:05,  2.28it/s]Step 4: SFNO tracks output shape torch.Size([1, 10, 5, 4])

Running inference:  29%|██▉       | 5/17 [00:02<00:05,  2.20it/s]Step 5: SFNO tracks output shape torch.Size([1, 10, 6, 4])

Running inference:  35%|███▌      | 6/17 [00:02<00:05,  2.14it/s]Step 6: SFNO tracks output shape torch.Size([1, 10, 7, 4])

Running inference:  41%|████      | 7/17 [00:03<00:04,  2.14it/s]Step 7: SFNO tracks output shape torch.Size([1, 11, 8, 4])

Running inference:  47%|████▋     | 8/17 [00:03<00:04,  2.11it/s]Step 8: SFNO tracks output shape torch.Size([1, 12, 9, 4])

Running inference:  53%|█████▎    | 9/17 [00:04<00:03,  2.09it/s]Step 9: SFNO tracks output shape torch.Size([1, 13, 10, 4])

Running inference:  59%|█████▉    | 10/17 [00:04<00:03,  2.11it/s]Step 10: SFNO tracks output shape torch.Size([1, 13, 11, 4])

Running inference:  65%|██████▍   | 11/17 [00:04<00:02,  2.14it/s]Step 11: SFNO tracks output shape torch.Size([1, 14, 12, 4])

Running inference:  71%|███████   | 12/17 [00:05<00:02,  2.14it/s]Step 12: SFNO tracks output shape torch.Size([1, 15, 13, 4])

Running inference:  76%|███████▋  | 13/17 [00:05<00:01,  2.12it/s]Step 13: SFNO tracks output shape torch.Size([1, 16, 14, 4])

Running inference:  82%|████████▏ | 14/17 [00:06<00:01,  2.10it/s]Step 14: SFNO tracks output shape torch.Size([1, 17, 15, 4])

Running inference:  88%|████████▊ | 15/17 [00:06<00:00,  2.13it/s]Step 15: SFNO tracks output shape torch.Size([1, 17, 16, 4])

Running inference:  94%|█████████▍| 16/17 [00:07<00:00,  2.12it/s]Step 16: SFNO tracks output shape torch.Size([1, 18, 17, 4])

Running inference: 100%|██████████| 17/17 [00:07<00:00,  2.09it/s]
Running inference: 100%|██████████| 17/17 [00:07<00:00,  2.18it/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)

Tropical Cyclone Tracks 2009-08-05 to 2009-08-09

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: (1 minutes 17.001 seconds)

Gallery generated by Sphinx-Gallery