Skip to content

UCast

GlobalMRF202640 GBPyTorch

Import path: earth2studio.models.px.UCast

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

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:

Parameters:

  • model (Module) –

    U-CAST Dhariwal U-Net core.

  • center (Tensor) –

    Per-variable ERA5 means in VARIABLES order.

  • scale (Tensor) –

    Per-variable ERA5 standard deviations in VARIABLES order.

  • residual_scale (Tensor) –

    Per-variable residual standard deviations in VARIABLES order.

  • static_condition (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, default: True ) –

    Keep dropout layers active during inference, by default True.

  • preload_static_fields (bool, default: True ) –

    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__ 

__call__(
    x: Tensor, coords: CoordSystem
) -> tuple[Tensor, CoordSystem]

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

create_iterator 

create_iterator(
    x: Tensor, coords: CoordSystem
) -> Iterator[tuple[Tensor, CoordSystem]]

Creates an iterator for autoregressive U-CAST inference.

load_default_package classmethod 

load_default_package() -> Package

Load the default package for the U-CAST model.

load_model classmethod 

load_model(
    package: Package, preload_static_fields: bool = True
) -> PrognosticModel

Load U-CAST from a package.

Parameters:

  • package (Package) –

    Model package to load from.

  • preload_static_fields (bool, default: True ) –

    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:

  • PrognosticModel –

    Loaded U-CAST model.