earth2studio.models.px.FuXi#
- class earth2studio.models.px.FuXi(ort_short, ort_medium, ort_long)[source]#
FuXi weather model consists of three auto-regressive U-net transfomer models with a time-step size of 6 hours. The three models are trained to predict short (5days), medium (10 days) and longer (15 days) forecasts respectively. FuXi operates on 0.25 degree lat-lon grid (south-pole including) equirectangular grid with 70 atmospheric/surface variables. This model uses two time-steps as an input.
Note
This model uses the ONNX checkpoint from the original publication repository. For additional information see the following resources:
Note
To avoid ONNX init session overhead of this model we recommend setting the default Pytorch device to the correct target prior to model construction.
- Parameters:
ort_short (str) – Path to FuXi short model onnx file
ort_medium (str) – Path to FuXi medium model onnx file
ort_long (str) – Path to FuXi long model onnx file
- __call__(x, coords)[source]#
Runs short prognostic model 1 step.
- Parameters:
x (torch.Tensor) – Input tensor
coords (CoordSystem) – Input coordinate system
- Returns:
Output tensor and coordinate system 6 hours in the future
- Return type:
tuple[torch.Tensor, CoordSystem]
- create_iterator(x, coords)[source]#
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]]]