⌘KCtrl+K
Generative downscaling model for generating high resolution regional scale weather fields.
Follow the steps below to download and run the NVIDIA NIM inference microservice for this model on your infrastructure of choice.
Pull and run the NVIDIA Earth-2 CorrDiff NIM with the command below.
docker pull nvcr.io/nim/nvidia/corrdiff:latest
This will download the optimized model for your infrastructure.
export NGC_API_KEY=<NGC API Key>
docker run --rm --runtime=nvidia --gpus all --shm-size 4g \
-p 8000:8000 \
-e NGC_API_KEY \
nvcr.io/nim/nvidia/corrdiff:latest
Check the health of the NIM with the following curl command:
curl -X 'GET' \
'http://localhost:8000/v1/health/ready' \
-H 'accept: application/json'
Generate an input numpy array for the model using the following Python script with Earth2Studio. See the NVIDIA NIM Docs quick start guide for more details.
from datetime import datetime, timedelta
import numpy as np
import torch
from earth2studio.data import GEFS_FX, GEFS_FX_721x1440
GEFS_SELECT_VARIABLES = ["u10m","v10m","t2m","r2m","sp","msl","tcwv"]
GEFS_VARIABLES = ["u1000","u925","u850","u700","u500","u250","v1000","v925","v850",\
"v700","v500","v250","z1000","z925","z850","z700","z500","z200","t1000","t925",\
"t850","t700","t500","t100","r1000","r925","r850","r700","r500","r100"]
ds_gefs = GEFS_FX(cache=True)
ds_gefs_select = GEFS_FX_721x1440(cache=True, product="gec00")
def fetch_input_gefs(
time: datetime, lead_time: timedelta, content_dtype: str = "float32"
):
dtype = np.dtype(getattr(np, content_dtype))
# Fetch high-res select GEFS input data
select_data = ds_gefs_select(time, lead_time, GEFS_SELECT_VARIABLES)
select_data = select_data.values
# Crop to bounding box [225, 21, 300, 53]
select_data = select_data[:, 0, :, 148:277, 900:1201].astype(dtype)
assert select_data.shape == (1, len(GEFS_SELECT_VARIABLES), 129, 301)
# Fetch GEFS input data
pressure_data = ds_gefs(time, lead_time, GEFS_VARIABLES)
# Interpolate to 0.25 grid
pressure_data = torch.nn.functional.interpolate(
torch.Tensor(pressure_data.values),
(len(GEFS_VARIABLES), 721, 1440),
mode="nearest",
)
pressure_data = pressure_data.numpy()
# Crop to bounding box [225, 21, 300, 53]
pressure_data = pressure_data[:, 0, :, 148:277, 900:1201].astype(dtype)
assert pressure_data.shape == (1, len(GEFS_VARIABLES), 129, 301)
# Create lead time field
lead_hour = int(lead_time.total_seconds() // (3 * 60 * 60)) * np.ones(
(1, 1, 129, 301)
).astype(dtype)
input_data = np.concatenate([select_data, pressure_data, lead_hour], axis=1)[None]
return input_data
input_array = fetch_input_gefs(datetime(2023, 1, 1), timedelta(hours=0))
np.save("corrdiff_inputs.npy", input_array)
You can now make a local API call using this curl command:
curl -X POST \
-F "input_array=@corrdiff_inputs.npy" \
-F "samples=2" \
-F "steps=12" \
-o output.tar \
http://localhost:8000/v1/infer
For more details on getting started with this NIM, visit the NVIDIA NIM Docs.
For more details on the model and its input / output tensors see the US CorrDiff Model Card.