neural_lam.datastore.mdp ======================== .. py:module:: neural_lam.datastore.mdp .. autoapi-nested-parse:: Datastore implementation wrapping ``mllam-data-prep`` outputs. Classes ------- .. autoapisummary:: neural_lam.datastore.mdp.MDPDatastore Module Contents --------------- .. py:class:: MDPDatastore(config_path: str | pathlib.Path, n_boundary_points: int = 30, reuse_existing: bool = True) Bases: :py:obj:`neural_lam.datastore.base.BaseDatastore` Datastore class for datasets made with the mllam_data_prep library (https://github.com/mllam/mllam-data-prep). This class wraps the `mllam_data_prep` library to do the necessary transforms to create the different categories (state/forcing/static) of data, with the actual transform to do being specified in the configuration file. Construct a new MDPDatastore from the configuration file at `config_path`. A boundary mask is created with `n_boundary_points` boundary points. If `reuse_existing` is True, the dataset is loaded from a zarr file if it exists (unless the config has been modified since the zarr was created), otherwise it is created from the configuration file. :param config_path: The path to the configuration file, this will be fed to the `mllam_data_prep.Config.from_yaml_file` method to then call `mllam_data_prep.create_dataset` to create the dataset. :type config_path: str or Path :param n_boundary_points: Width, in grid points, of the boundary along each x/y edge of the domain. 0 gives no boundary points. :type n_boundary_points: int :param reuse_existing: Whether to reuse an existing dataset zarr file if it exists and its creation date is newer than the configuration file. :type reuse_existing: bool :raises ValueError: If the dataset does not contain all of the required train/val/test splits. .. py:method:: get_dataarray(category: str, split: str | None, standardize: bool = False) -> xarray.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). "state" is the only required category, for other categories, the method will return `None` if the category is not found in the datastore. The returned dataarray will 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 will 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 will have an additional `ensemble_member` dimension. If `has_ensemble_forcing=True`, the returned forcing dataarray will also have an additional `ensemble_member` dimension. :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 the longitude, latitude coordinates of the dataset as numpy array for a given category of data. Override in MDP to use lat/lons directly from xr.Dataset, if available. :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The longitude, latitude coordinates of the dataset with shape `[n_grid_points, 2]`. :rtype: np.ndarray .. py:method:: get_num_data_vars(category: str) -> int Return the number of variables in the given category. :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The number of variables in the given category. :rtype: int .. py:method:: get_standardization_dataarray(category: str) -> xarray.Dataset Return the standardization 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. :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] Return the long names of the variables in the given category. :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The long names of the variables in the given category. :rtype: list[str] .. py:method:: get_vars_names(category: str) -> list[str] Return the names of the variables in the given category. :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The names of the variables in the given category. :rtype: list[str] .. py:method:: get_vars_units(category: str) -> list[str] Return the units of the variables in the given category. :param category: The category of the dataset (state/forcing/static). :type category: str :returns: The units of the variables in the given category. :rtype: list[str] .. py:method:: get_xy(category: str, stacked: bool) -> numpy.ndarray Return the x, y coordinates of the dataset (i.e. the Cartesian coordinates of the regular gridded dataset). :param category: The category of the dataset (state/forcing/static). :type category: str :param stacked: Whether to stack the x, y coordinates. `stacked=False` requires `is_on_rectilinear_grid`. :type stacked: bool :returns: The x, y coordinates of the dataset, returned differently based on the value of `stacked`: * `stacked==True`: shape `(num_grid_nodes, 2)`, with num_grid_nodes=N_x*N_y on a rectilinear grid. * `stacked==False`: shape `(N_x, N_y, 2)` :rtype: np.ndarray .. py:attribute:: SHORT_NAME :value: 'mdp' .. py:property:: boundary_mask :type: xarray.DataArray Produce a 0/1 mask for the boundary points of the dataset, these will sit at the edges of the domain (in x/y extent) and will be used to mask out the boundary points from the loss function and to overwrite the boundary points from the prediction. For now this is created when the mask is requested, but in the future this could be saved to the zarr file. :returns: A 0/1 mask for the boundary points of the dataset, where 1 is a boundary point and 0 is not. :rtype: xr.DataArray :raises NotImplementedError: If the datastore provides no ``state`` data. .. py:property:: config :type: mllam_data_prep.Config The configuration of the dataset. :returns: The configuration of the dataset. :rtype: mdp.Config .. py:property:: coords_projection :type: cartopy.crs.Projection Return the projection of the coordinates. NOTE: currently this expects the projection information to be in the `extra` section of the configuration file, with a `projection` key containing a `class_name` and `kwargs` for constructing the `cartopy.crs.Projection` object. This is a temporary solution until the projection information can be parsed in the produced dataset itself. `mllam-data-prep` ignores the contents of the `extra` section of the config file which is why we need to check that the necessary parts are there. :returns: The projection of the coordinates. :rtype: ccrs.Projection :raises ValueError: If the `projection` entry is missing from the `extra` section of the config, or if its `class_name` or `kwargs` keys are missing. .. py:property:: grid_shape_state :type: neural_lam.datastore.base.CartesianGridShape The shape of the cartesian grid for the state variables. :returns: The shape of the cartesian grid for the state variables. :rtype: CartesianGridShape :raises NotImplementedError: If the datastore is not on a rectilinear grid, e.g. after ``domain_cropping``. Use `num_grid_points` instead. .. py:attribute:: has_ensemble_forcing .. py:attribute:: is_ensemble .. py:property:: is_on_rectilinear_grid :type: bool Whether ``grid_index`` can be unstacked into a complete 2D ``x``/``y`` rectilinear grid, i.e. every combination of the distinct ``x`` and ``y`` values is present exactly once. The spacing need not be uniform. A datastore cropped with ``domain_cropping`` is not. :returns: ``True`` when every cell of the ``x``/``y`` grid is present exactly once. :rtype: bool .. py:property:: num_grid_points :type: int Return the number of grid points in the dataset. :returns: The number of grid points in the dataset. :rtype: int .. py:property:: root_path :type: pathlib.Path The root path of the dataset. :returns: The root path of the dataset. :rtype: Path .. py:attribute:: spatial_coordinates :value: None .. py:property:: step_length :type: datetime.timedelta The length of the time steps as a time interval. :returns: The length of the time steps as a datetime.timedelta object. :rtype: timedelta