使用 NVIDIA Earth2Studio 建置自訂批次系集天氣預報

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.Copy CodeCopiedUse a different Browserimport importlib.util, os, subprocess, sys if importlib.util.

findspec("earth2studio") is None: import numpy as np, torch as torch cfile = os.path.join(os.getcwd(), "e2sconstraints.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, "PIPCONSTRAINT": cfile} subprocess.

checkcall( [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("EARTH2STUDIOCACHE", "/content/e2scache") os.

makedirs("outputs", existok=True) from collections import OrderedDict from datetime import datetime, timedelta, timezone from tqdm.auto import tqdm from earth2studio.data import GFS, fetchdata from earth2studio.io import ZarrBackend from earth2studio.models.

batch import batchcoords, batchfunc from earth2studio.models.px import FCN from earth2studio.statistics import rmse from earth2studio.utils import handshakecoords, handshakedim from earth2studio.utils.coords import mapcoords from earth2studio.utils.time import totimearray 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 BATCHSIZE = 2 NSTEPS = 8 SAVEVARS = ["t2m", "z500", "u10m", "v10m", "tcwv"] VERIFYVARS = ["t2m", "z500", "u10m"] INIT = (datetime.now(timezone.

utc) - timedelta(days=7)).replace() INITSTR = INIT.strftime("%Y-%m-%dT%H:%M:%S") POI = ("New Delhi", 28.61, 77.21) print(f"Initialization: {INITSTR} | 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.Copy CodeCopiedUse a different Browserclass 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, cutin=3.0, rated=12.0, cutout=25.0): super().__init__() self.lat, self.lon = lat, lon self.hub, self.alpha = hub, alpha self.cutin, self.rated, self.

cutout = cutin, rated, cutout def inputcoords(self) -> CoordSystem: return OrderedDict({ "batch": np.empty(0), "variable": np.array(["u10m", "v10m"]), "lat": self.lat, "lon": self.lon, }) @batchcoords() def outputcoords(self, inputcoords: CoordSystem) -> CoordSystem: target = self.

inputcoords() for i, (key, ) in enumerate(target.items()): if key != "batch": handshakedim(inputcoords, key, i) handshakecoords(inputcoords, target, key) oc = OrderedDict({ "batch": np.empty(0), "variable": np.array(["windcf"]), "lat": self.lat, "lon": self.

lon, }) oc["batch"] = inputcoords["batch"] return oc @batch_func() def __call__(self, x: torch.Tensor, coords: CoordSystem): oc = self.outputcoords(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.

cutin 3) / (self.rated 3 - self.cutin 3) cf = torch.zeroslike(ws) cf = torch.where((ws >= self.cutin) & (ws < self.rated), ramp.clamp(0, 1), cf) cf = torch.where((ws >= self.rated) & (ws <= self.cutout), torch.

oneslike(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, controlmember: bool = True): self.amplitudes, self.default, self.

control = amplitudes, default, controlmember try: from earth2studio.perturbation import SphericalGaussian self.sampler, self.kind = SphericalGaussian(noiseamplitude=1.0), "SphericalGaussian" except Exception: from earth2studio.perturbation import Brown self.sampler, self.

kind = Brown(noiseamplitude=1.0), "Brown" def call(self, x: torch.Tensor, coords: CoordSystem): noise, = self.sampler(torch.zeroslike(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, coords We 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.

Copy CodeCopiedUse a different Browserdef writevars(io, x, coords, names): """Write selected channels of a (…, variable, lat, lon) tensor to the IO backend.""" vax = list(coords).index("variable") sub = OrderedDict((k, v) for k, v in coords.items() if k != "variable") for name in names: hit = np.

where(np.asarray(coords["variable"]) == name)[0] if hit.size: io.write(x.select(vax, int(hit[0])).cpu(), sub, name) def runensemble(time, nsteps, nensemble, batchsize, prognostic, diagnostic, perturbation, data, io, savevars, device): time = totimearray(time) ic = prognostic.

inputcoords() x0, c0 = fetchdata(source=data, time=time, leadtime=ic["leadtime"], variable=ic["variable"], device=device) print(f"Initial condition tensor: {tuple(x0.shape)} dims={list(c0)}") oc = prognostic.

outputcoords(ic) dt = oc["leadtime"] progvars = [v for v in savevars if v in set(map(str, oc["variable"]))] total = OrderedDict({ "ensemble": np.arange(nensemble), "time": time, "leadtime": np.asarray([dt i for i in range(nsteps + 1)]).flatten(), "lat": oc["lat"], "lon": oc["lon"], }) io.

addarray(total, progvars + ["windcf"]) dxtarget = OrderedDict((k, v) for k, v in diagnostic.inputcoords().items() if k != "batch") nbatch = int(np.ceil(nensemble / batchsize)) with torch.

inferencemode(): for b in tqdm(range(nbatch), desc="ensemble batches"): lo = b batchsize n = min(batchsize, nensemble - lo) x = x0.unsqueeze(0).repeat(n, ([1] * x0.ndim)) coords = OrderedDict({"ensemble": np.

arange(lo, lo + n), **c0}) x, coords = perturbation(x, coords) x, coords = mapcoords(x, coords, ic) for step, (xs, cs) in enumerate(progn