earth2studio.models.px.DLWP#
- class earth2studio.models.px.DLWP(core_model, landsea_mask, orography, latgrid, longrid, cubed_sphere_transform, cubed_sphere_inverse, center, scale)[source]#
Deep learning weather prediction (DLWP) prognostic model. This is a parsimonious global forecast model with a time-step size of 6 hours. The core model is a convolutional encoder-decoder trained on [64,64] cubed sphere data that has an input of 18 fields (2x7 atmos variables + 4 prescriptive) and outputs 14 fields (2x7 atmos variables). This implementation provides a wrapper that accepts [721,1440] lat-lon equirectangular grid of just the atmospheric varaibles as an input for better compatability with common data sources. Prescriptive fields are added inside the model wrapper.
Note
For more information about this model see:
- Parameters:
core_model (torch.nn.Module) – Core cubed-sphere DLWP model.
landsea_mask (torch.Tensor) – Land sea mask in cubed sphere form [6,64,64]
orography (torch.Tensor) – Surface geopotential (orography) in cubed sphere form [6,64,64]
latgrid (torch.Tensor) – Cubed sphere latitude coordinates [6,64,64]
longrid (torch.Tensor) – Cubed sphere longitude coordinates [6,64,64]
cubed_sphere_transform (torch.Tensor) – Sparse pytorch tensor to transform equirectangular fields to cubed sphere of size [24576, 1038240]
cubed_sphere_inverse (torch.Tensor) – Sparse pytorch tensor to transform cubed sphere fields to equirectangular of size [1038240, 24576]
center (torch.Tensor) – Model atmospheric variable center normalization tensor of size [1,7,1,1]
scale (torch.Tensor) – Model atmospheric variable scale normalization tensor of size [1,7,1,1]
- __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]]]
