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
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

Console output9 lines

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

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")

Console output116 lines

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.74it/s]
Fetching ARCO data:  60%|██████    | 3/5 [00:00<00:00,  8.21it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.21it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:03<00:00,  1.32it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:03<00:00,  1.59it/s]
Step 0: ARCO tracks output shape torch.Size([1, 7, 1, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  6.96it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.77it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.14it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.13it/s]
Step 1: ARCO tracks output shape torch.Size([1, 7, 2, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.05it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:01<00:00,  1.91it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.21it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.23it/s]
Step 2: ARCO tracks output shape torch.Size([1, 7, 3, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  5.97it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.53it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.95it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.91it/s]
Step 3: ARCO tracks output shape torch.Size([1, 8, 4, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  5.24it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.48it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.87it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.83it/s]
Step 4: ARCO tracks output shape torch.Size([1, 9, 5, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.00it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.73it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.30it/s]
Step 5: ARCO tracks output shape torch.Size([1, 10, 6, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  7.11it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.55it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.63it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.69it/s]
Step 6: ARCO tracks output shape torch.Size([1, 11, 7, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  8.56it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.81it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.53it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.67it/s]
Step 7: ARCO tracks output shape torch.Size([1, 12, 8, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  5.37it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:03<00:00,  1.22it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:03<00:00,  1.58it/s]
Step 8: ARCO tracks output shape torch.Size([1, 13, 9, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  5.75it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:01<00:00,  1.99it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.70it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.84it/s]
Step 9: ARCO tracks output shape torch.Size([1, 14, 10, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  40%|████      | 2/5 [00:00<00:00, 11.88it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.63it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.66it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.85it/s]
Step 10: ARCO tracks output shape torch.Size([1, 15, 11, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  6.90it/s]
Fetching ARCO data:  60%|██████    | 3/5 [00:00<00:00,  9.07it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:01<00:00,  1.72it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.00it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.33it/s]
Step 11: ARCO tracks output shape torch.Size([1, 16, 12, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  6.51it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:01<00:00,  1.98it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:01<00:00,  2.56it/s]
Step 12: ARCO tracks output shape torch.Size([1, 17, 13, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  5.06it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:01<00:00,  2.05it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.79it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.92it/s]
Step 13: ARCO tracks output shape torch.Size([1, 18, 14, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  5.71it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.76it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.26it/s]
Step 14: ARCO tracks output shape torch.Size([1, 19, 15, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  5.49it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.74it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.15it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  2.12it/s]
Step 15: ARCO tracks output shape torch.Size([1, 20, 16, 4])

Fetching ARCO data:   0%|          | 0/5 [00:00<?, ?it/s]
Fetching ARCO data:  20%|██        | 1/5 [00:00<00:00,  7.70it/s]
Fetching ARCO data:  80%|████████  | 4/5 [00:02<00:00,  1.47it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.87it/s]
Fetching ARCO data: 100%|██████████| 5/5 [00:02<00:00,  1.84it/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")

Console output43 lines

Fetching ARCO data:   0%|          | 0/13 [00:00<?, ?it/s]
Fetching ARCO data:  31%|███       | 4/13 [00:00<00:00, 12.71it/s]
Fetching ARCO data:  85%|████████▍ | 11/13 [00:01<00:00,  7.16it/s]
Fetching ARCO data:  92%|█████████▏| 12/13 [00:02<00:00,  4.52it/s]
Fetching ARCO data: 100%|██████████| 13/13 [00:02<00:00,  4.62it/s]
Fetching ARCO data: 100%|██████████| 13/13 [00:02<00:00,  5.34it/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:05,  3.03it/s]Step 1: SFNO tracks output shape torch.Size([1, 8, 2, 4])

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Running inference: 100%|██████████| 17/17 [00:08<00:00,  2.77it/s]
Running inference: 100%|██████████| 17/17 [00:08<00:00,  2.12it/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.

---

Download Python sourceDownload Jupyter notebookexecuted in 157.9s

Execution profile

Runtime telemetry

Total runtime2m 38s

Execution environment

CPUAMD EPYC 7313P 16-Core Processor

GPUNVIDIA H100 PCIe · 79.6 GiB

System RAM58.5 GiB

PlatformLinux 6.8.0-136-generic

Python3.13.13

GPU driver / CUDADriver 595.84 · CUDA support 13.2