neural_lam.datastore.base ========================= .. py:module:: neural_lam.datastore.base .. autoapi-nested-parse:: Abstract base classes describing Neural-LAM datastore APIs. Classes ------- .. autoapisummary:: neural_lam.datastore.base.BaseDatastore neural_lam.datastore.base.CartesianGridShape Module Contents --------------- .. py:class:: BaseDatastore Bases: :py:obj:`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. .. py:method:: expected_dim_order(category: str | None = None) -> tuple[str, Ellipsis] 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. :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The expected dimension order for the dataarray or dataset. :rtype: list[str] .. py:method:: get_dataarray(category: str, split: str | None, standardize: bool = False) -> xarray.DataArray | None :abstractmethod: 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. :param category: The category of the dataset (state/forcing/static). :type category: str :param split: The time split to filter the dataset (train/val/test). :type split: str :param standardize: If the dataarray should be returned standardized :type standardize: bool :returns: The xarray DataArray object with processed dataset. :rtype: xr.DataArray or None .. py:method:: get_lat_lon(category: str) -> numpy.ndarray Return stacked longitude/latitude pairs for the requested category. .. py:method:: get_num_data_vars(category: str) -> int :abstractmethod: Get the number of data variables in the given category. :param category: The category of the variables (state/forcing/static). :type category: str :returns: The number of data variables. :rtype: int .. py:method:: get_standardization_dataarray(category: str) -> xarray.Dataset :abstractmethod: 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). :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The standardization dataarray for the given category, with variables for the mean and standard deviation of the variables (and differences for state variables). :rtype: xr.Dataset .. py:method:: get_vars_long_names(category: str) -> list[str] :abstractmethod: Get the long names of the variables in the given category. :param category: The category of the variables (state/forcing/static). :type category: str :returns: The long names of the variables. :rtype: list[str] .. py:method:: get_vars_names(category: str) -> list[str] :abstractmethod: Get the names of the variables in the given category. :param category: The category of the variables (state/forcing/static). :type category: str :returns: The names of the variables. :rtype: list[str] .. py:method:: get_vars_units(category: str) -> list[str] :abstractmethod: Get the units of the variables in the given category. :param category: The category of the variables (state/forcing/static). :type category: str :returns: The units of the variables. :rtype: list[str] .. py:method:: get_xy(category: str, stacked: bool) -> numpy.ndarray :abstractmethod: Return the x, y coordinates of the dataset as a numpy arrays for a given category of data. :param category: The category of the dataset (state/forcing/static). :type category: str :param stacked: Whether to stack the x, y coordinates. `stacked=False` is only meaningful for grid points on a regular-grid. :type stacked: bool :returns: The x, y coordinates of the dataset with shape `[n_grid_points, 2]`. :rtype: np.ndarray .. py:method:: get_xy_extent(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. :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The extent of the x, y coordinates. :rtype: list[float] .. py:method:: stack_grid_coords(da_or_ds: xarray.DataArray) -> xarray.DataArray 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). :param da_or_ds: The dataarray or dataset to stack the grid coordinates of. :type da_or_ds: xr.DataArray or xr.Dataset :returns: The dataarray or dataset with the grid coordinates stacked. :rtype: xr.DataArray or xr.Dataset :raises NotImplementedError: If the datastore is not on a rectilinear grid. .. py:method:: unstack_grid_coords(da_or_ds: xarray.DataArray) -> xarray.DataArray 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). :param da_or_ds: The dataarray or dataset to unstack the grid coordinates of. :type da_or_ds: xr.DataArray or xr.Dataset :returns: The dataarray or dataset with the grid coordinates unstacked. :rtype: xr.DataArray or xr.Dataset :raises NotImplementedError: If the datastore is not on a rectilinear grid. .. py:property:: boundary_mask :type: 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',)`. :rtype: xr.DataArray .. py:property:: config :type: collections.abc.Mapping :abstractmethod: The configuration of the datastore. :returns: The configuration of the datastore, any dict like object can be returned. :rtype: collections.abc.Mapping .. py:property:: coords_projection :type: 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. :rtype: cartopy.crs.Projection .. py:property:: grid_shape_state :type: 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. :rtype: CartesianGridShape :raises NotImplementedError: If the datastore is not on a rectilinear grid. .. py:attribute:: has_ensemble_forcing :type: bool :value: False .. py:attribute:: is_ensemble :type: bool :value: False .. py:attribute:: is_forecast :type: bool :value: False .. py:property:: is_on_rectilinear_grid :type: 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. :rtype: bool .. py:property:: num_grid_points :type: 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. :rtype: int :raises NotImplementedError: If the datastore is not on a rectilinear grid and does not override this. .. py:property:: root_path :type: 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. :rtype: pathlib.Path .. py:attribute:: spatial_coordinates :type: tuple[str, str] :value: ('x', 'y') .. py:property:: state_feature_weights_values :type: 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. :rtype: list of float .. py:property:: step_length :type: datetime.timedelta :abstractmethod: The step length of the dataset as a time interval. :returns: The step length of the dataset. :rtype: datetime.timedelta .. py:class:: CartesianGridShape Dataclass to store the shape of a grid. .. py:attribute:: x :type: int .. py:attribute:: y :type: int