earth2studio.models.px.SFNO

class earth2studio.models.px.SFNO(core_model, center, scale)

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:

• https://arxiv.org/abs/2306.03838

• NVIDIA/modulus-makani

• https://catalog.ngc.nvidia.com/orgs/nvidia/teams/modulus/models/sfno_73ch_small

Parameters:

• core_model (torch.nn.Module) – Core PyTorch model with loaded weights

• center (torch.Tensor) – Model center normalization tensor of size [73]

• scale (torch.Tensor) – Model scale normalization tensor of size [73]

__call__(x, coords)

Runs prognostic model 1 step.

Parameters:

• x (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Returns:

• x (torch.Tensor)

• coords (CoordSystem)

Return type:

tuple[Tensor, OrderedDict[str, ndarray]]

create_iterator(x, coords)

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 (torch.Tensor) – Input tensor

• coords (CoordSystem) – Input coordinate system

Yields:

Iterator[tuple[torch.Tensor, CoordSystem]] – Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

Return type:

Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]

classmethod load_default_package()

Load prognostic package

Return type:

Package

classmethod load_model(package)

Load prognostic from package

Parameters:

package (Package)

Return type:

PrognosticModel