用 NVIDIA Earth2Studio 构建批量集合天气预报工作流
Building Custom Batched Ensemble Weather Forecasting with NVIDIA Earth2Studio
做气象预报或地球科学 AI 的同学必看，这篇给出了从环境配置到集合扰动、验证与可视化的完整可照做流程，直接照着跑就能搭出自己的集合预报管道。
In this tutorial, we build an ensemble weather forecasting workflow with NVIDIA Earth2Studio. We install the required Earth2Studio components while preserving Colab’s existing CUDA-enabled PyTorch environment, load the FCN prognostic model, and retrieve atmospheric initial conditions from GFS. We then implement a custom wind-power diagnostic that converts 10-meter wind components into turbine capacity factors, along with a variable-scaled perturbation system that applies physically appropriate noise amplitudes to different atmospheric variables while retaining an unperturbed control member. Using Earth2Studio’s low-level iterator, coordinate-mapping, batching, and Zarr APIs, we construct our own ensemble execution pipeline, write forecast and diagnostic fields to a coordinate-aware data store, and verify the forecasts against GFS analyses using latitude-weighted RMSE, fair CRPS, ensemble spread, and spread-skill ratios. Finally, we visualize ensemble uncertainty through spatial maps, geopotential-height spaghetti contours, point-based fan charts, wind-capacity-factor forecasts, and lead-time skill curves.
在本教程中，我们使用NVIDIA Earth2Studio构建一个集合天气预报工作流。我们在保留Colab现有支持CUDA的PyTorch环境的同时，安装所需的Earth2Studio组件，加载FCN预报模型，并从GFS获取大气初始条件。然后，我们实现一个自定义的风力发电诊断模块，将10米风分量转换为涡轮机容量因子，以及一个变量缩放扰动系统，该系统对不同的大气变量施加物理上适当的噪声幅度，同时保留一个未扰动的控制成员。利用Earth2Studio的低级迭代器、坐标映射、批处理和Zarr API，我们构建自己的集合执行流水线，将预报和诊断字段写入坐标感知的数据存储，并使用纬度加权均方根误差、公平连续排名概率评分、集合离散度和离散度-技巧比率，对预报与GFS分析进行验证。最后，我们通过空间地图、位势高度意大利面条图、基于点的扇形图、风力容量因子预报和提前时间技巧曲线来可视化集合不确定性。
import importlib.util, os, subprocess, sys
if importlib.util.find_spec("earth2studio") is None:
import numpy as _np, torch as _torch
cfile = os.path.join(os.getcwd(), "e2s_constraints.txt")
with open(cfile, "w") as f:
f.write(f"torch=={_torch.__version__.split('+')[0]}\n")
f.write(f"numpy=={_np.__version__}\n")
env = {**os.environ, "PIP_CONSTRAINT": cfile}
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "-q",
"earth2studio[fcn,data,perturbation,statistics]"], env=env)
print("\n>>> Install done. If the imports below fail: Runtime > Restart session, re-run.\n")
os.environ.setdefault("EARTH2STUDIO_CACHE", "/content/e2s_cache")
os.makedirs("outputs", exist_ok=True)
from collections import OrderedDict
from datetime import datetime, timedelta, timezone
from tqdm.auto import tqdm
from earth2studio.data import GFS, fetch_data
from earth2studio.io import ZarrBackend
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.models.px import FCN
from earth2studio.statistics import rmse
from earth2studio.utils import handshake_coords, handshake_dim
from earth2studio.utils.coords import map_coords
from earth2studio.utils.time import to_time_array
from earth2studio.utils.type import CoordSystem
if DEVICE.type == "cpu":
print("!! No GPU detected — this will be very slow. Runtime > Change runtime type > T4 GPU")
NENSEMBLE = 8
BATCH_SIZE = 2
NSTEPS = 8
SAVE_VARS = ["t2m", "z500", "u10m", "v10m", "tcwv"]
VERIFY_VARS = ["t2m", "z500", "u10m"]
INIT = (datetime.now(timezone.utc) - timedelta(days=7)).replace()
INIT_STR = INIT.strftime("%Y-%m-%dT%H:%M:%S")
POI = ("New Delhi", 28.61, 77.21)
print(f"Initialization: {INIT_STR} | device: {DEVICE}")import importlib.util, os, subprocess, sys
if importlib.util.find_spec("earth2studio") is None:
import numpy as _np, torch as _torch
cfile = os.path.join(os.getcwd(), "e2s_constraints.txt")
with open(cfile, "w") as f:
f.write(f"torch=={_torch.__version__.split('+')[0]}\n")
f.write(f"numpy=={_np.__version__}\n")
env = {**os.environ, "PIP_CONSTRAINT": cfile}
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "-q",
"earth2studio[fcn,data,perturbation,statistics]"], env=env)
print("\n>>> Install done. If the imports below fail: Runtime > Restart session, re-run.\n")
os.environ.setdefault("EARTH2STUDIO_CACHE", "/content/e2s_cache")
os.makedirs("outputs", exist_ok=True)
from collections import OrderedDict
from datetime import datetime, timedelta, timezone
from tqdm.auto import tqdm
from earth2studio.data import GFS, fetch_data
from earth2studio.io import ZarrBackend
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.models.px import FCN
from earth2studio.statistics import rmse
from earth2studio.utils import handshake_coords, handshake_dim
from earth2studio.utils.coords import map_coords
from earth2studio.utils.time import to_time_array
from earth2studio.utils.type import CoordSystem
if DEVICE.type == "cpu":
print("!! No GPU detected — this will be very slow. Runtime > Change runtime type > T4 GPU")
NENSEMBLE = 8
BATCH_SIZE = 2
NSTEPS = 8
SAVE_VARS = ["t2m", "z500", "u10m", "v10m", "tcwv"]
VERIFY_VARS = ["t2m", "z500", "u10m"]
INIT = (datetime.now(timezone.utc) - timedelta(days=7)).replace()
INIT_STR = INIT.strftime("%Y-%m-%dT%H:%M:%S")
POI = ("New Delhi", 28.61, 77.21)
print(f"Initialization: {INIT_STR} | device: {DEVICE}")We install Earth2Studio while preserving Colab’s existing CUDA-enabled PyTorch and NumPy environment through package constraints. We configure the model cache, import the forecasting, data, statistics, plotting, and coordinate-management utilities, and detect the available compute device. We also define the ensemble size, batch size, forecast duration, saved variables, verification variables, initialization time, and New Delhi point of interest.
我们通过包约束安装Earth2Studio，同时保留Colab现有的支持CUDA的PyTorch和NumPy环境。我们配置模型缓存，导入预报、数据、统计、绘图和坐标管理工具，并检测可用的计算设备。我们还定义了集合大小、批大小、预报时长、保存变量、验证变量、初始化时间和新德里兴趣点。
class WindPowerCF(torch.nn.Module):
"""Turbine capacity factor [0,1] from 10 m winds via power-law shear + power curve."""
def __init__(self, lat, lon, hub=100.0, alpha=0.143,
cut_in=3.0, rated=12.0, cut_out=25.0):
super().__init__()
self.lat, self.lon = lat, lon
self.hub, self.alpha = hub, alpha
self.cut_in, self.rated, self.cut_out = cut_in, rated, cut_out
def input_coords(self) -> CoordSystem:
return OrderedDict({
"batch": np.empty(0),
"variable": np.array(["u10m", "v10m"]),
"lat": self.lat,
"lon": self.lon,
})
@batch_coords()
def output_coords(self, input_coords: CoordSystem) -> CoordSystem:
target = self.input_coords()
for i, (key, _) in enumerate(target.items()):
if key != "batch":
handshake_dim(input_coords, key, i)
handshake_coords(input_coords, target, key)
oc = OrderedDict({
"batch": np.empty(0),
"variable": np.array(["wind_cf"]),
"lat": self.lat,
"lon": self.lon,
})
oc["batch"] = input_coords["batch"]
return oc
@batch_func()
def __call__(self, x: torch.Tensor, coords: CoordSystem):
oc = self.output_coords(coords)
u, v = x[..., 0:1, :, :], x[..., 1:2, :, :]
ws10 = torch.sqrt(u * u + v * v)
ws = ws10 * (self.hub / 10.0) ** self.alpha
ramp = (ws ** 3 - self.cut_in ** 3) / (self.rated ** 3 - self.cut_in ** 3)
cf = torch.zeros_like(ws)
cf = torch.where((ws >= self.cut_in) & (ws < self.rated), ramp.clamp(0, 1), cf)
cf = torch.where((ws >= self.rated) & (ws <= self.cut_out), torch.ones_like(cf), cf)
return cf, oc
class VariableScaledNoise:
"""Spatially correlated noise with per-variable amplitudes + control member."""
def __init__(self, amplitudes: dict, default: float = 0.0, control_member: bool = True):
self.amplitudes, self.default, self.control = amplitudes, default, control_member
try:
from earth2studio.perturbation import SphericalGaussian
self.sampler, self.kind = SphericalGaussian(noise_amplitude=1.0), "SphericalGaussian"
except Exception:
from earth2studio.perturbation import Brown
self.sampler, self.kind = Brown(noise_amplitude=1.0), "Brown"
def __call__(self, x: torch.Tensor, coords: CoordSystem):
noise, _ = self.sampler(torch.zeros_like(x), coords)
vax = list(coords).index("variable")
amps = torch.tensor([self.amplitudes.get(str(v), self.default)
for v in coords["variable"]], device=x.device, dtype=x.dtype)
shape = [1] * x.ndim; shape[vax] = amps.numel()
pert = noise * amps.reshape(shape)
if self.control and "ensemble" in coords:
eax = list(coords).index("ensemble")
mask = torch.tensor((np.asarray(coords["ensemble"]) != 0).astype(np.float32),
device=x.device, dtype=x.dtype)
mshape = [1] * x.ndim; mshape[eax] = mask.numel()
pert = pert * mask.reshape(mshape)
return x + pert, coordsclass WindPowerCF(torch.nn.Module):
"""Turbine capacity factor [0,1] from 10 m winds via power-law shear + power curve."""
def __init__(self, lat, lon, hub=100.0, alpha=0.143,
cut_in=3.0, rated=12.0, cut_out=25.0):
super().__init__()
self.lat, self.lon = lat, lon
self.hub, self.alpha = hub, alpha
self.cut_in, self.rated, self.cut_out = cut_in, rated, cut_out
def input_coords(self) -> CoordSystem:
return OrderedDict({
"batch": np.empty(0),
"variable": np.array(["u10m", "v10m"]),
"lat": self.lat,
"lon": self.lon,
})
@batch_coords()
def output_coords(self, input_coords: CoordSystem) -> CoordSystem:
target = self.input_coords()
for i, (key, _) in enumerate(target.items()):
if key != "batch":
handshake_dim(input_coords, key, i)
handshake_coords(input_coords, target, key)
oc = OrderedDict({
"batch": np.empty(0),
"variable": np.array(["wind_cf"]),
"lat": self.lat,
"lon": self.lon,
})
oc["batch"] = input_coords["batch"]
return oc
@batch_func()
def __call__(self, x: torch.Tensor, coords: CoordSystem):
oc = self.output_coords(coords)
u, v = x[..., 0:1, :, :], x[..., 1:2, :, :]
ws10 = torch.sqrt(u * u + v * v)
ws = ws10 * (self.hub / 10.0) ** self.alpha
ramp = (ws ** 3 - self.cut_in ** 3) / (self.rated ** 3 - self.cut_in ** 3)
cf = torch.zeros_like(ws)
cf = torch.where((ws >= self.cut_in) & (ws < self.rated), ramp.clamp(0, 1), cf)
cf = torch.where((ws >= self.rated) & (ws <= self.cut_out), torch.ones_like(cf), cf)
return cf, oc
class VariableScaledNoise:
"""Spatially correlated noise with per-variable amplitudes + control member."""
def __init__(self, amplitudes: dict, default: float = 0.0, control_member: bool = True):
self.amplitudes, self.default, self.control = amplitudes, default, control_member
try:
from earth2studio.perturbation import SphericalGaussian
self.sampler, self.kind = SphericalGaussian(noise_amplitude=1.0), "SphericalGaussian"
except Exception:
from earth2studio.perturbation import Brown
self.sampler, self.kind = Brown(noise_amplitude=1.0), "Brown"
def __call__(self, x: torch.Tensor, coords: CoordSystem):
noise, _ = self.sampler(torch.zeros_like(x), coords)
vax = list(coords).index("variable")
amps = torch.tensor([self.amplitudes.get(str(v), self.default)
for v in coords["variable"]], device=x.device, dtype=x.dtype)
shape = [1] * x.ndim; shape[vax] = amps.numel()
pert = noise * amps.reshape(shape)
if self.control and "ensemble" in coords:
eax = list(coords).index("ensemble")
mask = torch.tensor((np.asarray(coords["ensemble"]) != 0).astype(np.float32),
device=x.device, dtype=x.dtype)
mshape = [1] * x.ndim; mshape[eax] = mask.numel()
pert = pert * mask.reshape(mshape)
return x + pert, coordsWe create a custom diagnostic model that converts 10-meter wind components into hub-height wind speed and turbine capacity factor. We validate coordinate compatibility through Earth2Studio’s handshake utilities and support batched inputs with the provided decorators. We also implement variable-specific spatial perturbations that retain member zero as an unperturbed control forecast.
我们创建一个自定义诊断模型，将10米风分量转换为轮毂高度风速和涡轮机容量因子。我们通过Earth2Studio的握手工具验证坐标兼容性，并使用提供的装饰器支持批处理输入。我们还实现特定变量的空间扰动，保留第零个成员作为未扰动的控制预报。
更进一步：量化金融体系
看懂新闻只是起点——沿量化金融路径，把它变成能交付的工程能力