neural_lam.datastore.npyfilesmeps.store

neural_lam.datastore.npyfilesmeps.store#

Numpy-files based datastore to support the MEPS example dataset introduced in neural-lam v0.1.0.

Module Contents#

class neural_lam.datastore.npyfilesmeps.store.NpyFilesDatastoreMEPS(config_path)#

Bases: neural_lam.datastore.base.BaseRegularGridDatastore

Represents a dataset stored as numpy files on disk. The dataset is assumed to be stored in a directory structure where each sample is stored in a separate file. The file-name format is assumed to be described by STATE_FILENAME_FORMAT.

The MEPS dataset is organised into three splits: train, val, and test. Each split has a set of files which are:

  • STATE_FILENAME_FORMAT:

    The state variables for a forecast started at analysis_time with member id member_id. The dimensions of the array are [forecast_timestep, y, x, feature].

  • TOA_SW_DOWN_FLUX_FILENAME_FORMAT:

    The top-of-atmosphere downwelling shortwave flux at time. The dimensions of the array are [forecast_timestep, y, x].

  • OPEN_WATER_FILENAME_FORMAT:

    The open water fraction at time. The dimensions of the array are [y, x].

Folder structure:

meps_example_reduced ├── data_config.yaml ├── samples │ ├── test │ │ ├── nwp_2022090100_mbr000.npy │ │ ├── nwp_2022090100_mbr001.npy │ │ ├── nwp_2022090112_mbr000.npy │ │ ├── nwp_2022090112_mbr001.npy │ │ ├── … │ │ ├── nwp_toa_downwelling_shortwave_flux_2022090100.npy │ │ ├── nwp_toa_downwelling_shortwave_flux_2022090112.npy │ │ ├── … │ │ ├── wtr_2022090100.npy │ │ ├── wtr_2022090112.npy │ │ └── … │ ├── train │ │ ├── nwp_2022040100_mbr000.npy │ │ ├── nwp_2022040100_mbr001.npy │ │ ├── … │ │ ├── nwp_2022040112_mbr000.npy │ │ ├── nwp_2022040112_mbr001.npy │ │ ├── … │ │ ├── nwp_toa_downwelling_shortwave_flux_2022040100.npy │ │ ├── nwp_toa_downwelling_shortwave_flux_2022040112.npy │ │ ├── … │ │ ├── wtr_2022040100.npy │ │ ├── wtr_2022040112.npy │ │ └── … │ └── val │ ├── nwp_2022060500_mbr000.npy │ ├── nwp_2022060500_mbr001.npy │ ├── … │ ├── nwp_2022060512_mbr000.npy │ ├── nwp_2022060512_mbr001.npy │ ├── … │ ├── nwp_toa_downwelling_shortwave_flux_2022060500.npy │ ├── nwp_toa_downwelling_shortwave_flux_2022060512.npy │ ├── … │ ├── wtr_2022060500.npy │ ├── wtr_2022060512.npy │ └── … └── static

├── border_mask.npy ├── diff_mean.pt ├── diff_std.pt ├── flux_stats.pt ├── grid_features.pt ├── nwp_xy.npy ├── parameter_mean.pt ├── parameter_std.pt ├── parameter_weights.npy └── surface_geopotential.npy

For the MEPS dataset: N_t’ = 65 N_t = 65//subsample_step (= 21 for 3h steps) dim_y = 268 dim_x = 238 N_grid = 268x238 = 63784 d_features = 17 (d_features’ = 18) d_forcing = 5

For the MEPS reduced dataset: N_t’ = 65 N_t = 65//subsample_step (= 21 for 3h steps) dim_y = 134 dim_x = 119 N_grid = 134x119 = 15946 d_features = 8 d_forcing = 1

Create a new NpyFilesDatastore using the configuration file at the given path. The config file should be a YAML file and will be loaded into an instance of the NpyDatastoreConfig dataclass.

Internally, the datastore uses dask.delayed to load the data from the numpy files, so that the data isn’t actually loaded until it’s needed.

Parameters:

config_path (str) – The path to the configuration file for the datastore.

coords_projection() → cartopy.crs.Projection#

The projection of the spatial coordinates.

Returns:

The projection of the spatial coordinates.

Return type:

ccrs.Projection

get_dataarray(category: str, split: str | None, standardize: bool = False) → xarray.core.dataarray.DataArray#

Get the data array for the given category and split of data. If the category is ‘state’, the data array will be a concatenation of the data arrays for all ensemble members. The data will be loaded as a dask array, so that the data isn’t actually loaded until it’s needed.

Parameters:

  • category (str) – The category of the data to load. One of ‘state’, ‘forcing’, or ‘static’.

  • split (str) – The dataset split to load the data for. One of ‘train’, ‘val’, or ‘test’.

  • standardize (bool) – If the dataarray should be returned standardized

Returns:

The data array for the given category and split, with dimensions per category: state: `[elapsed_forecast_duration, analysis_time, grid_index,

feature, ensemble_member]`

forcing: `[elapsed_forecast_duration, analysis_time, grid_index,

feature]`

static: [grid_index, feature]

Return type:

xr.DataArray

get_num_data_vars(category: str) → int#

Return the number of variables available in category.

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.

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

get_vars_long_names(category: str) → List[str]#

Return descriptive names for the variables in category.

get_vars_names(category: str) → List[str]#

Return canonical short names for the variables in category.

get_vars_units(category: str) → List[str]#

Return unit strings for the variables in category.

get_xy(category: str, stacked: bool) → numpy.ndarray#

Return the x, y coordinates of the dataset.

Parameters:

  • category (str) – The category of the dataset (state/forcing/static).

  • stacked (bool) – Whether to stack the x, y coordinates.

Returns:

The x, y coordinates of the dataset (with x first then y second), returned differently based on the value of stacked: - stacked==True: shape (n_grid_points, 2) where

n_grid_points=N_x*N_y.

  • stacked==False: shape (N_x, N_y, 2)

Return type:

np.ndarray

SHORT_NAME = 'npyfilesmeps'#
property boundary_mask: xarray.DataArray#

The boundary mask for the dataset. This is a binary mask that is 1 where the grid cell is on the boundary of the domain, and 0 otherwise.

Returns:

The boundary mask for the dataset, with dimensions [grid_index].

Return type:

xr.DataArray

property config: neural_lam.datastore.npyfilesmeps.config.NpyDatastoreConfig#

The configuration for the datastore.

Returns:

The configuration for the datastore.

Return type:

NpyDatastoreConfig

property grid_shape_state: 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.

Return type:

CartesianGridShape

is_ensemble = True#
is_forecast = True#
property root_path: pathlib.Path#

The root path of the datastore on disk. This is the directory relative to which graphs and other files can be stored.

Returns:

The root path of the datastore

Return type:

Path

property step_length: datetime.timedelta#

The length of each time step as a time interval.

Returns:

The length of each time step as a datetime.timedelta object.

Return type:

timedelta

neural_lam.datastore.npyfilesmeps.store.OPEN_WATER_FILENAME_FORMAT = 'wtr_{analysis_time:%Y%m%d%H}.npy'#
neural_lam.datastore.npyfilesmeps.store.STATE_FILENAME_FORMAT = 'nwp_{analysis_time:%Y%m%d%H}_mbr{member_id:03d}.npy'#
neural_lam.datastore.npyfilesmeps.store.TOA_SW_DOWN_FLUX_FILENAME_FORMAT = 'nwp_toa_downwelling_shortwave_flux_{analysis_time:%Y%m%d%H}.npy'#
Contents