UCast

class earth2studio.models.px.UCast(

model,

center,

scale,

residual_scale,

static_condition,

sst_fill_value,

stochastic=True,

preload_static_fields=True,

)

GlobalMRF202640 GB

U-CAST 1.5 degree global probabilistic weather model.

U-CAST is a 12-hour autoregressive U-Net forecaster trained on WeatherBench2 ERA5. The model uses two history steps, 83 dynamic variables, four time/solar progress forcing channels, and two static fields. Ensemble members can be generated by adding an ensemble dimension ahead of time; dropout remains active by default during inference so members diverge autoregressively.

Note

For additional information see the following resources:

• https://arxiv.org/abs/2604.09041

• Rose-STL-Lab/u-cast

• https://huggingface.co/salv47/u-cast

Parameters:

• model (torch.nn.Module) – U-CAST Dhariwal U-Net core.

• center (torch.Tensor) – Per-variable ERA5 means in VARIABLES order.

• scale (torch.Tensor) – Per-variable ERA5 standard deviations in VARIABLES order.

• residual_scale (torch.Tensor) – Per-variable residual standard deviations in VARIABLES order.

• static_condition (torch.Tensor) – Static fields in (field, lat, lon) order. Empty if preload_static_fields=False.

• sst_fill_value (float) – Fill value used for SST NaNs over land before normalization.

• stochastic (bool) – Keep dropout layers active during inference, by default True.

• preload_static_fields (bool) – If True, static fields are fetched from WeatherBench2 at load time and cached. If False, these fields must be provided as input variables (lsm and z), allowing use of static fields from alternative sources for exact reproducibility, by default True.

__call__(x, coords)

Runs the 12-hour U-CAST prognostic model one step.

Parameters:

• x (Tensor)

• coords (OrderedDict[str, ndarray])

Return type:

tuple[Tensor, OrderedDict[str, ndarray]]

create_iterator(x, coords)

Creates an iterator for autoregressive U-CAST inference.

Parameters:

• x (Tensor)

• coords (OrderedDict[str, ndarray])

Return type:

Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]

classmethod load_default_package()

Load the default package for the U-CAST model.

Return type:

Package

classmethod load_model(package, preload_static_fields=True)

Load U-CAST from a package.

Parameters:

• package (Package) – Model package to load from.

• preload_static_fields (bool, optional) – If True, static fields (lsm and z) are fetched from WeatherBench2 at load time and cached. If False, these fields must be provided as input variables, allowing use of static fields from alternative sources for exact reproducibility with reference implementations, by default True.

Returns:

Loaded U-CAST model.

Return type:

PrognosticModel