"""Abstract base classes describing Neural-LAM datastore APIs."""
# Standard library
import abc
import collections
import dataclasses
import functools
from datetime import timedelta
from functools import cached_property
from pathlib import Path
from typing import overload
# Third-party
import cartopy.crs as ccrs
import numpy as np
import xarray as xr
from pandas.core.indexes.multi import MultiIndex
[docs]
@dataclasses.dataclass
class CartesianGridShape:
"""Dataclass to store the shape of a grid."""
x: int
y: int
[docs]
class BaseDatastore(abc.ABC):
"""
Base class for weather data used in the neural-lam package. A datastore
defines the interface for accessing weather data by providing methods to
access the data in a processed format that can be used for training and
evaluation of neural networks.
NOTE: All methods return either primitive types, `numpy.ndarray`,
`xarray.DataArray` or `xarray.Dataset` objects, not `pytorch.Tensor`
objects. Conversion to `pytorch.Tensor` objects should be done in the
`weather_dataset.WeatherDataset` class (which inherits from
`torch.utils.data.Dataset` and uses the datastore to access the data).
Forecast vs analysis data
-------------------------
If the datastore is used to represent forecast rather than analysis
data, then the ``is_forecast`` attribute should be set to True, and
returned data from ``get_dataarray`` is assumed to have `analysis_time`
and `forecast_time` dimensions (rather than just `time`).
Ensemble vs deterministic data
------------------------------
If the datastore is used to present an ensemble of state realisations, for
example for forecast ensembles, then the `is_ensemble` attribute should be
set to `True` and returned state data from ``get_dataarray`` is expected to
have an `ensemble_member` dimension. If each ensemble member has its own
forcing values, then `has_ensemble_forcing` should be set to `True`, and
returned forcing data from ``get_dataarray`` is expected to have an
`ensemble_member` dimension; otherwise forcing data is expected not to have
one.
Grid index
----------
All methods that return data specific to a grid point (like
``get_dataarray``) should have a single dimension named `grid_index` that
represents the spatial grid index of the data. The actual x, y coordinates
of the grid points should be stored in the `x` and `y` coordinates of the
dataarray or dataset with the `grid_index` dimension as the coordinate for
each of the `x` and `y` coordinates.
Rectilinear grid
----------------
A datastore whose `grid_index` can be reshaped back into a complete 2D
`x`/`y` grid, with possibly non-uniform spacing, should set
`is_on_rectilinear_grid` to `True` and implement `grid_shape_state`.
`grid_shape_state`, `stack_grid_coords` and `unstack_grid_coords` raise
`NotImplementedError` otherwise.
"""
is_ensemble: bool = False
has_ensemble_forcing: bool = False
is_forecast: bool = False
spatial_coordinates: tuple[str, str] = ("x", "y")
@property
@abc.abstractmethod
def root_path(self) -> Path:
"""
The root path to the datastore. It is relative to this that any derived
files (for example the graph components) are stored.
Returns
-------
pathlib.Path
The root path to the datastore.
"""
@property
@abc.abstractmethod
def config(self) -> collections.abc.Mapping:
"""The configuration of the datastore.
Returns
-------
collections.abc.Mapping
The configuration of the datastore, any dict like object can be
returned.
"""
@property
@abc.abstractmethod
def step_length(self) -> timedelta:
"""The step length of the dataset as a time interval.
Returns
-------
datetime.timedelta
The step length of the dataset.
"""
[docs]
@abc.abstractmethod
def get_vars_units(self, category: str) -> list[str]:
"""Get the units of the variables in the given category.
Parameters
----------
category : str
The category of the variables (state/forcing/static).
Returns
-------
list[str]
The units of the variables.
"""
[docs]
@abc.abstractmethod
def get_vars_names(self, category: str) -> list[str]:
"""Get the names of the variables in the given category.
Parameters
----------
category : str
The category of the variables (state/forcing/static).
Returns
-------
list[str]
The names of the variables.
"""
[docs]
@abc.abstractmethod
def get_vars_long_names(self, category: str) -> list[str]:
"""Get the long names of the variables in the given category.
Parameters
----------
category : str
The category of the variables (state/forcing/static).
Returns
-------
list[str]
The long names of the variables.
"""
[docs]
@abc.abstractmethod
def get_num_data_vars(self, category: str) -> int:
"""Get the number of data variables in the given category.
Parameters
----------
category : str
The category of the variables (state/forcing/static).
Returns
-------
int
The number of data variables.
"""
[docs]
@abc.abstractmethod
def get_standardization_dataarray(self, category: str) -> xr.Dataset:
"""
Return the standardization (i.e. scaling to mean of 0.0 and standard
deviation of 1.0) dataarray for the given category. This should contain
a `{category}_mean` and `{category}_std` variable for each variable in
the category.
For `category=="state"`, the dataarray should also contain a
`state_diff_mean_standardized` and `state_diff_std_standardized`
variable for the one-step differences of the state variables.
The returned dataarray should at least have dimensions of
`({category}_feature)`, but can also include for example `grid_index`
(if the standardization is done per grid point for example).
Parameters
----------
category : str
The category of the dataset (state/forcing/static).
Returns
-------
xr.Dataset
The standardization dataarray for the given category, with variables
for the mean and standard deviation of the variables (and
differences for state variables).
"""
def _standardize_datarray(
self, da: xr.DataArray, category: str
) -> xr.DataArray:
"""
Helper function to standardize a dataarray before returning it.
Parameters
----------
da: xr.DataArray
The dataarray to standardize
category : str
The category of the dataarray (state/forcing/static), to load
standardization statistics for.
Returns
-------
xr.Dataarray
The standardized dataarray
"""
standard_da = self.get_standardization_dataarray(category=category)
mean = standard_da[f"{category}_mean"]
std = standard_da[f"{category}_std"]
eps = np.finfo(std.dtype).eps
return (da - mean) / std.where(std > eps, other=eps)
[docs]
@abc.abstractmethod
def get_dataarray(
self,
category: str,
split: str | None,
standardize: bool = False,
) -> xr.DataArray | None:
"""
Return the processed data (as a single `xr.DataArray`) for the given
category of data and test/train/val-split that covers all the data (in
space and time) of a given category (state/forcing/static). A
datastore must be able to return for the "state" category, but
"forcing" and "static" are optional (in which case the method should
return `None`). For the "static" category the `split` is allowed to be
`None` because the static data is the same for all splits.
The returned dataarray is expected to at minimum have dimensions of
`(grid_index, {category}_feature)` so that any spatial dimensions have
been stacked into a single dimension and all variables and levels have
been stacked into a single feature dimension named by the `category` of
data being loaded.
For categories of data that have a time dimension (i.e. not static
data), the dataarray is expected additionally have `(analysis_time,
elapsed_forecast_duration)` dimensions if `is_forecast` is True, or
`(time)` if `is_forecast` is False.
If we have multiple ensemble members of state data, the returned state
dataarray is expected to have an additional `ensemble_member`
dimension. If `has_ensemble_forcing=True`, the returned forcing
dataarray is expected to have an additional `ensemble_member`
dimension; otherwise it is expected not to have one.
Parameters
----------
category : str
The category of the dataset (state/forcing/static).
split : str
The time split to filter the dataset (train/val/test).
standardize: bool
If the dataarray should be returned standardized
Returns
-------
xr.DataArray or None
The xarray DataArray object with processed dataset.
"""
@cached_property
@abc.abstractmethod
def boundary_mask(self) -> xr.DataArray:
"""
Return the boundary mask for the dataset, with spatial dimensions
stacked. Where the value is 1, the grid point is a boundary point, and
where the value is 0, the grid point is not a boundary point.
Returns
-------
xr.DataArray
The boundary mask for the dataset, with dimensions
`('grid_index',)`.
"""
[docs]
@abc.abstractmethod
def get_xy(self, category: str, stacked: bool) -> np.ndarray:
"""
Return the x, y coordinates of the dataset as a numpy arrays for a
given category of data.
Parameters
----------
category : str
The category of the dataset (state/forcing/static).
stacked : bool
Whether to stack the x, y coordinates. `stacked=False` is only
meaningful for grid points on a regular-grid.
Returns
-------
np.ndarray
The x, y coordinates of the dataset with shape `[n_grid_points, 2]`.
"""
@property
@abc.abstractmethod
def coords_projection(self) -> ccrs.Projection:
"""Return the projection object for the coordinates.
The projection object is used to plot the coordinates on a map.
Returns
-------
cartopy.crs.Projection:
The projection object.
"""
[docs]
@functools.lru_cache
def get_xy_extent(self, category: str) -> list[float]:
"""
Return the extent of the x, y coordinates for a given category of data.
The extent should be returned as a list of 4 floats with `[xmin, xmax,
ymin, ymax]` which can then be used to set the extent of a plot.
Parameters
----------
category : str
The category of the dataset (state/forcing/static).
Returns
-------
list[float]
The extent of the x, y coordinates.
"""
xy = self.get_xy(category, stacked=True)
extent = [
xy[:, 0].min(),
xy[:, 0].max(),
xy[:, 1].min(),
xy[:, 1].max(),
]
return [float(v) for v in extent]
[docs]
@functools.lru_cache
def get_lat_lon(self, category: str) -> np.ndarray:
"""
Return stacked longitude/latitude pairs for the requested category.
"""
xy = self.get_xy(category=category, stacked=True)
if xy.size == 0:
return xy
lon_lat = ccrs.PlateCarree().transform_points(
self.coords_projection,
xy[:, 0],
xy[:, 1],
)[:, :2]
return lon_lat
@property
def num_grid_points(self) -> int:
"""Return the number of grid points in the dataset.
Derived from `grid_shape_state`, so datastores not on a rectilinear
grid must override this.
Returns
-------
int
The number of grid points in the dataset.
Raises
------
NotImplementedError
If the datastore is not on a rectilinear grid and does not
override this.
"""
if not self.is_on_rectilinear_grid:
raise NotImplementedError(
f"{type(self).__name__} is not on a rectilinear grid, "
"so `num_grid_points` cannot be derived from "
"`grid_shape_state` and must be overridden."
)
return self.grid_shape_state.x * self.grid_shape_state.y
@cached_property
@abc.abstractmethod
def state_feature_weights_values(self) -> list[float]:
"""
Return the weights for each state feature as a list of floats.
The weights are defined by the user in a config file for the
datastore. Implementations must assert that there is one weight
for each state feature. The weights can be used to scale the
loss function for each state variable.
Returns
-------
list of float
The weight for each state feature.
"""
[docs]
@functools.lru_cache
def expected_dim_order(
self,
category: str | None = None,
) -> tuple[str, ...]:
"""
Return the expected dimension order for the dataarray or dataset
returned by `get_dataarray` for the given category of data. The
dimension order is the order of the dimensions in the dataarray or
dataset, and is used to check that the data is in the expected format.
This is necessary so that when stacking and unstacking the spatial grid
we can ensure that the dimension order is the same as what is returned
from `get_dataarray`. And also ensures that downstream uses of a
datastore (e.g. WeatherDataset) sees the data in a common structure.
If the category is None, then the it assumed that data only represents
a 1D scalar field varying with grid-index.
The order is constructed to match the order in `pytorch.Tensor` objects
that will be constructed from the data so that the last two dimensions
are always the grid-index and feature dimensions (i.e. the order is
`[..., grid_index, {category}_feature]`), with any time-related and
ensemble-number dimension(s) coming before these two.
Parameters
----------
category : str
The category of the dataset (state/forcing/static).
Returns
-------
list[str]
The expected dimension order for the dataarray or dataset.
"""
dim_order = []
if category is not None:
if category != "static":
# static data does not vary in time
if self.is_forecast:
dim_order.extend(
[
"analysis_time",
"elapsed_forecast_duration",
]
)
elif not self.is_forecast:
dim_order.append("time")
if (
category == "state"
and self.is_ensemble
or category == "forcing"
and self.has_ensemble_forcing
):
dim_order.append("ensemble_member")
dim_order.append("grid_index")
if category is not None:
dim_order.append(f"{category}_feature")
return tuple(dim_order)
@property
def is_on_rectilinear_grid(self) -> bool:
"""Whether ``grid_index`` can be unstacked into a complete 2D
``x``/``y`` rectilinear grid. The spacing need not be uniform.
Returns
-------
bool
``False`` unless a subclass opts in.
"""
return False
@cached_property
def grid_shape_state(self) -> CartesianGridShape:
"""The shape of the grid for the state variables.
Returns
-------
CartesianGridShape:
The shape of the grid for the state variables, which has `x` and
`y` attributes.
Raises
------
NotImplementedError
If the datastore is not on a rectilinear grid.
"""
raise NotImplementedError(
f"{type(self).__name__} does not define `grid_shape_state`; "
"this is only available for datastores where "
"`is_on_rectilinear_grid` is True."
)
def _check_on_rectilinear_grid(self, operation: str) -> None:
"""Guard an operation that requires a complete 2D ``x``/``y`` grid.
Parameters
----------
operation : str
Name of the calling method, used in the error message.
Raises
------
NotImplementedError
If the datastore's grid points do not form a complete 2D grid.
"""
if not self.is_on_rectilinear_grid:
raise NotImplementedError(
f"`{operation}` requires the grid points to form a complete "
"2D grid (`is_on_rectilinear_grid`); "
f"{type(self).__name__} reports that they do not."
)
@overload
def unstack_grid_coords(self, da_or_ds: xr.DataArray) -> xr.DataArray:
"""Unstack spatial grid coordinates of DataArray."""
@overload
def unstack_grid_coords(self, da_or_ds: xr.Dataset) -> xr.Dataset:
"""Unstack spatial grid coordinates of Dataset."""
[docs]
def unstack_grid_coords(
self, da_or_ds: xr.DataArray | xr.Dataset
) -> xr.DataArray | xr.Dataset:
"""
Unstack the spatial grid coordinates from `grid_index` into separate `x`
and `y` dimensions to create a 2D grid (if the spatial coordinates have
different names, those are used instead). Only performs unstacking if
the data is currently stacked (has grid_index dimension).
Parameters
----------
da_or_ds : xr.DataArray or xr.Dataset
The dataarray or dataset to unstack the grid coordinates of.
Returns
-------
xr.DataArray or xr.Dataset
The dataarray or dataset with the grid coordinates unstacked.
Raises
------
NotImplementedError
If the datastore is not on a rectilinear grid.
"""
# Return original data if already unstacked (no grid_index dimension)
if "grid_index" not in da_or_ds.dims:
return da_or_ds
self._check_on_rectilinear_grid("unstack_grid_coords")
# Check whether `grid_index` is a multi-index
if not isinstance(da_or_ds.indexes.get("grid_index"), MultiIndex):
da_or_ds = da_or_ds.set_index(grid_index=self.spatial_coordinates)
da_or_ds_unstacked = da_or_ds.unstack("grid_index")
# Ensure that the x, y dimensions are in the correct order
dims = list(da_or_ds_unstacked.dims)
xy_dim_order = [d for d in dims if d in self.spatial_coordinates]
if xy_dim_order != self.spatial_coordinates:
# work out where the first spatial coordinate is located
# so that we can insert the second spatial coordinate next to it in
# the correct order. Although this looks verbose, it ensures that
# we don't change the order of any other dimensions.
first_xy_dim_index = min(
dims.index(self.spatial_coordinates[0]),
dims.index(self.spatial_coordinates[1]),
)
new_dim_order = list(dims)
new_dim_order.remove(self.spatial_coordinates[0])
new_dim_order.remove(self.spatial_coordinates[1])
new_dim_order.insert(
first_xy_dim_index, self.spatial_coordinates[0]
)
new_dim_order.insert(
first_xy_dim_index + 1, self.spatial_coordinates[1]
)
da_or_ds_unstacked = da_or_ds_unstacked.transpose(*new_dim_order)
return da_or_ds_unstacked
@overload
def stack_grid_coords(self, da_or_ds: xr.DataArray) -> xr.DataArray:
"""Stack spatial grid coordinates of DataArray."""
@overload
def stack_grid_coords(self, da_or_ds: xr.Dataset) -> xr.Dataset:
"""Stack spatial grid coordinates of Dataset."""
[docs]
def stack_grid_coords(
self, da_or_ds: xr.DataArray | xr.Dataset
) -> xr.DataArray | xr.Dataset:
"""
Stack the spatial grid coordinates (x and y) into a single `grid_index`
dimension. Only performs stacking if the data is currently unstacked
(has x and y dimensions).
Parameters
----------
da_or_ds : xr.DataArray or xr.Dataset
The dataarray or dataset to stack the grid coordinates of.
Returns
-------
xr.DataArray or xr.Dataset
The dataarray or dataset with the grid coordinates stacked.
Raises
------
NotImplementedError
If the datastore is not on a rectilinear grid.
"""
# Return original data if already stacked (has grid_index dimension)
if "grid_index" in da_or_ds.dims:
return da_or_ds
self._check_on_rectilinear_grid("stack_grid_coords")
da_or_ds_stacked = da_or_ds.stack(grid_index=self.spatial_coordinates)
# infer what category of data the array represents by finding the
# dimension named in the format `{category}_feature`
category = None
for dim in da_or_ds_stacked.dims:
if isinstance(dim, str) and dim.endswith("_feature"):
if category is not None:
raise ValueError(
"Multiple dimensions ending with '_feature' found in "
f"dataarray: {da_or_ds_stacked}. Cannot infer category."
)
category = dim.split("_")[0]
dim_order = self.expected_dim_order(category=category)
return da_or_ds_stacked.transpose(*dim_order)
