earth2studio.models.px.Pangu24#
- class earth2studio.models.px.Pangu24(ort_24hr)[source]#
Pangu Weather 24 hour model. This model consists of single auto-regressive model with a time-step size of 24 hours. 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
- __call__(x, coords)[source]#
Runs 24 hour prognostic model 1 step.
- Parameters:
x (torch.Tensor) – Input tensor
coords (CoordSystem) – Input coordinate system
- Returns:
Output tensor and coordinate system 24 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]]]
