借助 NVIDIA Earth-2，将您的最新观测结果转化为及时的天气决策

from datetime import datetime
from earth2studio.data import GHCNHourly
from earth2studio.models.da import CorrDiffCosmoEra5SDA

domain = dict(lat_min=50.2, lat_max=53.8, lon_min=4.6, lon_max=10.4)
sda = CorrDiffCosmoEra5SDA.load_model(
    CorrDiffCosmoEra5SDA.load_default_package(),
    assimilate_variables=("u10m", "v10m"),
    resolution="rea2",
    domain=domain,
    number_of_samples=1,
    sampler_steps=12,
    amp=True,          
).to("cuda")

# Fetch and regrid ERA5 inputs onto the high resolution regional grid
# Follow the link to the example below for the full implementation
init_time = datetime(2024, 1, 26)
x = fetch_and_regrid_era5(init_time, domain)

# Fetch GHCN hourly 10-m wind observations over the model domain
lat, lon = sda.model.lat_output_numpy, sda.model.lon_output_numpy
bbox = lat.min(), lon.min(), lat.max(), lon.max()
ghcn = GHCNHourly(stations=GHCNHourly.get_stations_bbox(bbox))
obs = ghcn(init_time, ["u10m", "v10m"]).dropna(subset=["observation"])

prior = sda(x)  # free downscaling, no observations
analysis = sda(x, obs)  # guide the diffusion toward the observations

import numpy as np
from earth2studio.data import GHCNHourly
from earth2studio.models.px import StormCastCONUS

# Limit the domain to the central U.S.
hrrr_lat_lim, hrrr_lon_lim  = (305, 785), (595, 1203)
model = StormCastCONUS.load_model(
    StormCastCONUS.load_default_package(),
    hrrr_lat_lim=hrrr_lat_lim,  # comment out for full CONUS domain
    hrrr_lon_lim=hrrr_lon_lim,  # comment out for full CONUS domain
    num_diffusion_steps=18,
    num_sda_diffusion_steps=96,  # more steps for SDA for better stability
    sda_std_obs=0.15,
    sda_gamma=1e-3,
).to("cuda")

# Fetch HRRR initial conditions
# Follow the link to the example below for the full implementation
init_time = datetime(2026, 4, 17, 18)
x, coords = fetch_hrrr(init_time)

# Define GHCN hourly data source for the model domain
lat, lon = model.lat, model.lon
bbox = lat.min(), lon.min(), lat.max(), lon.max()
ghcn = GHCNHourly(
    stations=GHCNHourly.get_stations_bbox(bbox),
    time_tolerance=timedelta(minutes=15),
)

# Initialize generator and get the first output (analysis passthrough)
gen = model.create_generator(x.clone(), coords.copy())
x, coords = next(gen)

# Run the first part of the rollout with SDA
for step in range(nsteps_sda):
    valid_time = np.array(
        [coords["time"][0] + coords["lead_time"][0] + np.timedelta64(1, "h")]
    )
    obs = ghcn(valid_time, ["u10m", "v10m", "t2m"])
    x, coords = gen.send(obs)  # advance one step with observations

# Run the remaining rollout without SDA
for step in range(nsteps_non_sda):
    x, coords = next(gen)  # advance one step without observations

from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler
from physicsnemo.diffusion.samplers import sample

# Construct diffusion scheduler
sigma_min, sigma_max = 0.01, 100
scheduler = EDMNoiseScheduler(sigma_min=sigma_min, sigma_max=sigma_max)

# Get denoiser from scheduler
denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor)

# Generate sample
latents = sigma_max * torch.randn(shape)
sample(denoiser, latents, noise_scheduler=scheduler, num_steps=num_steps)

from physicsnemo.diffusion.guidance import (
    DataConsistencyDPSGuidance,
    DPSScorePredictor,
)

# Setup SDA guidance
guidance = DataConsistencyDPSGuidance(
    mask=mask,  # binary mask that identifies pixels with observations
    y=y_obs,  # gridded observations
    std_y=sda_std_obs,  # the remaining parameters are SDA settings
    norm=sda_dps_norm,
    gamma=sda_gamma,
    sigma_fn=scheduler.sigma,
    alpha_fn=scheduler.alpha,
)

# Convert x0_predictor to score predictor
score_predictor = DPSScorePredictor( 
    x0_predictor=x0_predictor,
    x0_to_score_fn=scheduler.x0_to_score,
    guidances=guidance,
)
denoiser = scheduler.get_denoiser(score_predictor=score_predictor)

# Generate sample (identical to non-SDA example)
latents = sigma_max * torch.randn(shape)
sample(denoiser, latents, noise_scheduler=scheduler, num_steps=num_steps)

from datetime import timedelta

import numpy as np
from earth2studio.data import UFSObsConv, UFSObsSat, fetch_dataframe
from earth2studio.models.da import HealDA

model = HealDA.load_model(
    HealDA.load_default_package(),
    lat_lon=True,  # regrid from HEALPix to regular lat/lon
).to("cuda")

# HealDA was trained on the UFS replay window: 21h before to 3h after analysis time
time_tolerance = (timedelta(hours=-21), timedelta(hours=3))
analysis_time = np.array([np.datetime64("2024-01-01T00:00")])

# input_coords() returns the schemas the two observation DataFrames must satisfy
conv_schema, sat_schema = model.input_coords()

# fetch_dataframe attaches the request_time metadata the model needs
conv_df = fetch_dataframe(
    UFSObsConv(time_tolerance=time_tolerance),
    time=analysis_time,
    variable=np.array(conv_schema["variable"]),
    fields=np.array(list(conv_schema.keys())),
)
sat_df = fetch_dataframe(
    UFSObsSat(time_tolerance=time_tolerance),
    time=analysis_time,
    variable=np.array(sat_schema["variable"]),
    fields=np.array(list(sat_schema.keys())),
)

# stateless model - call it directly for a one-shot analysis, or use
# create_generator for cycled assimilation
analysis = model(conv_obs=conv_df, sat_obs=sat_df)

era5 = NCAR_ERA5()
da_era5 = era5(datetime(2025, 7, 15), ["t2m", "z500"])
print(da_era5.shape)  # (1, 2, 721, 1440)

ifs = IFS_FX()
da_ifs = ifs(datetime(2026, 7, 15), timedelta(hours=48), ["t2m"])
print(da_ifs.shape)  # (1, 1, 1, 721, 1440)

goes = GOES(satellite="goes19", scan_mode="C")
da_goes = goes(datetime(2026, 7, 15), ["abi01c", "abi02c", "abi03c"])
print(da_goes.shape)  # (1, 3, 1500, 2500)

ghcn = GHCNHourly(stations=["USW00013301"])
df_ghcn = ghcn(datetime(2026, 6, 15), ["t2m", "ws10m"])
print(df_ghcn.shape)  # (10, 7)