earth2studio.models.px.Pangu3#
- class earth2studio.models.px.Pangu3(ort_24hr, ort_6hr, ort_3hr)[source]#
Pangu Weather 3 hour model. This model consists of three underlying auto-regressive models with a time-step size of 24, 6 and 3 hours. These three models are interweaved during prediction. Pangu Weather operates on 0.25 degree lat-lon grid (south-pole including) equirectangular grid with 69 atmospheric/surface variables.
Note
This model uses the ONNX checkpoints from the original publication. 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.
Warning
We encourage users to familiarize themselves with the license restrictions of this model’s checkpoints.
- Parameters:
ort_24hr (str) – Path to Pangu 24 hour onnx file
ort_6hr (str) – Path to Pangu 6 hour onnx file
ort_3hr (str) – Path to Pangu 3 hour onnx file
- __call__(x, coords)[source]#
Runs 3 hour prognostic model 1 step.
- Parameters:
x (torch.Tensor) – Input tensor
coords (CoordSystem) – Input coordinate system
- Returns:
Output tensor and coordinate system 3 hours in the future
- Return type:
tuple[torch.Tensor, CoordSystem]
- 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]]]