earth2studio.models.px.Aurora#
- class earth2studio.models.px.Aurora(core_model, z, slt, lsm)[source]#
Aurora 0.25 degree global forecast model. This model consists of single auto-regressive model with a time-step size of 6 hours. This model operates on 0.25 degree lat-lon grid (720, 1440) equirectangular grid with 4 surface-level variables, 5 atmospheric variables with 13 pressure levels and 3 static variables.
Note
This model uses the checkpoints from the original publication. For additional information see the following resources:
Warning
We encourage users to familiarize themselves with the license restrictions of this model’s checkpoints.
- Parameters:
core_model (torch.nn.Module) – Core Aurora model
z (torch.Tensor) – Geopotential
slt (torch.Tensor) – Soil type
lsm (torch.Tensor) – Land sea mask
- __call__(x, coords)[source]#
Runs 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]]]