Skip to content

InterpModAFNO

GlobalMRF202440 GBNVIDIAPyTorch

Import path: earth2studio.models.px.InterpModAFNO

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

ModAFNO interpolation for global prognostic models. Interpolates a forecast model to a shorter time-step size (by default from 6 to 1 hour). Operates on 0.25 degree lat-lon equirectangular grid with 73 variables.

Note

For more information on the model, please refer to:

Warning

The model requires a base forecast model to be set before execution. This can be done by setting the px_model attribute or using the load_model method.

Parameters:

  • interp_model (Module) –

    The interpolation model that performs the time interpolation

  • center (Tensor) –

    Model center normalization tensor

  • scale (Tensor) –

    Model scale normalization tensors

  • geop (Tensor) –

    Geopotential height data used as a static feature

  • lsm (Tensor) –

    Land-sea mask data used as a static feature

  • px_model (PrognosticModel, default: None ) –

    The base forecast model that produces the coarse time resolution forecasts. If not provide, should be set by the user before executing the model, by default None.

  • num_interp_steps (int, default: 6 ) –

    Number of interpolation steps to perform between forecast steps, by default 6

__call__ 

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

Runs prognostic model 1 step

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Returns:

  • tuple[Tensor, CoordSystem] –

    Output tensor and coordinate system 1 hour in the future

create_iterator 

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

Creates a iterator which can be used to perform time-integration of the prognostic model. Will return the initial condition first (0th step).

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Yields:

  • Iterator[tuple[Tensor, CoordSystem]] –

    Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

load_default_package classmethod 

load_default_package() -> Package

Load prognostic package

load_model classmethod 

load_model(
    package: Package,
    px_model: PrognosticModel | None = None,
) -> PrognosticModel

Load prognostic from package

Parameters:

  • package (Package) –

    Package to load model from

  • px_model (PrognosticModel | None, default: None ) –

    The base forecast model that produces the coarse time resolution forecasts. If None, should be set manually, by default None

Returns:

  • PrognosticModel –

    Prognostic model

Examples using earth2studio.models.px.InterpModAFNO

Preview of Temporal InterpolationTemporal InterpolationTemporal Interpolation inference using InterpModAFNO model.