WeatherNext2Cyclones

GlobalMRF202680 GBGoogleJAX

Import path: earth2studio.models.px.WeatherNext2Cyclones

View source on GitHub View install commands

Documentation

Bases: _WeatherNext2Base

WeatherNext 2 Cyclones operational medium-range forecast model.

This wrapper uses Google DeepMind's operational 0.25 degree WeatherNextCyclones_<2025 checkpoint family. These are the models that ran during the 2025 Atlantic hurricane season. Four trained checkpoint members are available; load_model selects member 1 by default.

The model requires two input states, valid at -6h and 0h lead time, and predicts 6 hours forward per model call. Cyclone tracking can be enabled with track_cyclones=True to accumulate WeatherNext's tropical cyclone diagnostics in the cyclone_tracks property.

Note

To avoid JAX preallocating GPU memory and use the CUDA virtual memory management allocator, set these variables before importing JAX or Earth2Studio:

.. code-block:: console

export XLA_PYTHON_CLIENT_PREALLOCATE=false
export XLA_PYTHON_CLIENT_ALLOCATOR=vmm

Note

For more information see the following references:

• doi.org/10.1038/s41586-026-10953-2
• google-deepmind/weathernext#provided-pretrained-models
• huggingface.co/kashif/weathernext2
• docs.jax.dev/en/latest/gpu_memory_allocation.html

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 = WeatherNext2Cyclones.load_model(
...     WeatherNext2Cyclones.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:

• Iterator[tuple[Tensor, CoordSystem]] –

  Iterator that generates model time steps.

load_default_package classmethod

load_default_package() -> Package

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

Returns:

• Package –

  Model package.

load_model classmethod

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

Load the operational prognostic model from a 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.

• checkpoint_member (int, default: 1 ) –

  Operational checkpoint member from 1 through 4, by default 1.

Returns:

• PrognosticModel –

  Prognostic model.