neural_lam.weather_dataset

neural_lam.weather_dataset#

Dataset helpers wrapping Neural-LAM datastores for PyTorch Lightning.

Classes#

`WeatherDataModule`	DataModule for weather data.
`WeatherDataset`	Dataset class for weather data.

Module Contents#

class neural_lam.weather_dataset.WeatherDataModule(datastore: neural_lam.datastore.base.BaseDatastore, ar_steps_train: int = 3, ar_steps_eval: int = 25, num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, num_past_boundary_steps: int = 1, num_future_boundary_steps: int = 1, datastore_boundary: neural_lam.datastore.base.BaseDatastore | None = None, load_single_member: bool = False, batch_size: int = 4, num_workers: int = 16, eval_split: str = 'test')[source]#

Bases: pytorch_lightning.LightningDataModule

DataModule for weather data.

Parameters:

datastore (BaseDatastore) – Datastore used for all splits.
ar_steps_train (int, optional) – Number of autoregressive steps for training batches. Default 3.
ar_steps_eval (int, optional) – Number of autoregressive steps for validation/test batches. Default 25.
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.
num_past_boundary_steps (int, optional) – Number of past boundary forcing steps to include. Default 1.
num_future_boundary_steps (int, optional) – Number of future boundary forcing steps to include. Default 1.
datastore_boundary (BaseDatastore, optional) – Separate datastore providing boundary forcing, used for all splits. If None, the boundary tensor is empty.
load_single_member (bool, optional) – If True, load only a single ensemble member per sample. Default False.
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: str | None = None) → None[source]#

Instantiate datasets for the requested trainer stage.

Parameters:: stage (str or None, optional) – Trainer stage identifier ("fit"/"test"/None). When None, both the training split and the validation/test evaluation splits are prepared.

test_dataloader() → torch.utils.data.DataLoader[source]#

Load test dataset.

Returns:: Dataloader over self.test_dataset, not shuffled.
Return type:: torch.utils.data.DataLoader

train_dataloader() → torch.utils.data.DataLoader[source]#

Load train dataset.

Returns:: Dataloader over self.train_dataset, shuffled.
Return type:: torch.utils.data.DataLoader

val_dataloader() → torch.utils.data.DataLoader[source]#

Load validation dataset.

Returns:: Dataloader over self.val_dataset, not shuffled.
Return type:: torch.utils.data.DataLoader

ar_steps_eval = 25#

ar_steps_train = 3#

batch_size = 4#

eval_split = 'test'#

load_single_member = False#

multiprocessing_context: str | None = None#

num_future_boundary_steps = 1#

num_future_forcing_steps = 1#

num_past_boundary_steps = 1#

num_past_forcing_steps = 1#

num_workers: int = 16#

test_dataset: WeatherDataset | None = None#

train_dataset: WeatherDataset | None = None#

val_dataset: WeatherDataset | None = 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, num_past_boundary_steps: int = 1, num_future_boundary_steps: int = 1, datastore_boundary: neural_lam.datastore.base.BaseDatastore | None = None, load_single_member: bool = False)[source]#

Bases: torch.utils.data.Dataset

Dataset class for weather data.

Loads and processes weather data from a given datastore, optionally with boundary forcing from a separate datastore. See __init__() for the full parameter list.

Parameters:

datastore (BaseDatastore) – Datastore providing access to state/forcing/static arrays.
split (str, optional) – Data split ("train", "val", or "test"). Default "train".
ar_steps (int, optional) – Number of autoregressive steps per training sample. Default 3.
num_past_forcing_steps (int, optional) – Past forcing window length i so that [t-i, ..., t] forcings are concatenated. Default 1.
num_future_forcing_steps (int, optional) – Future forcing window length j so that [t, ..., t+j] forcings are available. Default 1.
num_past_boundary_steps (int, optional) – Past window length for boundary forcing, in boundary time steps. Default 1.
num_future_boundary_steps (int, optional) – Future window length for boundary forcing. Default 1.
datastore_boundary (BaseDatastore, optional) – Datastore providing boundary forcing. If None, the boundary tensor is empty.
load_single_member (bool, optional) – If False and the datastore returns an ensemble of state realisations, treat each state ensemble member as an independent sample. If True, only ensemble member 0 is used. Default False.

Raises:

ValueError – If the datastore does not provide state data, if the configured ar_steps and forcing windows leave zero samples in split, or if the state/forcing dimension order does not match the datastore’s expected dimension order.
NotImplementedError – If datastore_boundary has ensemble forcing.

create_dataarray_from_tensor(tensor: torch.Tensor, time: datetime.datetime | list[datetime.datetime] | numpy.ndarray, category: str) → xarray.DataArray[source]#

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

INIT_STEPS = 2#

ar_steps = 3#

da_boundary_forcing = None#

da_forcing#

da_state#

datastore#

datastore_boundary = None#

load_single_member = False#

num_future_boundary_steps = 1#

num_future_forcing_steps = 1#

num_past_boundary_steps = 1#

num_past_forcing_steps = 1#

split = 'train'#