docs(dataset): migrate WeatherDataset to Google-style and add type hints

.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

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,35 +21,41 @@ 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
 
 
 @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,82 +109,66 @@ 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
     training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig)
 
     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

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:

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
 

neural_lam/datastore/mdp.py

@@ -1,3 +1,5 @@
+"""Datastore implementation wrapping ``mllam-data-prep`` outputs."""
+
 # Standard library
 import copy
 import functools

neural_lam/datastore/npyfilesmeps/__init__.py

@@ -1,2 +1,4 @@
+"""MEPS-specific datastore exposing numpy-based datasets."""
+
 # Local
 from .store import NpyFilesDatastoreMEPS  # noqa