Skip to content

WeatherNext2CyclonesMini

GlobalMRF202640 GBGoogleJAX

Import path: earth2studio.models.px.WeatherNext2CyclonesMini

View source on GitHub View install commands

Documentation

Bases: _WeatherNext2Base

WeatherNext 2 Cyclones Mini medium-range forecast model.

WeatherNext 2 is Google DeepMind's global medium-range weather forecasting model family. This wrapper uses the public WeatherNextCyclones_Mini checkpoint and 1 degree sample grid.

The model requires two input states, valid at -6h and 0h lead time, and predicts 6 hours forward per model call. By default this wrapper returns only the gridded weather fields expected by Earth2Studio prognostic models. Cyclone tracking can be enabled with track_cyclones=True to accumulate WeatherNext's tropical cyclone track diagnostics in the cyclone_tracks property without changing the model output type.

Note

For more information see the following references:

Warning

The cyclone-tracking interface, including track_cyclones and cyclone_tracks, is in beta and may change in future releases as Earth2Studio standardizes how track diagnostics are exposed.

We encourage users to familiarize themselves with the license restrictions of this model's checkpoints.

Parameters:

  • ckpt (CheckPoint) –

    Model checkpoint containing weights.

  • land_sea_mask (ndarray) –

    Land-sea mask on the WeatherNext grid.

  • geopotential_at_surface (ndarray) –

    Surface geopotential on the WeatherNext grid.

  • seed (int, default: 0 ) –

    Initial random seed for the stochastic FGN noise generator, by default 0.

  • jit_compile (bool, default: True ) –

    JIT-compile the model forward pass, by default True.

  • track_cyclones (bool, default: False ) –

    Accumulate tropical cyclone tracks in the cyclone_tracks property, by default False.

Examples:

Access tropical cyclone tracks after a model call:

>>> model = WeatherNext2CyclonesMini.load_model(
...     WeatherNext2CyclonesMini.load_default_package(),
...     track_cyclones=True,
... )
>>> x, coords = model(x, coords)
>>> tracks = model.cyclone_tracks
>>> tracks[["track_id", "lead_time", "lat", "lon", "tcmsl", "tcw10m"]]

The tcmsl and tcw10m columns provide Earth2Studio-compatible names for the minimum sea-level pressure and surface wind speed diagnostics.

The tracker filters short-lived cyclogenesis tracks, so short rollouts can return an empty dataframe even when cyclone tracking is active. The active duration threshold is set by model._cyclone_tracker.cyclogenesis_minimum_duration.

__call__ 

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

Runs prognostic model one step.

Parameters:

  • x (Tensor) –

    Input tensor.

  • coords (CoordSystem) –

    Input coordinate system.

Returns:

  • tuple[Tensor, CoordSystem] –

    Output tensor and coordinate system 6 hours in the future.

create_iterator 

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

Create a time-integration iterator for the prognostic model.

Parameters:

  • x (Tensor) –

    Input tensor.

  • coords (CoordSystem) –

    Input coordinate system.

Yields:

load_default_package classmethod 

load_default_package() -> Package

Load default pre-trained WeatherNext 2 package from Google Cloud.

Returns:

load_model classmethod 

load_model(
    package: Package,
    seed: int = 0,
    jit_compile: bool = True,
    track_cyclones: bool = False,
) -> PrognosticModel

Load prognostic model from package.

Parameters:

  • package (Package) –

    Package to load model from.

  • seed (int, default: 0 ) –

    Initial random seed for the stochastic FGN noise generator, by default 0.

  • jit_compile (bool, default: True ) –

    JIT-compile the model forward pass, by default True.

  • track_cyclones (bool, default: False ) –

    Accumulate tropical cyclone tracks in the cyclone_tracks property, by default False.

Returns:

  • PrognosticModel –

    Prognostic model.