earth2studio.models.px.SFNO

class earth2studio.models.px.SFNO(core_model, center, scale, variables=array(['u10m', 'v10m', 'u100m', 'v100m', 't2m', 'sp', 'msl', 'tcwv', 'u50', 'u100', 'u150', 'u200', 'u250', 'u300', 'u400', 'u500', 'u600', 'u700', 'u850', 'u925', 'u1000', 'v50', 'v100', 'v150', 'v200', 'v250', 'v300', 'v400', 'v500', 'v600', 'v700', 'v850', 'v925', 'v1000', 'z50', 'z100', 'z150', 'z200', 'z250', 'z300', 'z400', 'z500', 'z600', 'z700', 'z850', 'z925', 'z1000', 't50', 't100', 't150', 't200', 't250', 't300', 't400', 't500', 't600', 't700', 't850', 't925', 't1000', 'q50', 'q100', 'q150', 'q200', 'q250', 'q300', 'q400', 'q500', 'q600', 'q700', 'q850', 'q925', 'q1000'], dtype='<U5'))

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
- https://github.com/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

• scale (torch.Tensor) – Model scale normalization tensor

• variables (np.array, optional) – Variables associated with model, by default 73 variable model.

__call__(x, coords)

Runs prognostic model 1 step. :param x: Input tensor :type x: torch.Tensor :param coords: Input coordinate system :type coords: CoordSystem

Returns:

• x (torch.Tensor)

• coords (CoordSystem)

Parameters:

• x (Tensor)

• coords (OrderedDict[str, ndarray])

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). :param x: Input tensor :type x: torch.Tensor :param coords: Input coordinate system :type coords: CoordSystem

Yields:

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

Parameters:

• x (Tensor)

• coords (OrderedDict[str, ndarray])

Return type:

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

classmethod load_default_package()

Load prognostic package

Return type:

Package

classmethod load_model(package, variables=['u10m', 'v10m', 'u100m', 'v100m', 't2m', 'sp', 'msl', 'tcwv', 'u50', 'u100', 'u150', 'u200', 'u250', 'u300', 'u400', 'u500', 'u600', 'u700', 'u850', 'u925', 'u1000', 'v50', 'v100', 'v150', 'v200', 'v250', 'v300', 'v400', 'v500', 'v600', 'v700', 'v850', 'v925', 'v1000', 'z50', 'z100', 'z150', 'z200', 'z250', 'z300', 'z400', 'z500', 'z600', 'z700', 'z850', 'z925', 'z1000', 't50', 't100', 't150', 't200', 't250', 't300', 't400', 't500', 't600', 't700', 't850', 't925', 't1000', 'q50', 'q100', 'q150', 'q200', 'q250', 'q300', 'q400', 'q500', 'q600', 'q700', 'q850', 'q925', 'q1000'])

Load prognostic from package

Parameters:

• package (Package)

• variables (list)

Return type:

PrognosticModel