Generative Downscaling¶

Generative downscaling over Taiwan using CorrDiff diffusion model.

This example will demonstrate how to user Nvidia's CorrDiff model, trained for predicting weather over Taiwan, to perform generative downscaling from quarter degree global forecast data to ~3km.

This checkpoint was trained on ERA5 data and WRF data that spans 2018-2021 at one hour time resolution. In this example, we demonstrate an application to GFS data for a typhoon super-resolution from 2023. The model's performance on GFS data and on data from this year has not been evaluated.

In this example you will learn:

Creating a custom workflow for running CorrDiff inference
Creating a data-source for CorrDiff's input
Initializing and running CorrDiff diagnostic model
Post-processing results.

Creating a Simple CorrDiff Workflow¶

As usual, we start with creating a simple workflow to run CorrDiff in. To maximize the generalization of this workflow, we use dependency injection following the pattern provided inside earth2studio.run. Since CorrDiff is a diagnostic model, this workflow won't predict a time-series, rather just an instantaneous prediction.

For this workflow, we specify

time: Input list of datetimes / strings to run inference for
corrdiff: The initialized CorrDiffTaiwan model
data: Initialized data source to fetch initial conditions from
io: IOBackend
number_of_samples: Number of samples to generate from the model

import os

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

load_dotenv()  # TODO: make common example prep function

from collections import OrderedDict
from datetime import datetime

import numpy as np
import torch
from loguru import logger

from earth2studio.data import DataSource, prep_data_array
from earth2studio.io import IOBackend
from earth2studio.models.dx import CorrDiffTaiwan
from earth2studio.utils.coords import map_coords, split_coords
from earth2studio.utils.time import to_time_array


def run(
    time: list[str] | list[datetime] | list[np.datetime64],
    corrdiff: CorrDiffTaiwan,
    data: DataSource,
    io: IOBackend,
    number_of_samples: int = 1,
) -> IOBackend:
    """CorrDiff infernce workflow

    Parameters
    ----------
    time : list[str] | list[datetime] | list[np.datetime64]
        List of string, datetimes or np.datetime64
    corrdiff : CorrDiffTaiwan
        CorrDiff mode
    data : DataSource
        Data source
    io : IOBackend
        IO object
    number_of_samples : int, optional
        Number of samples to generate, by default 1

    Returns
    -------
    IOBackend
        Output IO object
    """
    logger.info("Running corrdiff inference!")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Inference device: {device}")

    corrdiff = corrdiff.to(device)
    # Update the number of samples for corrdiff to generate
    corrdiff.number_of_samples = number_of_samples

    # Fetch data from data source and load onto device
    time = to_time_array(time)
    x, coords = prep_data_array(
        data(time, corrdiff.input_coords()["variable"]), device=device
    )
    x, coords = map_coords(x, coords, corrdiff.input_coords())

    logger.success(f"Fetched data from {data.__class__.__name__}")

    # Set up IO backend
    output_coords = corrdiff.output_coords(corrdiff.input_coords())
    total_coords = OrderedDict(
        {
            "time": coords["time"],
            "sample": output_coords["sample"],
            "lat": output_coords["lat"],
            "lon": output_coords["lon"],
        }
    )
    io.add_array(total_coords, output_coords["variable"])

    logger.info("Inference starting!")
    x, coords = corrdiff(x, coords)
    io.write(*split_coords(x, coords))

    logger.success("Inference complete")
    return io

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

Set Up¶

With the workflow defined, the next step is initializing the needed components from Earth-2 studio

It's clear we need the following:

Diagnostic Model: CorrDiff model for Taiwan earth2studio.models.dx.CorrDiffTaiwan.
Datasource: Pull data from the GFS data api earth2studio.data.GFS.
IO Backend: Save the outputs into a Zarr store earth2studio.io.ZarrBackend.

from earth2studio.data import GFS
from earth2studio.io import ZarrBackend

# Create CorrDiff model
package = CorrDiffTaiwan.load_default_package()
corrdiff = CorrDiffTaiwan.load_model(package)

# Create the data source
data = GFS()

# Create the IO handler, store in memory
io = ZarrBackend()

Console output26 lines

Downloading corrdiff_inference_package.zip: 0%|          | 0.00/684M [00:00<?, ?B/s]
Downloading corrdiff_inference_package.zip: 1%|          | 5.84M/684M [00:00<00:11, 60.7MB/s]
Downloading corrdiff_inference_package.zip: 2%|▏         | 11.7M/684M [00:00<00:12, 55.4MB/s]
Downloading corrdiff_inference_package.zip: 7%|▋         | 50.0M/684M [00:00<00:03, 205MB/s] 
Downloading corrdiff_inference_package.zip: 12%|█▏        | 79.6M/684M [00:00<00:02, 245MB/s]
Downloading corrdiff_inference_package.zip: 15%|█▌        | 104M/684M [00:00<00:02, 232MB/s] 
Downloading corrdiff_inference_package.zip: 19%|█▉        | 128M/684M [00:00<00:02, 240MB/s]
Downloading corrdiff_inference_package.zip: 24%|██▎       | 161M/684M [00:00<00:02, 271MB/s]
Downloading corrdiff_inference_package.zip: 27%|██▋       | 187M/684M [00:00<00:02, 253MB/s]
Downloading corrdiff_inference_package.zip: 32%|███▏      | 216M/684M [00:00<00:01, 267MB/s]
Downloading corrdiff_inference_package.zip: 35%|███▌      | 242M/684M [00:01<00:01, 266MB/s]
Downloading corrdiff_inference_package.zip: 40%|███▉      | 270M/684M [00:01<00:01, 276MB/s]
Downloading corrdiff_inference_package.zip: 43%|████▎     | 297M/684M [00:01<00:01, 268MB/s]
Downloading corrdiff_inference_package.zip: 48%|████▊     | 326M/684M [00:01<00:01, 280MB/s]
Downloading corrdiff_inference_package.zip: 52%|█████▏    | 358M/684M [00:01<00:01, 294MB/s]
Downloading corrdiff_inference_package.zip: 57%|█████▋    | 388M/684M [00:01<00:01, 300MB/s]
Downloading corrdiff_inference_package.zip: 61%|██████    | 417M/684M [00:01<00:00, 302MB/s]
Downloading corrdiff_inference_package.zip: 65%|██████▌   | 446M/684M [00:01<00:00, 288MB/s]
Downloading corrdiff_inference_package.zip: 70%|██████▉   | 477M/684M [00:01<00:00, 299MB/s]
Downloading corrdiff_inference_package.zip: 74%|███████▍  | 506M/684M [00:02<00:00, 289MB/s]
Downloading corrdiff_inference_package.zip: 79%|███████▉  | 541M/684M [00:02<00:00, 311MB/s]
Downloading corrdiff_inference_package.zip: 84%|████████▍ | 575M/684M [00:02<00:00, 325MB/s]
Downloading corrdiff_inference_package.zip: 89%|████████▉ | 609M/684M [00:02<00:00, 333MB/s]
Downloading corrdiff_inference_package.zip: 94%|█████████▍| 645M/684M [00:02<00:00, 347MB/s]
Downloading corrdiff_inference_package.zip: 99%|█████████▉| 678M/684M [00:02<00:00, 327MB/s]
Downloading corrdiff_inference_package.zip: 100%|██████████| 684M/684M [00:02<00:00, 281MB/s]

Execute the Workflow¶

With all components initialized, running the workflow is a single line of Python code. Workflow will return the provided IO object back to the user, which can be used to then post process. Some have additional APIs that can be handy for post-processing or saving to file. Check the API docs for more information.

For the inference we will predict 1 sample for a particular timestamp representing Typhoon Koinu.

io = run(["2023-10-04T18:00:00"], corrdiff, data, io, number_of_samples=1)

Console output12 lines

2026-08-15 05:08:51.702 | INFO     | __main__:run:49 - Running corrdiff inference!
2026-08-15 05:08:51.702 | INFO     | __main__:run:51 - Inference device: cuda

Fetching GFS data:   0%|          | 0/12 [00:00<?, ?it/s]
Fetching GFS data:   8%|▊         | 1/12 [00:00<00:04,  2.33it/s]
Fetching GFS data:  17%|█▋        | 2/12 [00:00<00:02,  3.79it/s]
Fetching GFS data:  42%|████▏     | 5/12 [00:00<00:00,  8.78it/s]
Fetching GFS data: 100%|██████████| 12/12 [00:00<00:00, 21.96it/s]
Fetching GFS data: 100%|██████████| 12/12 [00:00<00:00, 13.97it/s]
2026-08-15 05:08:53.269 | SUCCESS  | __main__:run:64 - Fetched data from GFS
2026-08-15 05:08:53.302 | INFO     | __main__:run:78 - Inference starting!
2026-08-15 05:08:55.207 | SUCCESS  | __main__:run:82 - Inference complete

Post Processing¶

The last step is to post process our results. Cartopy is a great library for plotting fields on projections of a sphere.

Notice that the Zarr IO function has additional APIs to interact with the stored data.

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

projection = ccrs.LambertConformal(
    central_longitude=io["lon"][:].mean(),
)

fig = plt.figure(figsize=(4 * 8, 8))

ax0 = fig.add_subplot(1, 3, 1, projection=projection)
c = ax0.pcolormesh(
    io["lon"],
    io["lat"],
    io["mrr"][0, 0],
    transform=ccrs.PlateCarree(),
    cmap="inferno",
)
plt.colorbar(c, ax=ax0, shrink=0.6, label="mrr dBz")
ax0.coastlines()
ax0.gridlines()
ax0.set_title("Radar Reflectivity")

ax1 = fig.add_subplot(1, 3, 2, projection=projection)
c = ax1.pcolormesh(
    io["lon"],
    io["lat"],
    io["t2m"][0, 0],
    transform=ccrs.PlateCarree(),
    cmap="RdBu_r",
)
plt.colorbar(c, ax=ax1, shrink=0.6, label="K")
ax1.coastlines()
ax1.gridlines()
ax1.set_title("2-meter Temperature")

ax2 = fig.add_subplot(1, 3, 3, projection=projection)
c = ax2.pcolormesh(
    io["lon"],
    io["lat"],
    np.sqrt(io["u10m"][0, 0] ** 2 + io["v10m"][0, 0] ** 2),
    transform=ccrs.PlateCarree(),
    cmap="Greens",
)
plt.colorbar(c, ax=ax2, shrink=0.6, label="w10m m s^-1")
ax2.coastlines()
ax2.gridlines()
ax2.set_title("10-meter Wind Speed")

plt.savefig("outputs/04_corr_diff_prediction.jpg")

Console output2 lines

/__w/earth2studio/earth2studio/.venv/lib/python3.13/site-packages/cartopy/io/__init__.py:242: DownloadWarning: Downloading: https://naturalearth.s3.amazonaws.com/10m_physical/ne_10m_coastline.zip
  warnings.warn(f'Downloading: {url}', DownloadWarning)

Output from Generative Downscaling

Download Python source Download Jupyter notebookexecuted in 39.2s

Execution profile

Runtime telemetry

Total runtime39.2 s

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