`SFNO`¶

GlobalMRF202340 GBNVIDIAPyTorch

Import path: earth2studio.models.px.SFNO

View source on GitHub View install commands

Documentation¶

Bases: Module, AutoModelMixin, PrognosticMixin

Spherical Fourier Operator Network global prognostic model. Consists of a single model with a time-step size of 6 hours. FourCastNet operates on 0.25 degree lat-lon grid (south-pole excluding) equirectangular grid with 73 variables.

Note

This model and checkpoint are trained using Modulus-Makani. For more information see the following references:

Parameters:

core_model (Module) –

Core PyTorch model with loaded weights
variables (array, default: array(VARIABLES) ) –

Variables associated with model, by default 73 variable model.

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

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,
    variables: list = VARIABLES,
    device: str = "cpu",
) -> PrognosticModel

Load prognostic from package

Parameters:

package (Package) –

Package to load model from
variables (list, default: VARIABLES ) –

Model variable override, by default VARIABLES for SFNO 73 channel