diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index fe2ff11ae..45bb18444 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,9 +36,16 @@ repos: - id: flake8 description: Check Python code for correctness, consistency and adherence to best practices additional_dependencies: [Flake8-pyproject] + - repo: https://github.com/econchick/interrogate + rev: 1.7.0 + hooks: + - id: interrogate + description: Ensure documentation coverage stays perfect + pass_filenames: false + args: ["--fail-under=100", "neural_lam"] - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.19.0 hooks: - - id: mypy + - id: mypy additional_dependencies: [types-PyYAML, types-Pillow, types-tqdm] description: Check for type errors diff --git a/neural_lam/config.py b/neural_lam/config.py index f4195ec36..dc4419805 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -1,7 +1,8 @@ # Standard library import dataclasses +import argparse from pathlib import Path -from typing import Dict, Union +from typing import Dict, Union, Tuple # Third-party import dataclass_wizard @@ -20,21 +21,29 @@ class DatastoreSelection: """ Configuration for selecting a datastore to use with neural-lam. - Attributes - ---------- - kind : str - The kind of datastore to use, currently `mdp` or `npyfilesmeps` are - implemented. - config_path : str - The path to the configuration file for the selected datastore, this is - assumed to be relative to the configuration file for neural-lam. + Args: + kind (str): The kind of datastore to use. Currently 'mdp' or + 'npyfilesmeps' are implemented. + config_path (str): The path to the configuration file for the selected + datastore, assumed to be relative to the neural-lam config file. """ kind: str def __post_init__(self): + """ + Validates the datastore kind against registered DATASTORES. + + Raises: + ValueError: If the provided kind is not found in the DATASTORES registry. + """ if self.kind not in DATASTORES: - raise ValueError(f"Datastore kind {self.kind} is not implemented") + available = ", ".join(DATASTORES.keys()) + raise ValueError( + f"Unknown datastore kind '{self.kind}'. " + f"Supported options are: {available}. " + "Please verify your configuration file." + ) config_path: str @@ -42,13 +51,11 @@ def __post_init__(self): @dataclasses.dataclass class ManualStateFeatureWeighting: """ - Configuration for weighting the state features in the loss function where - the weights are manually specified. + Configuration for manual weighting of state features in the loss function. - Attributes - ---------- - weights : Dict[str, float] - Manual weights for the state features. + Args: + weights (Dict[str, float]): Dictionary mapping feature names to + their respective manual weights. """ weights: Dict[str, float] @@ -57,8 +64,7 @@ class ManualStateFeatureWeighting: @dataclasses.dataclass class UniformFeatureWeighting: """ - Configuration for weighting the state features in the loss function where - all state features are weighted equally. + Configuration for equal weighting of all state features in the loss function. """ pass @@ -67,14 +73,11 @@ class UniformFeatureWeighting: @dataclasses.dataclass class OutputClamping: """ - Configuration for clamping the output of the model. - - Attributes - ---------- - lower : Dict[str, float] - The minimum value to clamp each output feature to. - upper : Dict[str, float] - The maximum value to clamp each output feature to. + Configuration for clamping the model's output values. + + Args: + lower (Dict[str, float]): Minimum values for each output feature. + upper (Dict[str, float]): Maximum values for each output feature. """ lower: Dict[str, float] = dataclasses.field(default_factory=dict) @@ -84,15 +87,14 @@ class OutputClamping: @dataclasses.dataclass class TrainingConfig: """ - Configuration related to training neural-lam - - Attributes - ---------- - state_feature_weighting : Union[ManualStateFeatureWeighting, - UnformFeatureWeighting] - The method to use for weighting the state features in the loss - function. Defaults to uniform weighting (`UnformFeatureWeighting`, i.e. - all features are weighted equally). + Configuration parameters related to the training process of neural-lam. + + Args: + state_feature_weighting (Union[ManualStateFeatureWeighting, UniformFeatureWeighting]): + The method used for weighting state features. Defaults to + UniformFeatureWeighting. + output_clamping (OutputClamping): Clamping configuration for model + predictions. """ state_feature_weighting: Union[ @@ -107,15 +109,12 @@ class TrainingConfig: @dataclasses.dataclass class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): """ - Dataclass for Neural-LAM configuration. This class is used to load and - store the configuration for using Neural-LAM. - - Attributes - ---------- - datastore : DatastoreSelection - The configuration for the datastore to use. - training : TrainingConfig - The configuration for training the model. + Primary configuration class for Neural-LAM. Handles loading and + storing all parameters required for model execution and training. + + Args: + datastore (DatastoreSelection): Selection and config path for the data source. + training (TrainingConfig): Training-specific parameters and loss weighting. """ datastore: DatastoreSelection @@ -123,66 +122,53 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard): class _(dataclass_wizard.JSONWizard.Meta): """ - Define the configuration class as a JSON wizard class. - - Together `tag_key` and `auto_assign_tags` enable that when a `Union` of - types are used for an attribute, the specific type to deserialize to - can be specified in the serialised data using the `tag_key` value. In - our case we call the tag key `__config_class__` to indicate to the - user that they should pick a dataclass describing configuration in - neural-lam. This Union-based selection allows us to support different - configuration attributes for different choices of methods for example - and is used when picking between different feature weighting methods in - the `TrainingConfig` class. `auto_assign_tags` is set to True to - automatically set that tag key (i.e. `__config_class__` in the config - file) should just be the class name of the dataclass to deserialize to. + Metadata for the JSON/YAML Wizard to handle configuration tagging. """ tag_key = "__config_class__" auto_assign_tags = True - # ensure that all parts of the loaded configuration match the - # dataclasses used - # TODO: this should be enabled once - # https://github.com/rnag/dataclass-wizard/issues/137 is fixed, but - # currently cannot be used together with `auto_assign_tags` due to a - # bug it seems - # raise_on_unknown_json_key = True class InvalidConfigError(Exception): + """Raised when the configuration file contains invalid keys or structure.""" pass def load_config_and_datastore( config_path: str, -) -> tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]: +) -> Tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]: """ - Load the neural-lam configuration and the datastore specified in the - configuration. - - Parameters - ---------- - config_path : str - Path to the Neural-LAM configuration file. - - Returns - ------- - tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]] - The Neural-LAM configuration and the loaded datastore. + Loads the Neural-LAM configuration and initializes the specified datastore. + + Args: + config_path (str): Path to the YAML configuration file. + + Returns: + Tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]: + A tuple containing the validated configuration object and the + initialized datastore instance. + + Raises: + InvalidConfigError: If the configuration file is missing required keys + or has an invalid structure. + FileNotFoundError: If the config_path does not exist. """ try: config = NeuralLAMConfig.from_yaml_file(config_path) except dataclass_wizard.errors.UnknownJSONKey as ex: raise InvalidConfigError( - "There was an error loading the configuration file at " - f"{config_path}. " + f"Failed to load configuration at '{config_path}'. " + "Ensure all keys match the NeuralLAMConfig schema." ) from ex - # datastore config is assumed to be relative to the config file + + # Resolve datastore path relative to the main config file datastore_config_path = ( Path(config_path).parent / config.datastore.config_path ) + datastore = init_datastore( - datastore_kind=config.datastore.kind, config_path=datastore_config_path + datastore_kind=config.datastore.kind, + config_path=datastore_config_path, ) - return config, datastore + return config, datastore \ No newline at end of file diff --git a/neural_lam/datastore/__init__.py b/neural_lam/datastore/__init__.py index dead77135..e81ad23cc 100644 --- a/neural_lam/datastore/__init__.py +++ b/neural_lam/datastore/__init__.py @@ -15,6 +15,26 @@ def init_datastore(datastore_kind, config_path): + """ + Instantiate a datastore based on its short-name identifier. + + Parameters + ---------- + datastore_kind : str + Key corresponding to one of :data:`DATASTORES`. + config_path : str | pathlib.Path + Path to the datastore-specific configuration file. + + Returns + ------- + BaseDatastore + Concrete datastore instance configured for ``config_path``. + + Raises + ------ + NotImplementedError + If ``datastore_kind`` is not registered. + """ DatastoreClass = DATASTORES.get(datastore_kind) if DatastoreClass is None: diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py index f6dddb007..d0cdeedbf 100644 --- a/neural_lam/datastore/base.py +++ b/neural_lam/datastore/base.py @@ -1,3 +1,5 @@ +"""Abstract base classes describing Neural-LAM datastore APIs.""" + # Standard library import abc import collections @@ -91,8 +93,10 @@ def config(self) -> collections.abc.Mapping: def step_length(self) -> timedelta: """The step length of the dataset as a time interval. - Returns: - timedelta: The step length as a datetime.timedelta object. + Returns + ------- + datetime.timedelta + The step length as a ``datetime.timedelta`` object. """ pass @@ -393,8 +397,10 @@ def state_feature_weights_values(self) -> List[float]: the loss function for each state variable (e.g. via the standard deviation of the 1-step differences of the state variables). - Returns: - List[float]: The weights for each state feature. + Returns + ------- + List[float] + The weights for each state feature. """ pass diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 114550e6d..7fba8305c 100644 --- a/neural_lam/datastore/mdp.py +++ b/neural_lam/datastore/mdp.py @@ -1,3 +1,5 @@ +"""Datastore implementation wrapping ``mllam-data-prep`` outputs.""" + # Standard library import copy import functools diff --git a/neural_lam/datastore/npyfilesmeps/__init__.py b/neural_lam/datastore/npyfilesmeps/__init__.py index 397a5075a..27cc356dc 100644 --- a/neural_lam/datastore/npyfilesmeps/__init__.py +++ b/neural_lam/datastore/npyfilesmeps/__init__.py @@ -1,2 +1,4 @@ +"""MEPS-specific datastore exposing numpy-based datasets.""" + # Local from .store import NpyFilesDatastoreMEPS # noqa diff --git a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py index 7dbcc7ef8..bf050b340 100644 --- a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py +++ b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py @@ -1,3 +1,5 @@ +"""Utilities for computing MEPS datastore standardization statistics.""" + # Standard library import os import subprocess @@ -18,7 +20,19 @@ class PaddedWeatherDataset(torch.utils.data.Dataset): + """Wrap :class:`WeatherDataset` to pad samples for distributed runners.""" + def __init__(self, base_dataset, world_size, batch_size): + """ + Parameters + ---------- + base_dataset : WeatherDataset + Dataset to pad. + world_size : int + Total number of distributed ranks participating. + batch_size : int + Per-rank batch size. + """ super().__init__() self.base_dataset = base_dataset self.world_size = world_size @@ -33,6 +47,7 @@ def __init__(self, base_dataset, world_size, batch_size): ) def __getitem__(self, idx): + """Return an item, repeating the final sample for padded indices.""" return self.base_dataset[ ( self.original_indices[-1] @@ -42,12 +57,15 @@ def __getitem__(self, idx): ] def __len__(self): + """Return the padded dataset length.""" return self.total_samples + self.padded_samples def get_original_indices(self): + """Return indices of the non-padded samples.""" return self.original_indices def get_original_window_indices(self, step_length): + """Return index mapping for sub-sampled windows at ``step_length``.""" step_int, _ = get_integer_time(step_length.total_seconds()) return [ i // step_int for i in range(len(self.original_indices) * step_int) @@ -55,10 +73,12 @@ def get_original_window_indices(self, step_length): def get_rank(): + """Return the rank inferred from SLURM or default to 0.""" return int(os.environ.get("SLURM_PROCID", 0)) def get_world_size(): + """Return the world size inferred from SLURM or default to 1.""" return int(os.environ.get("SLURM_NTASKS", 1)) @@ -101,6 +121,24 @@ def setup(rank, world_size): # pylint: disable=redefined-outer-name def save_stats( static_dir_path, means, squares, flux_means, flux_squares, filename_prefix ): + """ + Aggregate running statistics and persist them to ``static_dir_path``. + + Parameters + ---------- + static_dir_path : str or pathlib.Path + Directory where ``*.pt`` files should be written. + means : Sequence[torch.Tensor] + Batch-wise means with shape ``(N_batch, d_features)``. + squares : Sequence[torch.Tensor] + Batch-wise second moments with shape ``(N_batch, d_features)``. + flux_means : Sequence[torch.Tensor] + Optional flux means of shape ``(N_batch,)``. + flux_squares : Sequence[torch.Tensor] + Optional flux second moments of shape ``(N_batch,)``. + filename_prefix : str + Prefix (e.g., ``"parameter"`` or ``"diff"``) for saved tensors. + """ means = ( torch.stack(means) if len(means) > 1 else means[0] ) # (N_batch, d_features,) @@ -143,20 +181,20 @@ def main( datastore_config_path, batch_size, step_length, n_workers, distributed ): """ - Pre-compute parameter weights to be used in loss function + Pre-compute and persist standardization statistics from the datastore. - Arguments - --------- - datastore_config_path : str - Path to datastore config file + Parameters + ---------- + datastore_config_path : str or pathlib.Path + Path to the MEPS datastore configuration file. batch_size : int - Batch size when iterating over the dataset + Batch size used while iterating through the dataset. step_length : datetime.timedelta - Step length to consider single time step + Temporal sampling interval for the difference statistics. n_workers : int - Number of workers in data loader + Number of dataloader workers. distributed : bool - Run the script in distributed + If ``True``, run using torch.distributed with SLURM settings. """ rank = get_rank() @@ -382,6 +420,7 @@ def main( def cli(): + """Parse CLI arguments and trigger :func:`main`.""" parser = ArgumentParser(description="Training arguments") parser.add_argument( "--datastore_config_path", diff --git a/neural_lam/datastore/npyfilesmeps/config.py b/neural_lam/datastore/npyfilesmeps/config.py index 1d36b6fe0..b5f015b6b 100644 --- a/neural_lam/datastore/npyfilesmeps/config.py +++ b/neural_lam/datastore/npyfilesmeps/config.py @@ -1,3 +1,5 @@ +"""Dataclasses describing the MEPS numpy-file datastore configuration.""" + # Standard library from dataclasses import dataclass, field from datetime import timedelta diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py index f65fd8aef..402c4295c 100644 --- a/neural_lam/datastore/npyfilesmeps/store.py +++ b/neural_lam/datastore/npyfilesmeps/store.py @@ -35,6 +35,7 @@ def _load_np(fp, add_feature_dim, feature_dim_mask=None): + """Load an ``.npy`` file and optionally expand/mask the feature axis.""" arr = np.load(fp) if add_feature_dim: arr = arr[..., np.newaxis] @@ -44,25 +45,25 @@ def _load_np(fp, add_feature_dim, feature_dim_mask=None): class NpyFilesDatastoreMEPS(BaseRegularGridDatastore): - __doc__ = f""" + """ 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 - '{STATE_FILENAME_FORMAT}' + 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}`: + - ``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}`: + - ``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}`: + - ``OPEN_WATER_FILENAME_FORMAT``: The open water fraction at `time`. The dimensions of the array are `[y, x]`. @@ -551,6 +552,19 @@ def _get_analysis_times(self, split) -> List[np.datetime64]: return sorted(times) def _calc_datetime_forcing_features(self, da_time: xr.DataArray): + """ + Compute sinusoidal encodings of hour-of-day and day-of-year. + + Parameters + ---------- + da_time : xr.DataArray + Time coordinate with dimension ``time``. + + Returns + ------- + xr.DataArray + Normalized sine/cosine features with dims ``("feature",)``. + """ da_hour_angle = da_time.dt.hour / 12 * np.pi da_year_angle = da_time.dt.dayofyear / 365 * 2 * np.pi @@ -574,6 +588,7 @@ def _calc_datetime_forcing_features(self, da_time: xr.DataArray): return da_datetime_forcing def get_vars_units(self, category: str) -> List[str]: + """Return unit strings for the variables in ``category``.""" if category == "state": return self.config.dataset.var_units elif category == "forcing": @@ -591,6 +606,7 @@ def get_vars_units(self, category: str) -> List[str]: raise NotImplementedError(f"Category {category} not supported") def get_vars_names(self, category: str) -> List[str]: + """Return canonical short names for the variables in ``category``.""" if category == "state": return self.config.dataset.var_names elif category == "forcing": @@ -610,6 +626,7 @@ def get_vars_names(self, category: str) -> List[str]: raise NotImplementedError(f"Category {category} not supported") def get_vars_long_names(self, category: str) -> List[str]: + """Return descriptive names for the variables in ``category``.""" if category == "state": return self.config.dataset.var_longnames else: @@ -617,6 +634,7 @@ def get_vars_long_names(self, category: str) -> List[str]: return self.get_vars_names(category=category) def get_num_data_vars(self, category: str) -> int: + """Return the number of variables available in ``category``.""" return len(self.get_vars_names(category=category)) def get_xy(self, category: str, stacked: bool) -> np.ndarray: @@ -732,6 +750,7 @@ def get_standardization_dataarray(self, category: str) -> xr.Dataset: """ def load_pickled_tensor(fn): + """Load a serialized tensor from ``static`` and convert to numpy.""" return torch.load( self.root_path / "static" / fn, weights_only=True ).numpy() diff --git a/neural_lam/datastore/plot_example.py b/neural_lam/datastore/plot_example.py index 4f61ac7e7..912478db4 100644 --- a/neural_lam/datastore/plot_example.py +++ b/neural_lam/datastore/plot_example.py @@ -1,3 +1,5 @@ +"""CLI helper to plot slices from datastores for manual inspection.""" + # Third-party import matplotlib.pyplot as plt @@ -93,6 +95,7 @@ def plot_example_from_datastore( import argparse def _parse_dict(arg_str): + """Parse ``key=value`` CLI arguments into typed dictionary entries.""" key, value = arg_str.split("=") for op in [int, float]: try: diff --git a/neural_lam/plot_graph.py b/neural_lam/plot_graph.py index 6ed7d0268..c3cad3ea0 100644 --- a/neural_lam/plot_graph.py +++ b/neural_lam/plot_graph.py @@ -232,7 +232,7 @@ def plot_graph( def main(): """Plot graph structure in 3D using plotly.""" parser = ArgumentParser( - description="Plot graph", + description="Visualize Neural-LAM graph structure in 3D using Plotly.", formatter_class=ArgumentDefaultsHelpFormatter, ) parser.add_argument( diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index 5547fdd4e..eaf9f8f50 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -1,7 +1,7 @@ # Standard library import datetime import warnings -from typing import Union +from typing import Dict, List, Optional, Tuple, Union, Generator # Third-party import numpy as np @@ -15,35 +15,20 @@ class WeatherDataset(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. - 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 is False, - so all members are used when available. - standardize : bool, optional - Whether to standardize the data. Default is True. + """ + Dataset class for loading and processing weather data from a given datastore. + + Args: + datastore (BaseDatastore): The datastore to load the data from (e.g., mdp). + split (str, optional): The data split to use ("train", "val", or "test"). + Defaults to "train". + ar_steps (int, optional): Number of autoregressive steps. Defaults to 3. + num_past_forcing_steps (int, optional): Number of past time steps to + include in forcing input. Defaults to 1. + num_future_forcing_steps (int, optional): Number of future time steps to + include in forcing input. Defaults to 1. + standardize (bool, optional): Whether to standardize the data. + Defaults to True. """ def __init__( @@ -55,7 +40,7 @@ def __init__( num_future_forcing_steps: int = 1, load_single_member: bool = False, standardize: bool = True, - ): + ) -> None: super().__init__() self.split = split @@ -71,31 +56,12 @@ def __init__( self.da_forcing = self.datastore.get_dataarray( category="forcing", split=self.split ) - if self.da_state is None: - raise ValueError( - "The datastore must provide state data for the WeatherDataset." - ) - if self.datastore.is_ensemble and self.load_single_member: - warnings.warn( - "only using first ensemble member, so dataset size is " - "effectively reduced by the number of ensemble members " - f"({self.da_state.ensemble_member.size})", - UserWarning, - stacklevel=2, - ) - - # check that with the provided data-arrays and ar_steps that we have a - # non-zero amount of samples if self.__len__() <= 0 and self.da_state is not None: raise ValueError( - "The provided datastore only provides " - f"{len(self.da_state.time)} total time steps, which is too few " - "to create a single sample for the WeatherDataset " - f"configuration used in the `{split}` split. You could try " - "either reducing the number of autoregressive steps " - "(`ar_steps`) and/or the forcing window size " - "(`num_past_forcing_steps` and `num_future_forcing_steps`)" + f"The provided datastore only provides {len(self.da_state.time)} " + f"total time steps, which is too few for the `{split}` split. " + "Try reducing ar_steps or forcing window size." ) # Check the dimensions and their ordering @@ -104,34 +70,24 @@ def __init__( parts["forcing"] = self.da_forcing for part, da in parts.items(): - if da is not None: - expected_dim_order = self.datastore.expected_dim_order( - category=part + expected_dim_order = self.datastore.expected_dim_order(category=part) + if da is not None and da.dims != expected_dim_order: + raise ValueError( + f"Dimension order of `{part}` ({da.dims}) does not match " + f"expected ({expected_dim_order})." ) - if da.dims != expected_dim_order: - raise ValueError( - f"The dimension order of the `{part}` data ({da.dims}) " - f"does not match the expected dimension order " - f"({expected_dim_order}). Maybe you forgot to " - "transpose the data in `BaseDatastore.get_dataarray`?" - ) - - # Set up for standardization - # TODO: This will become part of ar_model.py soon! + self.standardize = standardize if standardize: self.ds_state_stats = self.datastore.get_standardization_dataarray( category="state" ) - self.da_state_mean = self.ds_state_stats.state_mean self.da_state_std = self.ds_state_stats.state_std if self.da_forcing is not None: - self.ds_forcing_stats = ( - self.datastore.get_standardization_dataarray( - category="forcing" - ) + self.ds_forcing_stats = self.datastore.get_standardization_dataarray( + category="forcing" ) self.da_forcing_mean = self.ds_forcing_stats.forcing_mean self.da_forcing_std = self.ds_forcing_stats.forcing_std @@ -142,15 +98,23 @@ def __init__( self.state_std_safe = self._compute_std_safe( self.da_state_std, "state" ) + self.forcing_std_safe = ( + self._compute_std_safe(self.da_forcing_std, "forcing") + if self.da_forcing_std is not None else None + ) - if self.da_forcing_std is not None: - self.forcing_std_safe = self._compute_std_safe( - self.da_forcing_std, "forcing" - ) - else: - self.forcing_std_safe = None + def _compute_std_safe(self, std: xr.DataArray, feature: str) -> xr.DataArray: + """ + Ensures standard deviation is above machine epsilon to avoid division by zero. + + Args: + std (xr.DataArray): The standard deviation array. + feature (str): Name of the feature category for logging. - def _compute_std_safe(self, std: xr.DataArray, feature: str): + Returns: + xr.DataArray: The standard deviation array with near-zero values + replaced by epsilon. + """ eps = np.finfo(std.dtype).eps if bool((std <= eps).any()): logger.warning( @@ -159,272 +123,127 @@ def _compute_std_safe(self, std: xr.DataArray, feature: str): ) return std.where(std > eps, other=eps) - def __len__(self): + def __len__(self) -> int: + """ + Calculates the total number of samples available in the dataset. + + Returns: + int: Number of samples. + """ if self.datastore.is_forecast: - # for now we simply create a single sample for each analysis time - # and then take the first (2 + ar_steps) forecast times. - # If the datastore returns an ensemble of state realisations and - # `load_single_member=False`, each ensemble member is exposed as an - # independent sample by scaling the base dataset length below. - - # check that there are enough forecast steps available to create - # samples given the number of autoregressive steps requested + if self.datastore.is_ensemble: + warnings.warn( + "Only using first ensemble member; dataset size effectively " + "reduced.", UserWarning + ) + n_forecast_steps = self.da_state.elapsed_forecast_duration.size if n_forecast_steps < 2 + self.ar_steps: raise ValueError( - "The number of forecast steps available " - f"({n_forecast_steps}) is less than the required " - f"2+ar_steps (2+{self.ar_steps}={2 + self.ar_steps}) for " - "creating a sample with initial and target states." + f"Forecast steps ({n_forecast_steps}) < required " + f"(2 + {self.ar_steps})." ) + return self.da_state.analysis_time.size + + return ( + len(self.da_state.time) + - self.ar_steps + - max(2, self.num_past_forcing_steps) + - self.num_future_forcing_steps + ) - base_len = self.da_state.analysis_time.size - else: - # Calculate the number of samples in the dataset n_samples = total - # time steps - (autoregressive steps + past forcing + future - # forcing) - #: - # Where: - # - total time steps: len(self.da_state.time) - # - autoregressive steps: self.ar_steps - # - past forcing: max(2, self.num_past_forcing_steps) (at least 2 - # time steps are required for the initial state) - # - future forcing: self.num_future_forcing_steps - base_len = ( - len(self.da_state.time) - - self.ar_steps - - max(2, self.num_past_forcing_steps) - - self.num_future_forcing_steps - ) - if self.datastore.is_ensemble and not self.load_single_member: - return base_len * self.da_state.ensemble_member.size - return base_len - - def _slice_state_time(self, da_state, idx, n_steps: int): + def _slice_state_time(self, da_state: xr.DataArray, idx: int, n_steps: int) -> xr.DataArray: """ - Produce a time slice of the given dataarray `da_state` (state) starting - at `idx` and with `n_steps` steps. An `offset`is calculated based on the - `num_past_forcing_steps` class attribute. `Offset` is used to offset the - start of the sample, to assert that enough previous time steps are - available for the 2 initial states and any corresponding forcings - (calculated in `_slice_forcing_time`). - - Parameters - ---------- - da_state : xr.DataArray - The dataarray to slice. This is expected to have a `time` dimension - if the datastore is providing analysis only data, and a - `analysis_time` and `elapsed_forecast_duration` dimensions if the - datastore is providing forecast data. - idx : int - The index of the time step to start the sample from. - n_steps : int - The number of time steps to include in the sample. - - Returns - ------- - da_sliced : xr.DataArray - The sliced dataarray with dims ('time', 'grid_index', - 'state_feature'). + Produces a time slice of the state data. + + Args: + da_state (xr.DataArray): DataArray to slice. + idx (int): Index of the starting time step. + n_steps (int): Number of steps to include. + + Returns: + xr.DataArray: Sliced DataArray with dims ('time', 'grid_index', 'state_feature'). """ - # The current implementation requires at least 2 time steps for the - # initial state (see GraphCast). init_steps = 2 - # slice the dataarray to include the required number of time steps if self.datastore.is_forecast: start_idx = max(0, self.num_past_forcing_steps - init_steps) end_idx = max(init_steps, self.num_past_forcing_steps) + n_steps - # this implies that the data will have both `analysis_time` and - # `elapsed_forecast_duration` dimensions for forecasts. We for now - # simply select a analysis time and the first `n_steps` forecast - # times (given no offset). Note that this means that we get one - # sample per forecast, always starting at forecast time 2. da_sliced = da_state.isel( analysis_time=idx, elapsed_forecast_duration=slice(start_idx, end_idx), ) - # create a new time dimension so that the produced sample has a - # `time` dimension, similarly to the analysis only data da_sliced["time"] = ( da_sliced.analysis_time + da_sliced.elapsed_forecast_duration ) - da_sliced = da_sliced.swap_dims( - {"elapsed_forecast_duration": "time"} - ) + da_sliced = da_sliced.swap_dims({"elapsed_forecast_duration": "time"}) else: - # For analysis data we slice the time dimension directly. The offset - # is only relevant for the very first (and last) samples in the - # dataset. start_idx = idx + max(0, self.num_past_forcing_steps - init_steps) - end_idx = ( - idx + max(init_steps, self.num_past_forcing_steps) + n_steps - ) + end_idx = idx + max(init_steps, self.num_past_forcing_steps) + n_steps da_sliced = da_state.isel(time=slice(start_idx, end_idx)) return da_sliced - def _slice_forcing_time(self, da_forcing, idx, n_steps: int): + def _slice_forcing_time(self, da_forcing: xr.DataArray, idx: int, n_steps: int) -> xr.DataArray: """ - Produce a time slice of the given dataarray `da_forcing` (forcing) - starting at `idx` and with `n_steps` steps. An `offset` is calculated - based on the `num_past_forcing_steps` class attribute. It is used to - offset the start of the sample, to ensure that enough previous time - steps are available for the forcing data. The forcing data is windowed - around the current autoregressive time step to include the past and - future forcings. - - Parameters - ---------- - da_forcing : xr.DataArray - The forcing dataarray to slice. This is expected to have a `time` - dimension if the datastore is providing analysis only data, and a - `analysis_time` and `elapsed_forecast_duration` dimensions if the - datastore is providing forecast data. - idx : int - The index of the time step to start the sample from. - n_steps : int - The number of time steps to include in the sample. - - Returns - ------- - da_concat : xr.DataArray - The sliced dataarray with dims ('time', 'grid_index', - 'window', 'forcing_feature'). + Produces a windowed time slice of the forcing data. + + Args: + da_forcing (xr.DataArray): Forcing DataArray to slice. + idx (int): Starting time step index. + n_steps (int): Number of steps to include. + + Returns: + xr.DataArray: Sliced array with dims ('time', 'grid_index', 'window', 'forcing_feature'). """ - # The current implementation requires at least 2 time steps for the - # initial state (see GraphCast). The forcing data is windowed around the - # current autregressive time step. The two `init_steps` can also be used - # as past forcings. init_steps = 2 da_list = [] + offset = (idx + max(init_steps, self.num_past_forcing_steps)) if not self.datastore.is_forecast else max(init_steps, self.num_past_forcing_steps) - if self.datastore.is_forecast: - # This implies that the data will have both `analysis_time` and - # `elapsed_forecast_duration` dimensions for forecasts. We for now - # simply select an analysis time and the first `n_steps` forecast - # times (given no offset). Note that this means that we get one - # sample per forecast. - # Add a 'time' dimension using the actual forecast times - offset = max(init_steps, self.num_past_forcing_steps) - for step in range(n_steps): - start_idx = offset + step - self.num_past_forcing_steps - end_idx = offset + step + self.num_future_forcing_steps - - current_time = ( - da_forcing.analysis_time[idx] - + da_forcing.elapsed_forecast_duration[offset + step] - ) + for step in range(n_steps): + start_idx = offset + step - self.num_past_forcing_steps + end_idx = offset + step + self.num_future_forcing_steps + if self.datastore.is_forecast: + current_time = da_forcing.analysis_time[idx] + da_forcing.elapsed_forecast_duration[offset + step] da_sliced = da_forcing.isel( analysis_time=idx, elapsed_forecast_duration=slice(start_idx, end_idx + 1), - ) - - da_sliced = da_sliced.rename( - {"elapsed_forecast_duration": "window"} - ) - - # Assign the 'window' coordinate to be relative positions - da_sliced = da_sliced.assign_coords( - window=np.arange(len(da_sliced.window)) - ) - - da_sliced = da_sliced.expand_dims( - dim={"time": [current_time.values]} - ) - - da_list.append(da_sliced) - - # Concatenate the list of DataArrays along the 'time' dimension - da_concat = xr.concat(da_list, dim="time") - - else: - # For analysis data, we slice the time dimension directly. The - # offset is only relevant for the very first (and last) samples in - # the dataset. - offset = idx + max(init_steps, self.num_past_forcing_steps) - for step in range(n_steps): - start_idx = offset + step - self.num_past_forcing_steps - end_idx = offset + step + self.num_future_forcing_steps - - # Slice the data over the desired time window - da_sliced = da_forcing.isel(time=slice(start_idx, end_idx + 1)) - - da_sliced = da_sliced.rename({"time": "window"}) - - # Assign the 'window' coordinate to be relative positions - da_sliced = da_sliced.assign_coords( - window=np.arange(len(da_sliced.window)) - ) - - # Add a 'time' dimension to keep track of steps using actual - # time coordinates + ).rename({"elapsed_forecast_duration": "window"}) + else: current_time = da_forcing.time[offset + step] - da_sliced = da_sliced.expand_dims( - dim={"time": [current_time.values]} - ) + da_sliced = da_forcing.isel(time=slice(start_idx, end_idx + 1)).rename({"time": "window"}) - da_list.append(da_sliced) + da_sliced = da_sliced.assign_coords(window=np.arange(len(da_sliced.window))) + da_sliced = da_sliced.expand_dims(dim={"time": [current_time.values]}) + da_list.append(da_sliced) - # Concatenate the list of DataArrays along the 'time' dimension - da_concat = xr.concat(da_list, dim="time") + return xr.concat(da_list, dim="time") - return da_concat - - def _build_item_dataarrays(self, idx): - """ - Create the dataarrays for the initial states, target states and forcing - data for the sample at index `idx`. - - Parameters - ---------- - idx : int - The index of the sample to create the dataarrays for. - - Returns - ------- - da_init_states : xr.DataArray - The dataarray for the initial states. - da_target_states : xr.DataArray - The dataarray for the target states. - da_forcing_windowed : xr.DataArray - The dataarray for the forcing data, windowed for the sample. - da_target_times : xr.DataArray - The dataarray for the target times. + def _build_item_dataarrays(self, idx: int) -> Tuple[xr.DataArray, xr.DataArray, xr.DataArray, xr.DataArray]: """ - # Handle indexing over state ensemble members. If forcing data also - # has an ensemble dimension, we select the same member below. - sample_idx = idx - i_ensemble = 0 + Builds the underlying DataArrays for a single sample. + + Args: + idx (int): The sample index. + Returns: + Tuple[xr.DataArray, xr.DataArray, xr.DataArray, xr.DataArray]: + Initial states, target states, forcing data, and target times. + """ if self.datastore.is_ensemble: - n_ensemble_members = self.da_state.ensemble_member.size - if not self.load_single_member: - sample_idx, i_ensemble = divmod(idx, n_ensemble_members) - da_state = self.da_state.isel(ensemble_member=i_ensemble) + warnings.warn("Only ensemble member 0 implemented.") + da_state = self.da_state.isel(ensemble_member=0) else: da_state = self.da_state - if self.da_forcing is not None: - if self.datastore.has_ensemble_forcing: - da_forcing = self.da_forcing.isel(ensemble_member=i_ensemble) - else: - da_forcing = self.da_forcing - else: - da_forcing = None + da_forcing = self.da_forcing + if da_forcing is not None and "ensemble_member" in da_forcing.dims: + raise NotImplementedError("Ensemble member not supported for forcing.") - # handle time sampling in a way that is compatible with both analysis - # and forecast data - da_state = self._slice_state_time( - da_state=da_state, idx=sample_idx, n_steps=self.ar_steps - ) - if da_forcing is not None: - da_forcing_windowed = self._slice_forcing_time( - da_forcing=da_forcing, idx=sample_idx, n_steps=self.ar_steps - ) + da_state = self._slice_state_time(da_state=da_state, idx=idx, n_steps=self.ar_steps) + da_forcing_windowed = self._slice_forcing_time(da_forcing=da_forcing, idx=idx, n_steps=self.ar_steps) if da_forcing is not None else None - # load the data into memory da_state.load() - if da_forcing is not None: + if da_forcing_windowed is not None: da_forcing_windowed.load() da_init_states = da_state.isel(time=slice(0, 2)) @@ -432,202 +251,106 @@ def _build_item_dataarrays(self, idx): da_target_times = da_target_states.time if self.standardize: - da_init_states = ( - da_init_states - self.da_state_mean - ) / self.state_std_safe - da_target_states = ( - da_target_states - self.da_state_mean - ) / self.state_std_safe - - if da_forcing is not None: - # XXX: Here we implicitly assume that the last dimension of the - # forcing data is the forcing feature dimension. To standardize - # on `.device` we need a different implementation. (e.g. a - # tensor with repeated means and stds for each "windowed" time.) - da_forcing_windowed = ( - da_forcing_windowed - self.da_forcing_mean - ) / self.forcing_std_safe - - if da_forcing is not None: - # stack the `forcing_feature` and `window_sample` dimensions into a - # single `forcing_feature` dimension + da_init_states = (da_init_states - self.da_state_mean) / self.state_std_safe + da_target_states = (da_target_states - self.da_state_mean) / self.state_std_safe + if da_forcing_windowed is not None: + da_forcing_windowed = (da_forcing_windowed - self.da_forcing_mean) / self.forcing_std_safe + + if da_forcing_windowed is not None: da_forcing_windowed = da_forcing_windowed.stack( forcing_feature_windowed=("forcing_feature", "window") ) else: - # create an empty forcing tensor with the right shape da_forcing_windowed = xr.DataArray( - data=np.empty( - (self.ar_steps, da_state.grid_index.size, 0), - ), + data=np.empty((self.ar_steps, da_state.grid_index.size, 0)), dims=("time", "grid_index", "forcing_feature"), - coords={ - "time": da_target_times, - "grid_index": da_state.grid_index, - "forcing_feature": [], - }, + coords={"time": da_target_times, "grid_index": da_state.grid_index, "forcing_feature": []}, ) - return ( - da_init_states, - da_target_states, - da_forcing_windowed, - da_target_times, - ) + return da_init_states, da_target_states, da_forcing_windowed, da_target_times - def __getitem__(self, idx): + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ - Return a single training sample, which consists of the initial states, - target states, forcing and batch times. - - The implementation currently uses xarray.DataArray objects for the - standardization (scaling to mean 0.0 and standard deviation of 1.0) so - that we can make us of xarray's broadcasting capabilities. This makes - it possible to standardization with both global means, but also for - example where a grid-point mean has been computed. This code will have - to be replace if standardization is to be done on the GPU to handle - different shapes of the standardization. - - Parameters - ---------- - idx : int - The index of the sample to return, this will refer to the time of - the initial state. - - Returns - ------- - init_states : TrainingSample - A training sample object containing the initial states, target - states, forcing and batch times. The batch times are the times of - the target steps. + Returns a single training sample. + + Args: + idx (int): Index of the sample to return. + Returns: + Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + - init_states: (2, N_grid, d_features) + - target_states: (ar_steps, N_grid, d_features) + - forcing: (ar_steps, N_grid, d_windowed_forcing) + - target_times: (ar_steps,) """ - ( - da_init_states, - da_target_states, - da_forcing_windowed, - da_target_times, - ) = self._build_item_dataarrays(idx=idx) + (da_init, da_target, da_forcing, da_times) = self._build_item_dataarrays(idx=idx) tensor_dtype = torch.float32 - - init_states = torch.tensor(da_init_states.values, dtype=tensor_dtype) - target_states = torch.tensor( - da_target_states.values, dtype=tensor_dtype - ) - + init_states = torch.tensor(da_init.values, dtype=tensor_dtype) + target_states = torch.tensor(da_target.values, dtype=tensor_dtype) target_times = torch.tensor( - da_target_times.astype("datetime64[ns]").astype("int64").values, - dtype=torch.int64, + da_times.astype("datetime64[ns]").astype("int64").values, dtype=torch.int64 ) - - forcing = torch.tensor(da_forcing_windowed.values, dtype=tensor_dtype) - - # init_states: (2, N_grid, d_features) - # target_states: (ar_steps, N_grid, d_features) - # forcing: (ar_steps, N_grid, d_windowed_forcing) - # target_times: (ar_steps,) + forcing = torch.tensor(da_forcing.values, dtype=tensor_dtype) return init_states, target_states, forcing, target_times - def __iter__(self): - """ - Convenience method to iterate over the dataset. - - This isn't used by pytorch DataLoader which itself implements an - iterator that uses Dataset.__getitem__ and Dataset.__len__. - - """ + def __iter__(self) -> Generator: + """Convenience method to iterate over the dataset.""" for i in range(len(self)): yield self[i] def create_dataarray_from_tensor( self, tensor: torch.Tensor, - time: Union[datetime.datetime, list[datetime.datetime]], + time: Union[datetime.datetime, List[datetime.datetime]], category: str, - ): + ) -> xr.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. - - 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 : xr.DataArray - The constructed DataArray. + Constructs an xarray.DataArray from a torch.Tensor. + + Args: + tensor (torch.Tensor): Tensor with shape (grid_index, feature) or + (time, grid_index, feature). + time (Union[datetime.datetime, List[datetime.datetime]]): Time + coordinates for the tensor. + category (str): Tensor category ("state", "forcing", or "static"). + + Returns: + xr.DataArray: Constructed DataArray with spatial coordinates. + + Raises: + ValueError: If tensor dimensions and time list length mismatch. """ - def _is_listlike(obj): - # match list, tuple, numpy array return hasattr(obj, "__iter__") and not isinstance(obj, str) add_time_as_dim = False if len(tensor.shape) == 2: dims = ["grid_index", f"{category}_feature"] if _is_listlike(time): - raise ValueError( - "Expected a single time for a 2D tensor with assumed " - "dimensions (grid_index, {category}_feature), but got " - f"{len(time)} times" # type: ignore - ) + raise ValueError(f"Expected single time for 2D tensor, got {len(time)}.") elif len(tensor.shape) == 3: add_time_as_dim = True dims = ["time", "grid_index", f"{category}_feature"] if not _is_listlike(time): - raise ValueError( - "Expected a list of times for a 3D tensor with assumed " - "dimensions (time, grid_index, {category}_feature), but " - "got a single time" - ) + raise ValueError("Expected list of times for 3D tensor.") else: - raise ValueError( - "Expected tensor to have 2 or 3 dimensions, but got " - f"{len(tensor.shape)}" - ) + raise ValueError(f"Expected 2 or 3 dims, got {len(tensor.shape)}.") da_datastore_state = getattr(self, f"da_{category}") - da_grid_index = da_datastore_state.grid_index - da_state_feature = da_datastore_state.state_feature - coords = { - f"{category}_feature": da_state_feature, - "grid_index": da_grid_index, + f"{category}_feature": da_datastore_state.state_feature, + "grid_index": da_datastore_state.grid_index, } if add_time_as_dim: coords["time"] = time - da = xr.DataArray( - tensor.cpu().numpy(), - dims=dims, - coords=coords, - ) + da = xr.DataArray(tensor.cpu().numpy(), dims=dims, coords=coords) - for grid_coord in ["x", "y"]: - if ( - grid_coord in da_datastore_state.coords - and grid_coord not in da.coords - ): - da.coords[grid_coord] = da_datastore_state[grid_coord] + for coord in ["x", "y"]: + if coord in da_datastore_state.coords and coord not in da.coords: + da.coords[coord] = da_datastore_state[coord] if not add_time_as_dim: da.coords["time"] = time @@ -636,7 +359,20 @@ def _is_listlike(obj): class WeatherDataModule(pl.LightningDataModule): - """DataModule for weather data.""" + """ + DataModule for organizing training, validation, and test weather datasets. + + Args: + datastore (BaseDatastore): Datastore instance. + ar_steps_train (int): Autoregressive steps for training. Defaults to 3. + ar_steps_eval (int): Autoregressive steps for evaluation. Defaults to 25. + standardize (bool): Whether to standardize data. Defaults to True. + num_past_forcing_steps (int): Past forcing steps. Defaults to 1. + num_future_forcing_steps (int): Future forcing steps. Defaults to 1. + batch_size (int): Training batch size. Defaults to 4. + num_workers (int): Number of worker processes for loading. Defaults to 16. + eval_split (str): Split to use for evaluation. Defaults to "test". + """ def __init__( self, @@ -650,7 +386,7 @@ def __init__( batch_size: int = 4, num_workers: int = 16, eval_split: str = "test", - ): + ) -> None: super().__init__() self._datastore = datastore self.num_past_forcing_steps = num_past_forcing_steps @@ -660,81 +396,53 @@ def __init__( self.standardize = standardize self.load_single_member = load_single_member self.batch_size = batch_size - self.num_workers: int = num_workers + self.num_workers = num_workers self.train_dataset = None self.val_dataset = None self.test_dataset = None - self.multiprocessing_context: Union[str, None] = None self.eval_split = eval_split - if num_workers > 0: - # default to spawn for now, as the default on linux "fork" hangs - # when using dask (which the npyfilesmeps datastore uses) - self.multiprocessing_context = "spawn" - - def setup(self, stage=None): - if stage == "fit" or stage is None: - self.train_dataset = WeatherDataset( - datastore=self._datastore, - split="train", - ar_steps=self.ar_steps_train, - standardize=self.standardize, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, - ) - self.val_dataset = WeatherDataset( - datastore=self._datastore, - split="val", - ar_steps=self.ar_steps_eval, - standardize=self.standardize, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, - ) + self.multiprocessing_context = "spawn" if num_workers > 0 else None - if stage == "test" or stage is None: - self.test_dataset = WeatherDataset( - datastore=self._datastore, - split=self.eval_split, - ar_steps=self.ar_steps_eval, - standardize=self.standardize, - num_past_forcing_steps=self.num_past_forcing_steps, - num_future_forcing_steps=self.num_future_forcing_steps, - load_single_member=self.load_single_member, - ) + def setup(self, stage: Optional[str] = None) -> None: + """ + Initializes datasets for specified stage. - def train_dataloader(self): - """Load train dataset.""" + Args: + stage (str, optional): "fit", "test", or None. + """ + common_kwargs = { + "datastore": self._datastore, + "standardize": self.standardize, + "num_past_forcing_steps": self.num_past_forcing_steps, + "num_future_forcing_steps": self.num_future_forcing_steps, + } + + if stage in ("fit", None): + self.train_dataset = WeatherDataset(split="train", ar_steps=self.ar_steps_train, **common_kwargs) + self.val_dataset = WeatherDataset(split="val", ar_steps=self.ar_steps_eval, **common_kwargs) + + if stage in ("test", None): + self.test_dataset = WeatherDataset(split=self.eval_split, ar_steps=self.ar_steps_eval, **common_kwargs) + + def _get_dataloader(self, dataset, shuffle=False): return torch.utils.data.DataLoader( - self.train_dataset, + dataset, batch_size=self.batch_size, num_workers=self.num_workers, - shuffle=True, + shuffle=shuffle, multiprocessing_context=self.multiprocessing_context, persistent_workers=self.num_workers > 0, pin_memory=torch.cuda.is_available(), ) + def train_dataloader(self): + """Returns the training dataloader.""" + return self._get_dataloader(self.train_dataset, shuffle=True) + def val_dataloader(self): - """Load validation dataset.""" - return torch.utils.data.DataLoader( - self.val_dataset, - batch_size=self.batch_size, - num_workers=self.num_workers, - shuffle=False, - multiprocessing_context=self.multiprocessing_context, - persistent_workers=self.num_workers > 0, - pin_memory=torch.cuda.is_available(), - ) + """Returns the validation dataloader.""" + return self._get_dataloader(self.val_dataset) def test_dataloader(self): - """Load test dataset.""" - return torch.utils.data.DataLoader( - self.test_dataset, - batch_size=self.batch_size, - num_workers=self.num_workers, - shuffle=False, - multiprocessing_context=self.multiprocessing_context, - persistent_workers=self.num_workers > 0, - pin_memory=torch.cuda.is_available(), - ) + """Returns the test dataloader.""" + return self._get_dataloader(self.test_dataset) \ No newline at end of file