neural_lam.weather_dataset

neural_lam.weather_dataset#

Dataset helpers wrapping Neural-LAM datastores for PyTorch Lightning.

Module Contents#

class neural_lam.weather_dataset.WeatherDataModule(datastore: neural_lam.datastore.base.BaseDatastore, ar_steps_train: int = 3, ar_steps_eval: int = 25, standardize: bool = True, num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, batch_size: int = 4, num_workers: int = 16, eval_split: str = 'test')#

Bases: pytorch_lightning.LightningDataModule

DataModule for weather data.

Parameters:

  • datastore (BaseDatastore) – Datastore used for all splits.

  • ar_steps_train (int, optional) – Number of AR steps for training batches. Default 3.

  • ar_steps_eval (int, optional) – Number of AR steps for validation/test batches. Default 25.

  • standardize (bool, optional) – If True, datasets are returned standardized. Default True.

  • num_past_forcing_steps (int, optional) – Number of past forcing steps to include. Default 1.

  • num_future_forcing_steps (int, optional) – Number of future forcing steps to include. Default 1.

  • batch_size (int, optional) – Mini-batch size for dataloaders. Default 4.

  • num_workers (int, optional) – Number of background workers per dataloader. Default 16.

  • eval_split (str, optional) – Dataset split to use for test_dataloader. Default "test".

setup(stage=None)#

Instantiate datasets for the requested trainer stage.

Parameters:

stage (str or None, optional) – Trainer stage identifier ("fit"/"test"/None). When None, both train and evaluation datasets are created.

test_dataloader()#

Load test dataset.

train_dataloader()#

Load train dataset.

val_dataloader()#

Load validation dataset.

ar_steps_eval = 25#
ar_steps_train = 3#
batch_size = 4#
eval_split = 'test'#
multiprocessing_context: str | None = None#
num_future_forcing_steps = 1#
num_past_forcing_steps = 1#
num_workers: int = 16#
standardize = True#
test_dataset = None#
train_dataset = None#
val_dataset = None#
class neural_lam.weather_dataset.WeatherDataset(datastore: neural_lam.datastore.base.BaseDatastore, split: str = 'train', ar_steps: int = 3, num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, standardize: bool = True)#

Bases: torch.utils.data.Dataset

Dataset class for weather data.

This class loads and processes weather data from a given datastore.

Parameters:

  • datastore (BaseDatastore) – The datastore to load the data from (e.g. mdp).

  • split (str, optional) – The data split to use (“train”, “val” or “test”). Default is “train”.

  • ar_steps (int, optional) – The number of autoregressive steps. Default is 3.

  • num_past_forcing_steps (int, optional) – Number of past time steps to include in forcing input. If set to i, forcing from times t-i, t-i+1, …, t-1, t (and potentially beyond, given num_future_forcing_steps) are included as forcing inputs at time t Default is 1.

  • num_future_forcing_steps (int, optional) – Number of future time steps to include in forcing input. If set to j, forcing from times t, t+1, …, t+j-1, t+j (and potentially times before t, given num_past_forcing_steps) are included as forcing inputs at time t. Default is 1.

  • standardize (bool, optional) – Whether to standardize the data. Default is True.

:param : :type : param datastore: Datastore providing access to state/forcing/static arrays. :param : :type : type datastore: BaseDatastore :param : Default "train". :type : param split: Data split ("train", "val", or "test"). :param : :type : type split: str, optional :param : :type : param ar_steps: Number of autoregressive steps per training sample. Default 3. :param : :type : type ar_steps: int, optional :param : are concatenated. Default 1. :type : param num_past_forcing_steps: Past forcing window length i so that [t-i, ..., t] forcings :param : :type : type num_past_forcing_steps: int, optional :param : forcings are available. Default 1. :type : param num_future_forcing_steps: Future forcing window length j so that [t, ..., t+j] :param : :type : type num_future_forcing_steps: int, optional :param : :type : param standardize: If True, normalize state/forcing arrays via datastore stats. :param : :type : type standardize: bool, optional

create_dataarray_from_tensor(tensor: torch.Tensor, time: datetime.datetime | list[datetime.datetime], category: str)#

Construct a xarray.DataArray from a pytorch.Tensor with coordinates for grid_index, time and {category}_feature matching the shape and number of times provided and add the x/y coordinates from the datastore.

The number if times provided is expected to match the shape of the tensor. For a 2D tensor, the dimensions are assumed to be (grid_index, {category}_feature) and only a single time should be provided. For a 3D tensor, the dimensions are assumed to be (time, grid_index, {category}_feature) and a list of times should be provided.

Parameters:

  • tensor (torch.Tensor) – The tensor to construct the DataArray from, this assumed to have the same dimension ordering as returned by the __getitem__ method (i.e. time, grid_index, {category}_feature). The tensor will be copied to the CPU before constructing the DataArray.

  • time (datetime.datetime or list[datetime.datetime]) – The time or times of the tensor.

  • category (str) – The category of the tensor, either “state”, “forcing” or “static”.

Returns:

da – The constructed DataArray.

Return type:

xr.DataArray

ar_steps = 3#
da_forcing#
da_state#
datastore#
num_future_forcing_steps = 1#
num_past_forcing_steps = 1#
split = 'train'#
standardize = True#
Contents