FCN

GlobalMRF202240 GBNVIDIAPyTorch

Import path: earth2studio.models.px.FCN

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

FourCastNet 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 26 variables.

Note

This model is a retrained version on more atmospgeric variables from the FourCastNet paper. For additional information see the following resources:

• arxiv.org/abs/2202.11214
• huggingface.co/nvidia/fourcastnet1

Parameters:

• core_model (Module) –

  Core PyTorch model with loaded weights

• center (Tensor) –

  Model center normalization tensor of size [26]

• scale (Tensor) –

  Model scale normalization tensor of size [26]

__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 6 hours 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) -> PrognosticModel

Load prognostic from package

Examples using earth2studio.models.px.FCN

Running Deterministic InferenceBasic deterministic inference workflow. Running Diagnostic InferenceBasic prognostic + diagnostic inference workflow. Running Ensemble InferenceSimple ensemble inference workflow.

Checkpointing a Deterministic ForecastBasic inference workflow checkpointing. Tropical Cyclone TrackingTropical cyclone tracking with tracker diagnostic models. Extending Data SourcesImplementing a custom data source