earth2studio.models.px.FengWu#
- class earth2studio.models.px.FengWu(ort, center, scale)[source]#
FengWu (operational) weather model consists of single auto-regressive model with a time-step size of 6 hours. FengWu operates on 0.25 degree lat-lon grid (south-pole including) equirectangular grid with 69 atmospheric/surface variables. This model uses two time-steps as an input.
Note
This model uses the ONNX checkpoint from the original publication repository. This checkpoint is a operational version to the one used in the paper which requires less variables. 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 (str) – Path to FengWu 6 hour onnx file
center (torch.Tensor) – Model variable center normalization tensor of size [69]
scale (torch.Tensor) – Model variable scale normalization tensor of size [69]
- __call__(x, coords)[source]#
Runs 6 hour 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]]]