neural_lam.datastore.base

neural_lam.datastore.base#

Abstract base classes describing Neural-LAM datastore APIs.

Classes#

`BaseDatastore`	Base class for weather data used in the neural-lam package. A datastore
`CartesianGridShape`	Dataclass to store the shape of a grid.

Module Contents#

class neural_lam.datastore.base.BaseDatastore[source]#

Bases: 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.

expected_dim_order(category: str | None = None) → tuple[str, Ellipsis][source]#

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:: The expected dimension order for the dataarray or dataset.
Return type:: list[str]

abstract get_dataarray(category: str, split: str | None, standardize: bool = False) → xarray.DataArray | None[source]#

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:

The xarray DataArray object with processed dataset.

Return type:

xr.DataArray or None

get_lat_lon(category: str) → numpy.ndarray[source]#: Return stacked longitude/latitude pairs for the requested category.

abstract get_num_data_vars(category: str) → int[source]#

Get the number of data variables in the given category.

Parameters:: category (str) – The category of the variables (state/forcing/static).
Returns:: The number of data variables.
Return type:: int

abstract get_standardization_dataarray(category: str) → xarray.Dataset[source]#

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:: The standardization dataarray for the given category, with variables for the mean and standard deviation of the variables (and differences for state variables).
Return type:: xr.Dataset

abstract get_vars_long_names(category: str) → list[str][source]#

Get the long names of the variables in the given category.

Parameters:: category (str) – The category of the variables (state/forcing/static).
Returns:: The long names of the variables.
Return type:: list[str]

abstract get_vars_names(category: str) → list[str][source]#

Get the names of the variables in the given category.

Parameters:: category (str) – The category of the variables (state/forcing/static).
Returns:: The names of the variables.
Return type:: list[str]

abstract get_vars_units(category: str) → list[str][source]#

Get the units of the variables in the given category.

Parameters:: category (str) – The category of the variables (state/forcing/static).
Returns:: The units of the variables.
Return type:: list[str]

abstract get_xy(category: str, stacked: bool) → numpy.ndarray[source]#

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:

The x, y coordinates of the dataset with shape [n_grid_points, 2].

Return type:

np.ndarray

get_xy_extent(category: str) → list[float][source]#

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:: The extent of the x, y coordinates.
Return type:: list[float]

stack_grid_coords(da_or_ds: xarray.DataArray) → xarray.DataArray[source]#

stack_grid_coords(da_or_ds: xarray.Dataset) → xarray.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:: The dataarray or dataset with the grid coordinates stacked.
Return type:: xr.DataArray or xr.Dataset
Raises:: NotImplementedError – If the datastore is not on a rectilinear grid.

unstack_grid_coords(da_or_ds: xarray.DataArray) → xarray.DataArray[source]#

unstack_grid_coords(da_or_ds: xarray.Dataset) → xarray.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:: The dataarray or dataset with the grid coordinates unstacked.
Return type:: xr.DataArray or xr.Dataset
Raises:: NotImplementedError – If the datastore is not on a rectilinear grid.

property boundary_mask: xarray.DataArray#

Abstractmethod:

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:: The boundary mask for the dataset, with dimensions (‘grid_index’,).
Return type:: xr.DataArray

property config: collections.abc.Mapping#

Abstractmethod:

The configuration of the datastore.

Returns:: The configuration of the datastore, any dict like object can be returned.
Return type:: collections.abc.Mapping

property coords_projection: cartopy.crs.Projection#

Abstractmethod:

Return the projection object for the coordinates.

The projection object is used to plot the coordinates on a map.

Returns:: The projection object.
Return type:: cartopy.crs.Projection

property grid_shape_state: CartesianGridShape#

Abstractmethod:

The shape of the grid for the state variables.

Returns:: The shape of the grid for the state variables, which has x and y attributes.
Return type:: CartesianGridShape
Raises:: NotImplementedError – If the datastore is not on a rectilinear grid.

has_ensemble_forcing: bool = False#

is_ensemble: bool = False#

is_forecast: bool = False#

property is_on_rectilinear_grid: bool#

Whether grid_index can be unstacked into a complete 2D x/y rectilinear grid. The spacing need not be uniform.

Returns:: False unless a subclass opts in.
Return type:: bool

property num_grid_points: 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:: The number of grid points in the dataset.
Return type:: int
Raises:: NotImplementedError – If the datastore is not on a rectilinear grid and does not override this.

property root_path: pathlib.Path#

Abstractmethod:

The root path to the datastore. It is relative to this that any derived files (for example the graph components) are stored.

Returns:: The root path to the datastore.
Return type:: pathlib.Path

spatial_coordinates: tuple[str, str] = ('x', 'y')#

property state_feature_weights_values: list[float]#

Abstractmethod:

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:: The weight for each state feature.
Return type:: list of float

property step_length: datetime.timedelta#

Abstractmethod:

The step length of the dataset as a time interval.

Returns:: The step length of the dataset.
Return type:: datetime.timedelta

class neural_lam.datastore.base.CartesianGridShape[source]#

Dataclass to store the shape of a grid.

x: int#

y: int#