neural_lam.weather_dataset ========================== .. py:module:: neural_lam.weather_dataset .. autoapi-nested-parse:: Dataset helpers wrapping Neural-LAM datastores for PyTorch Lightning. Classes ------- .. autoapisummary:: neural_lam.weather_dataset.WeatherDataModule neural_lam.weather_dataset.WeatherDataset Module Contents --------------- .. py:class:: 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') Bases: :py:obj:`pytorch_lightning.LightningDataModule` DataModule for weather data. :param datastore: Datastore used for all splits. :type datastore: BaseDatastore :param ar_steps_train: Number of autoregressive steps for training batches. Default ``3``. :type ar_steps_train: int, optional :param ar_steps_eval: Number of autoregressive steps for validation/test batches. Default ``25``. :type ar_steps_eval: int, optional :param num_past_forcing_steps: Number of past forcing steps to include. Default ``1``. :type num_past_forcing_steps: int, optional :param num_future_forcing_steps: Number of future forcing steps to include. Default ``1``. :type num_future_forcing_steps: int, optional :param num_past_boundary_steps: Number of past boundary forcing steps to include. Default ``1``. :type num_past_boundary_steps: int, optional :param num_future_boundary_steps: Number of future boundary forcing steps to include. Default ``1``. :type num_future_boundary_steps: int, optional :param datastore_boundary: Separate datastore providing boundary forcing, used for all splits. If ``None``, the boundary tensor is empty. :type datastore_boundary: BaseDatastore, optional :param load_single_member: If ``True``, load only a single ensemble member per sample. Default ``False``. :type load_single_member: bool, optional :param batch_size: Mini-batch size for dataloaders. Default ``4``. :type batch_size: int, optional :param num_workers: Number of background workers per dataloader. Default ``16``. :type num_workers: int, optional :param eval_split: Dataset split to use for ``test_dataloader``. Default ``"test"``. :type eval_split: str, optional .. py:method:: setup(stage: str | None = None) -> None Instantiate datasets for the requested trainer stage. :param stage: Trainer stage identifier (``"fit"``/``"test"``/``None``). When ``None``, both the training split and the validation/test evaluation splits are prepared. :type stage: str or None, optional .. py:method:: test_dataloader() -> torch.utils.data.DataLoader Load test dataset. :returns: Dataloader over ``self.test_dataset``, not shuffled. :rtype: torch.utils.data.DataLoader .. py:method:: train_dataloader() -> torch.utils.data.DataLoader Load train dataset. :returns: Dataloader over ``self.train_dataset``, shuffled. :rtype: torch.utils.data.DataLoader .. py:method:: val_dataloader() -> torch.utils.data.DataLoader Load validation dataset. :returns: Dataloader over ``self.val_dataset``, not shuffled. :rtype: torch.utils.data.DataLoader .. py:attribute:: ar_steps_eval :value: 25 .. py:attribute:: ar_steps_train :value: 3 .. py:attribute:: batch_size :value: 4 .. py:attribute:: eval_split :value: 'test' .. py:attribute:: load_single_member :value: False .. py:attribute:: multiprocessing_context :type: str | None :value: None .. py:attribute:: num_future_boundary_steps :value: 1 .. py:attribute:: num_future_forcing_steps :value: 1 .. py:attribute:: num_past_boundary_steps :value: 1 .. py:attribute:: num_past_forcing_steps :value: 1 .. py:attribute:: num_workers :type: int :value: 16 .. py:attribute:: test_dataset :type: WeatherDataset | None :value: None .. py:attribute:: train_dataset :type: WeatherDataset | None :value: None .. py:attribute:: val_dataset :type: WeatherDataset | None :value: None .. py:class:: 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) Bases: :py:obj:`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 :meth:`__init__` for the full parameter list. :param datastore: Datastore providing access to state/forcing/static arrays. :type datastore: BaseDatastore :param split: Data split (``"train"``, ``"val"``, or ``"test"``). Default ``"train"``. :type split: str, optional :param ar_steps: Number of autoregressive steps per training sample. Default ``3``. :type ar_steps: int, optional :param num_past_forcing_steps: Past forcing window length ``i`` so that ``[t-i, ..., t]`` forcings are concatenated. Default ``1``. :type num_past_forcing_steps: int, optional :param num_future_forcing_steps: Future forcing window length ``j`` so that ``[t, ..., t+j]`` forcings are available. Default ``1``. :type num_future_forcing_steps: int, optional :param num_past_boundary_steps: Past window length for boundary forcing, in boundary time steps. Default ``1``. :type num_past_boundary_steps: int, optional :param num_future_boundary_steps: Future window length for boundary forcing. Default ``1``. :type num_future_boundary_steps: int, optional :param datastore_boundary: Datastore providing boundary forcing. If ``None``, the boundary tensor is empty. :type datastore_boundary: BaseDatastore, optional :param load_single_member: 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``. :type load_single_member: bool, optional :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. :raises NotImplementedError: If ``datastore_boundary`` has ensemble forcing. .. py:method:: create_dataarray_from_tensor(tensor: torch.Tensor, time: datetime.datetime | list[datetime.datetime] | numpy.ndarray, category: str) -> xarray.DataArray 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. :param 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. :type tensor: torch.Tensor :param time: The time or times of the tensor. :type time: datetime.datetime or list[datetime.datetime] :param category: The category of the tensor, either "state", "forcing" or "static". :type category: str :returns: **da** -- The constructed DataArray. :rtype: xr.DataArray .. py:attribute:: INIT_STEPS :value: 2 .. py:attribute:: ar_steps :value: 3 .. py:attribute:: da_boundary_forcing :value: None .. py:attribute:: da_forcing .. py:attribute:: da_state .. py:attribute:: datastore .. py:attribute:: datastore_boundary :value: None .. py:attribute:: load_single_member :value: False .. py:attribute:: num_future_boundary_steps :value: 1 .. py:attribute:: num_future_forcing_steps :value: 1 .. py:attribute:: num_past_boundary_steps :value: 1 .. py:attribute:: num_past_forcing_steps :value: 1 .. py:attribute:: split :value: 'train'