diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index fe2ff11a..b75a383f 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: ["-c", "pyproject.toml", "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/CHANGELOG.md b/CHANGELOG.md index 43ba2246..7e2768f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -74,6 +74,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Maintenance +- Establish 100% docstring coverage across `neural_lam/` via an `interrogate` pre-commit hook, add a Dimension Glossary to the README for canonical tensor-shape names, and rewrite public docstrings in NumPy style to serve as the entry point for the autoapi pipeline (#196 / #272). [\#252](https://github.com/mllam/neural-lam/pull/252) @Mohit-Lakra + - Register a `slow` pytest marker and apply it to `test_training` and `test_training_output_std` so contributors can skip long-running training tests during local iteration via `pytest -m "not slow"`. [\#651](https://github.com/mllam/neural-lam/pull/651) @sadamov - Add a short README pointer to [\#163](https://github.com/mllam/neural-lam/issues/163) for DGX Spark / PyTorch container compatibility notes, so users hitting `torch_scatter` errors know where to find the known-working / known-failing combos [\#266](https://github.com/mllam/neural-lam/pull/266) @Jayant-kernel diff --git a/README.md b/README.md index b852bedd..68f7faa8 100644 --- a/README.md +++ b/README.md @@ -578,6 +578,23 @@ In addition, hierarchical mesh graphs (`L > 1`) feature a few additional files w These files have the same list format as the ones above, but each list has length `L-1` (as these edges describe connections between levels). Entries 0 in these lists describe edges between the lowest levels 1 and 2. +## Dimension Glossary + +Canonical dimension names used in tensor shape annotations throughout the codebase: + +- `B` - batch size +- `pred_steps` - number of autoregressive prediction steps +- `num_grid_nodes` - number of nodes in the flattened spatial grid +- `num_mesh_nodes` - number of mesh nodes; indexed as `num_mesh_nodes[l]` for hierarchical level `l` +- `num_state_vars` - number of atmospheric state variables +- `num_forcing_vars` - number of forcing input variables +- `num_variables` - generic variable dimension used in metric functions +- `hidden_dim` - internal hidden representation size in GNN layers and MLPs +- `input_dim` - input feature dimensionality to a layer before transformation +- `num_edges` - number of edges in a graph (g2m, m2g, same-level, up, down) +- `num_send` - number of sender nodes in a message-passing step +- `num_rec` - number of receiver nodes in a message-passing step + # Development and Contributing Any push or Pull-Request to the main branch will trigger a selection of pre-commit hooks. These hooks will run a series of checks on the code, like formatting and linting. diff --git a/neural_lam/__init__.py b/neural_lam/__init__.py index 61a2aff2..b1d01939 100644 --- a/neural_lam/__init__.py +++ b/neural_lam/__init__.py @@ -1,3 +1,5 @@ +"""Neural-LAM: graph-based neural weather prediction models.""" + # Standard library import importlib.metadata diff --git a/neural_lam/config.py b/neural_lam/config.py index f4195ec3..1da43fff 100644 --- a/neural_lam/config.py +++ b/neural_lam/config.py @@ -1,3 +1,5 @@ +"""Configuration dataclasses and helpers for Neural-LAM experiments.""" + # Standard library import dataclasses from pathlib import Path @@ -31,13 +33,20 @@ class DatastoreSelection: """ kind: str + config_path: str def __post_init__(self): + """ + Validate that the selected datastore kind is implemented. + + Raises + ------ + ValueError + If the provided ``kind`` is not part of :data:`DATASTORES`. + """ if self.kind not in DATASTORES: raise ValueError(f"Datastore kind {self.kind} is not implemented") - config_path: str - @dataclasses.dataclass class ManualStateFeatureWeighting: @@ -89,10 +98,13 @@ class TrainingConfig: Attributes ---------- state_feature_weighting : Union[ManualStateFeatureWeighting, - UnformFeatureWeighting] + UniformFeatureWeighting] The method to use for weighting the state features in the loss - function. Defaults to uniform weighting (`UnformFeatureWeighting`, i.e. + function. Defaults to uniform weighting (`UniformFeatureWeighting`, i.e. all features are weighted equally). + output_clamping : OutputClamping + Per-feature lower / upper clamping bounds applied to the model output. + Defaults to an empty ``OutputClamping`` (no clamping). """ state_feature_weighting: Union[ @@ -107,15 +119,20 @@ 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. + Configuration for the Neural-LAM model and training pipeline. + + Loads and stores all settings needed to run Neural-LAM, including + datastore selection and training hyperparameters. Serialisation and + deserialisation from YAML/JSON is handled via ``dataclass_wizard``. Attributes ---------- datastore : DatastoreSelection - The configuration for the datastore to use. + Configuration specifying which datastore backend to use and its + associated settings. training : TrainingConfig - The configuration for training the model. + Configuration for training the model, including loss function and + feature-weighting strategy. Defaults to ``TrainingConfig()``. """ datastore: DatastoreSelection @@ -150,6 +167,8 @@ class _(dataclass_wizard.JSONWizard.Meta): class InvalidConfigError(Exception): + """Raised when the Neural-LAM configuration file is invalid or malformed.""" + pass diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index c0f47f75..c67e4397 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -1,3 +1,5 @@ +"""Graph construction utilities for Neural-LAM meshes and grids.""" + # Standard library import os from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser @@ -22,6 +24,21 @@ def plot_graph( graph: pyg.data.Data, title: Optional[str] = None ) -> tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]: + """ + Render a PyTorch Geometric graph using stored node coordinates. + + Parameters + ---------- + graph : torch_geometric.data.Data + Graph containing ``edge_index`` and ``pos`` attributes. + title : str or None, optional + Optional subplot title. + + Returns + ------- + tuple[matplotlib.figure.Figure, matplotlib.axes.Axes] + Figure and axis handles for further customization. + """ fig, axis = plt.subplots(figsize=(8, 8), dpi=200) # W,H edge_index = graph.edge_index pos = graph.pos @@ -73,6 +90,19 @@ def plot_graph( def sort_nodes_internally(nx_graph: networkx.Graph) -> networkx.DiGraph: + """ + Return a copy of ``nx_graph`` with deterministically ordered nodes. + + Parameters + ---------- + nx_graph : networkx.Graph + The input graph to sort nodes for. + + Returns + ------- + networkx.DiGraph + A directed graph with nodes sorted alphabetically by their labels. + """ # For some reason the networkx .nodes() return list can not be sorted, # but this is the ordering used by pyg when converting. # This function fixes this. @@ -83,6 +113,18 @@ def sort_nodes_internally(nx_graph: networkx.Graph) -> networkx.DiGraph: def save_edges(graph: pyg.data.Data, name: str, base_path: str) -> None: + """ + Persist edge indices/features for a PyG graph under ``base_path``. + + Parameters + ---------- + graph : torch_geometric.data.Data + The graph containing edge data. + name : str + The name prefix for the saved files. + base_path : str + The directory path where files should be saved. + """ torch.save( graph.edge_index, os.path.join(base_path, f"{name}_edge_index.pt") ) @@ -95,6 +137,18 @@ def save_edges(graph: pyg.data.Data, name: str, base_path: str) -> None: def save_edges_list( graphs: list[pyg.data.Data], name: str, base_path: str ) -> None: + """ + Persist edge indices/features for a list of graphs. + + Parameters + ---------- + graphs : list of torch_geometric.data.Data + The list of graphs containing edge data. + name : str + The name prefix for the saved files. + base_path : str + The directory path where files should be saved. + """ torch.save( [graph.edge_index for graph in graphs], os.path.join(base_path, f"{name}_edge_index.pt"), @@ -111,12 +165,44 @@ def save_edges_list( def from_networkx_with_start_index( nx_graph: networkx.Graph, start_index: int ) -> pyg.data.Data: + """ + Convert a NetworkX graph to PyG and offset node indices. + + Parameters + ---------- + nx_graph : networkx.Graph + The NetworkX graph to convert. + start_index : int + The value to add to each node index. + + Returns + ------- + pyg.data.Data + The converted PyG graph. + """ pyg_graph = from_networkx(nx_graph) pyg_graph.edge_index += start_index return pyg_graph def mk_2d_graph(xy: np.ndarray, nx: int, ny: int) -> networkx.DiGraph: + """ + Create a diagonal 2-D grid graph over the ``xy`` positions. + + Parameters + ---------- + xy : np.ndarray + The grid coordinates. + nx : int + Number of nodes in the x-dimension. + ny : int + Number of nodes in the y-dimension. + + Returns + ------- + networkx.DiGraph + The constructed directed 2-D grid graph. + """ xm, xM = np.amin(xy[:, :, 0][:, 0]), np.amax(xy[:, :, 0][:, 0]) ym, yM = np.amin(xy[:, :, 1][0, :]), np.amax(xy[:, :, 1][0, :]) @@ -156,6 +242,21 @@ def mk_2d_graph(xy: np.ndarray, nx: int, ny: int) -> networkx.DiGraph: def prepend_node_index(graph: networkx.Graph, new_index: int) -> networkx.Graph: + """ + Relabel each node by prepending ``new_index`` to its tuple identifier. + + Parameters + ---------- + graph : networkx.Graph + The graph to relabel. + new_index : int + The value to prepend to each node identifier. + + Returns + ------- + networkx.Graph + The relabeled graph. + """ # Relabel node indices in graph, insert (graph_level, i, j) ijk = [tuple((new_index,) + x) for x in graph.nodes] to_mapping = dict(zip(graph.nodes, ijk)) @@ -550,6 +651,22 @@ def create_graph_from_datastore( hierarchical: bool = False, create_plot: bool = False, ): + """ + Generate graph components for ``datastore`` and persist them on disk. + + Parameters + ---------- + datastore : BaseRegularGridDatastore + Datastore providing ``get_xy`` for state nodes. + output_root_path : str + Directory where the resulting ``*.pt`` graph files are stored. + n_max_levels : int or None, optional + Optional limit of hierarchical mesh levels to build. + hierarchical : bool, optional + If ``True``, create multi-level hierarchical graphs. Default ``False``. + create_plot : bool, optional + If ``True``, display matplotlib previews of the generated graphs. + """ if isinstance(datastore, BaseRegularGridDatastore): xy = datastore.get_xy(category="state", stacked=False) else: @@ -567,6 +684,15 @@ def create_graph_from_datastore( def cli(input_args: Optional[list[str]] = None) -> None: + """ + Parse CLI arguments and call :func:`create_graph_from_datastore`. + + Parameters + ---------- + input_args : list[str] or None, optional + Argument list forwarded to :class:`argparse.ArgumentParser`. When + ``None``, ``sys.argv`` is used. + """ parser = ArgumentParser( description="Graph generation for neural-lam", formatter_class=ArgumentDefaultsHelpFormatter, diff --git a/neural_lam/custom_loggers.py b/neural_lam/custom_loggers.py index 8dc7c73c..5b6000e6 100644 --- a/neural_lam/custom_loggers.py +++ b/neural_lam/custom_loggers.py @@ -1,3 +1,5 @@ +"""Custom logging utilities (e.g., MLFlow wrappers) used in Neural-LAM.""" + # Standard library import os from typing import Optional @@ -24,10 +26,26 @@ def __init__( run_name: str, save_dir: str, ) -> None: - """Initialize the logger and ensure ``save_dir`` exists on disk. + """ + Initialize the logger, ensure ``save_dir`` exists, and start the + MLflow run. + + Parameters + ---------- + experiment_name : str + Target MLflow experiment. + tracking_uri : str + MLflow tracking server URI. + run_name : str + Human-readable run name stored as ``mlflow.runName``. + save_dir : str + Directory where ``log_image`` writes temporary figure files. + Created eagerly with ``exist_ok=True``. - ``save_dir`` is created eagerly (with ``exist_ok=True``) so that - subsequent ``log_image`` calls can write temporary files there. + Notes + ----- + Starts the MLflow run with ``log_system_metrics=True`` and also + records ``run_id`` as an MLflow param. """ super().__init__( experiment_name=experiment_name, tracking_uri=tracking_uri @@ -75,6 +93,11 @@ def log_image( step : int or None, optional Step to associate with the log entry. ``None`` logs without a step suffix. + + Raises + ------ + SystemExit + If AWS credentials for the MLflow artifact store are missing. """ # Third-party from botocore.exceptions import NoCredentialsError diff --git a/neural_lam/datastore/__init__.py b/neural_lam/datastore/__init__.py index dead7713..e2117217 100644 --- a/neural_lam/datastore/__init__.py +++ b/neural_lam/datastore/__init__.py @@ -1,3 +1,5 @@ +"""Datastore backends for loading and serving weather model data.""" + # Local from .base import BaseDatastore # noqa from .mdp import MDPDatastore # noqa @@ -15,6 +17,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 a3870b0b..8376d38d 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 @@ -6,7 +8,6 @@ from datetime import timedelta from functools import cached_property from pathlib import Path -from typing import List, Optional, Union # Third-party import cartopy.crs as ccrs @@ -28,25 +29,28 @@ class BaseDatastore(abc.ABC): `weather_dataset.WeatherDataset` class (which inherits from `torch.utils.data.Dataset` and uses the datastore to access the data). - # Forecast vs analysis data - If the datastore is used to represent forecast rather than analysis data, - then the `is_forecast` attribute should be set to True, and returned data - from `get_dataarray` is assumed to have `analysis_time` and `forecast_time` - dimensions (rather than just `time`). + Forecast vs analysis data + ------------------------- + If the datastore is used to represent forecast rather than analysis + data, then the ``is_forecast`` attribute should be set to True, and + returned data from ``get_dataarray`` is assumed to have `analysis_time` + and `forecast_time` dimensions (rather than just `time`). - # Ensemble vs deterministic data + Ensemble vs deterministic data + ------------------------------ If the datastore is used to present an ensemble of state realisations, for example for forecast ensembles, then the `is_ensemble` attribute should be - set to `True` and returned state data from `get_dataarray` is expected to + set to `True` and returned state data from ``get_dataarray`` is expected to have an `ensemble_member` dimension. If each ensemble member has its own forcing values, then `has_ensemble_forcing` should be set to `True`, and - returned forcing data from `get_dataarray` is expected to have an + returned forcing data from ``get_dataarray`` is expected to have an `ensemble_member` dimension; otherwise forcing data is expected not to have one. - # Grid index + Grid index + ---------- All methods that return data specific to a grid point (like - `get_dataarray`) should have a single dimension named `grid_index` that + ``get_dataarray``) should have a single dimension named `grid_index` that represents the spatial grid index of the data. The actual x, y coordinates of the grid points should be stored in the `x` and `y` coordinates of the dataarray or dataset with the `grid_index` dimension as the coordinate for @@ -70,7 +74,6 @@ def root_path(self) -> Path: The root path to the datastore. """ - pass @property @abc.abstractmethod @@ -84,7 +87,6 @@ def config(self) -> collections.abc.Mapping: returned. """ - pass @property @abc.abstractmethod @@ -97,10 +99,9 @@ def step_length(self) -> timedelta: The step length of the dataset. """ - pass @abc.abstractmethod - def get_vars_units(self, category: str) -> List[str]: + def get_vars_units(self, category: str) -> list[str]: """Get the units of the variables in the given category. Parameters @@ -114,10 +115,9 @@ def get_vars_units(self, category: str) -> List[str]: The units of the variables. """ - pass @abc.abstractmethod - def get_vars_names(self, category: str) -> List[str]: + def get_vars_names(self, category: str) -> list[str]: """Get the names of the variables in the given category. Parameters @@ -131,10 +131,9 @@ def get_vars_names(self, category: str) -> List[str]: The names of the variables. """ - pass @abc.abstractmethod - def get_vars_long_names(self, category: str) -> List[str]: + def get_vars_long_names(self, category: str) -> list[str]: """Get the long names of the variables in the given category. Parameters @@ -148,7 +147,6 @@ def get_vars_long_names(self, category: str) -> List[str]: The long names of the variables. """ - pass @abc.abstractmethod def get_num_data_vars(self, category: str) -> int: @@ -165,7 +163,6 @@ def get_num_data_vars(self, category: str) -> int: The number of data variables. """ - pass @abc.abstractmethod def get_standardization_dataarray(self, category: str) -> xr.Dataset: @@ -194,7 +191,6 @@ def get_standardization_dataarray(self, category: str) -> xr.Dataset: differences for state variables). """ - pass def _standardize_datarray( self, da: xr.DataArray, category: str @@ -228,9 +224,9 @@ def _standardize_datarray( def get_dataarray( self, category: str, - split: Optional[str], + split: str | None, standardize: bool = False, - ) -> Union[xr.DataArray, None]: + ) -> xr.DataArray | None: """ Return the processed data (as a single `xr.DataArray`) for the given category of data and test/train/val-split that covers all the data (in @@ -272,7 +268,6 @@ def get_dataarray( The xarray DataArray object with processed dataset. """ - pass @cached_property @abc.abstractmethod @@ -289,7 +284,6 @@ def boundary_mask(self) -> xr.DataArray: `('grid_index',)`. """ - pass @abc.abstractmethod def get_xy(self, category: str, stacked: bool) -> np.ndarray: @@ -324,10 +318,9 @@ def coords_projection(self) -> ccrs.Projection: The projection object. """ - pass @functools.lru_cache - def get_xy_extent(self, category: str) -> List[float]: + def get_xy_extent(self, category: str) -> list[float]: """ Return the extent of the x, y coordinates for a given category of data. The extent should be returned as a list of 4 floats with `[xmin, xmax, @@ -381,11 +374,10 @@ def num_grid_points(self) -> int: The number of grid points in the dataset. """ - pass @cached_property @abc.abstractmethod - def state_feature_weights_values(self) -> List[float]: + def state_feature_weights_values(self) -> list[float]: """ Return the weights for each state feature as a list of floats. @@ -399,12 +391,11 @@ def state_feature_weights_values(self) -> List[float]: list of float The weight for each state feature. """ - pass @functools.lru_cache def expected_dim_order( self, - category: Optional[str] = None, + category: str | None = None, ) -> tuple[str, ...]: """ Return the expected dimension order for the dataarray or dataset @@ -443,14 +434,20 @@ def expected_dim_order( # static data does not vary in time if self.is_forecast: dim_order.extend( - ["analysis_time", "elapsed_forecast_duration"] + [ + "analysis_time", + "elapsed_forecast_duration", + ] ) elif not self.is_forecast: dim_order.append("time") - if category == "state" and self.is_ensemble: - dim_order.append("ensemble_member") - elif category == "forcing" and self.has_ensemble_forcing: + if ( + category == "state" + and self.is_ensemble + or category == "forcing" + and self.has_ensemble_forcing + ): dim_order.append("ensemble_member") dim_order.append("grid_index") @@ -513,7 +510,6 @@ def grid_shape_state(self) -> CartesianGridShape: `y` attributes. """ - pass @abc.abstractmethod def get_xy(self, category: str, stacked: bool) -> np.ndarray: @@ -535,11 +531,10 @@ def get_xy(self, category: str, stacked: bool) -> np.ndarray: n_grid_points=N_x*N_y. - `stacked==False`: shape `(N_x, N_y, 2)` """ - pass def unstack_grid_coords( - self, da_or_ds: Union[xr.DataArray, xr.Dataset] - ) -> Union[xr.DataArray, xr.Dataset]: + self, da_or_ds: xr.DataArray | xr.Dataset + ) -> xr.DataArray | xr.Dataset: """ Unstack the spatial grid coordinates from `grid_index` into separate `x` and `y` dimensions to create a 2D grid (if the spatial coordinates have @@ -593,8 +588,8 @@ def unstack_grid_coords( return da_or_ds_unstacked def stack_grid_coords( - self, da_or_ds: Union[xr.DataArray, xr.Dataset] - ) -> Union[xr.DataArray, xr.Dataset]: + self, da_or_ds: xr.DataArray | xr.Dataset + ) -> xr.DataArray | xr.Dataset: """ Stack the spatial grid coordinates (x and y) into a single `grid_index` dimension. Only performs stacking if the data is currently unstacked diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py index 88f45e37..7cad45d7 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 @@ -58,6 +60,12 @@ def __init__( Whether to reuse an existing dataset zarr file if it exists and its creation date is newer than the configuration file. + Raises + ------ + ValueError + If the dataset does not contain all of the required + train/val/test splits. + """ self._config_path = Path(config_path) self._root_path = self._config_path.parent @@ -417,6 +425,12 @@ def coords_projection(self) -> ccrs.Projection: ccrs.Projection The projection of the coordinates. + Raises + ------ + ValueError + If the `projection` entry is missing from the `extra` section of + the config, or if its `class_name` or `kwargs` keys are missing. + """ if "projection" not in self._config.extra: raise ValueError( diff --git a/neural_lam/datastore/npyfilesmeps/__init__.py b/neural_lam/datastore/npyfilesmeps/__init__.py index 397a5075..27cc356d 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 531db87a..9e694f14 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,22 +57,54 @@ 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_rank(): + """ + Return the rank inferred from SLURM or default to 0. + + Returns + ------- + int + The current process rank. + """ return int(os.environ.get("SLURM_PROCID", 0)) def get_world_size(): + """ + Return the world size inferred from SLURM or default to 1. + + Returns + ------- + int + The total number of processes in the distributed group. + """ return int(os.environ.get("SLURM_NTASKS", 1)) def setup(rank, world_size): # pylint: disable=redefined-outer-name - """Initialize the distributed group.""" + """ + Initialize the distributed group. + + Parameters + ---------- + rank : int + The rank of the current process. + world_size : int + The total number of processes. + + Raises + ------ + RuntimeError + If ``SLURM_JOB_NODELIST`` is set but no hostnames can be retrieved. + """ if "SLURM_JOB_NODELIST" in os.environ: nodelist = os.environ["SLURM_JOB_NODELIST"] hostnames = subprocess.check_output( @@ -95,12 +142,44 @@ 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] + Shape ``(B, num_state_vars)``. Batch-wise means. Sequence of one or + more tensors each of shape ``(B_i, num_state_vars)``; concatenated / + stacked along the batch dim before reduction. Typical callers pass a + single already-gathered tensor. The list can have length 0 to skip + saving. + squares : Sequence[torch.Tensor] + Shape ``(B, num_state_vars)``. Batch-wise second moments. Sequence of + one or more tensors each of shape ``(B_i, num_state_vars)``; + concatenated / stacked along the batch dim before reduction. Typical + callers pass a single already-gathered tensor. The list can have + length 0 to skip saving. + flux_means : Sequence[torch.Tensor] + Shape ``(B,)``. Flux means. Sequence of one or more tensors each of + shape ``(B_i,)``; concatenated / stacked along the batch dim before + reduction. Typical callers pass a single already-gathered tensor. + The list can have length 0 to skip saving. + flux_squares : Sequence[torch.Tensor] + Shape ``(B,)``. Flux second moments. Sequence of one or more tensors + each of shape ``(B_i,)``; concatenated / stacked along the batch dim + before reduction. Typical callers pass a single already-gathered + tensor. The list can have length 0 to skip saving. + 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,) + ) # (B, d_features,) squares = ( torch.stack(squares) if len(squares) > 1 else squares[0] - ) # (N_batch, d_features,) + ) # (B, d_features,) mean = torch.mean(means, dim=0) # (d_features,) second_moment = torch.mean(squares, dim=0) # (d_features,) std = torch.sqrt(second_moment - mean**2) # (d_features,) @@ -119,10 +198,10 @@ def save_stats( return flux_means = ( torch.stack(flux_means) if len(flux_means) > 1 else flux_means[0] - ) # (N_batch,) + ) # (B,) flux_squares = ( torch.stack(flux_squares) if len(flux_squares) > 1 else flux_squares[0] - ) # (N_batch,) + ) # (B,) flux_mean = torch.mean(flux_means) # (,) flux_second_moment = torch.mean(flux_squares) # (,) flux_std = torch.sqrt(flux_second_moment - flux_mean**2) # (,) @@ -137,20 +216,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() @@ -210,15 +289,15 @@ def main( target_batch.to(device), forcing_batch.to(device), ) - # (N_batch, N_t, N_grid, d_features) + # (B, N_t, num_grid_nodes, d_features) batch = torch.cat((init_batch, target_batch), dim=1) # Flux at 1st windowed position is index 0 in forcing flux_batch = forcing_batch[:, :, :, 0] - # (N_batch, d_features,) + # (B, d_features,) means.append(torch.mean(batch, dim=(1, 2)).cpu()) squares.append( torch.mean(batch**2, dim=(1, 2)).cpu() - ) # (N_batch, d_features,) + ) # (B, d_features,) flux_means.append(torch.mean(flux_batch).cpu()) # (,) flux_squares.append(torch.mean(flux_batch**2).cpu()) # (,) @@ -263,10 +342,10 @@ def main( ) ] else: - means = [torch.cat(means, dim=0)] # (N_batch, d_features,) - squares = [torch.cat(squares, dim=0)] # (N_batch, d_features,) - flux_means = [torch.tensor(flux_means)] # (N_batch,) - flux_squares = [torch.tensor(flux_squares)] # (N_batch,) + means = [torch.cat(means, dim=0)] # (B, d_features,) + squares = [torch.cat(squares, dim=0)] # (B, d_features,) + flux_means = [torch.tensor(flux_means)] # (B,) + flux_squares = [torch.tensor(flux_squares)] # (B,) if rank == 0: save_stats( @@ -338,7 +417,7 @@ def main( ) init_batch = (init_batch - state_mean) / state_std target_batch = (target_batch - state_mean) / state_std - # (N_batch, N_t', N_grid, d_features) + # (B, N_t', num_grid_nodes, num_state_vars) batch = torch.cat((init_batch, target_batch), dim=1) # Note: batch contains only 1h-steps stepped_batch = torch.cat( @@ -348,14 +427,14 @@ def main( ], dim=0, ) - # (N_batch', N_t, N_grid, d_features), - # N_batch' = step_length*N_batch + # (B', N_t, num_grid_nodes, d_features), + # B' = step_length*B batch_diffs = stepped_batch[:, 1:] - stepped_batch[:, :-1] - # (N_batch', N_t-1, N_grid, d_features) + # (B', N_t-1, num_grid_nodes, d_features) diff_means.append(torch.mean(batch_diffs, dim=(1, 2)).cpu()) - # (N_batch', d_features,) + # (B', d_features,) diff_squares.append(torch.mean(batch_diffs**2, dim=(1, 2)).cpu()) - # (N_batch', d_features,) + # (B', d_features,) if distributed and world_size > 1: dist.barrier() @@ -379,8 +458,8 @@ def main( diff_means = [diff_means_gathered[:n_original_windows]] diff_squares = [diff_squares_gathered[:n_original_windows]] - diff_means = [torch.cat(diff_means, dim=0)] # (N_batch', d_features,) - diff_squares = [torch.cat(diff_squares, dim=0)] # (N_batch', d_features,) + diff_means = [torch.cat(diff_means, dim=0)] # (B', d_features,) + diff_squares = [torch.cat(diff_squares, dim=0)] # (B', d_features,) if rank == 0: save_stats(static_dir_path, diff_means, diff_squares, [], [], "diff") @@ -390,6 +469,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 1d36b6fe..b423b4a1 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 @@ -36,6 +38,11 @@ class Dataset: var_units: A list of units for each variable. var_longnames: A list of long, descriptive names for each variable. num_forcing_features: The number of forcing features in the dataset. + num_timesteps: The number of timesteps per analysis sample. + step_length: The time delta between consecutive timesteps. + num_ensemble_members: The number of ensemble members in the dataset. + remove_state_features_with_index: Indices of state features to drop + when loading the dataset. """ diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py index f65fd8ae..e14bbc74 100644 --- a/neural_lam/datastore/npyfilesmeps/store.py +++ b/neural_lam/datastore/npyfilesmeps/store.py @@ -35,6 +35,23 @@ def _load_np(fp, add_feature_dim, feature_dim_mask=None): + """ + Load an ``.npy`` file and optionally expand/mask the feature axis. + + Parameters + ---------- + fp : str or Path + The file path to load. + add_feature_dim : bool + Whether to add a new feature dimension at the end. + feature_dim_mask : list of int or None, optional + Mask to apply to the feature dimension. + + Returns + ------- + np.ndarray + The loaded and optionally processed array. + """ arr = np.load(fp) if add_feature_dim: arr = arr[..., np.newaxis] @@ -44,25 +61,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}' + `nwp_{analysis_time:%Y%m%d%H}_mbr{member_id:03d}.npy`. The MEPS dataset is organised into three splits: train, val, and test. Each split has a set of files which are: - - `{STATE_FILENAME_FORMAT}`: + - `nwp_{analysis_time:%Y%m%d%H}_mbr{member_id:03d}.npy`: 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}`: + - `nwp_toa_downwelling_shortwave_flux_{analysis_time:%Y%m%d%H}.npy`: The top-of-atmosphere downwelling shortwave flux at `time`. The dimensions of the array are `[forecast_timestep, y, x]`. - - `{OPEN_WATER_FILENAME_FORMAT}`: + - `wtr_{analysis_time:%Y%m%d%H}.npy`: The open water fraction at `time`. The dimensions of the array are `[y, x]`. @@ -127,19 +144,20 @@ class NpyFilesDatastoreMEPS(BaseRegularGridDatastore): N_t = 65//subsample_step (= 21 for 3h steps) dim_y = 268 dim_x = 238 - N_grid = 268x238 = 63784 - d_features = 17 (d_features' = 18) - d_forcing = 5 + num_grid_nodes = 268x238 = 63784 + num_state_vars = 17 (num_state_vars' = 18) + num_forcing_vars = 5 For the MEPS reduced dataset: N_t' = 65 N_t = 65//subsample_step (= 21 for 3h steps) dim_y = 134 dim_x = 119 - N_grid = 134x119 = 15946 - d_features = 8 - d_forcing = 1 + num_grid_nodes = 134x119 = 15946 + num_state_vars = 8 + num_forcing_vars = 1 """ + SHORT_NAME = "npyfilesmeps" is_forecast = True @@ -551,6 +569,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 +605,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 +623,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 +643,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 +651,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: @@ -635,7 +670,7 @@ def get_xy(self, category: str, stacked: bool) -> np.ndarray: The x, y coordinates of the dataset (with x first then y second), returned differently based on the value of `stacked`: - `stacked==True`: shape `(n_grid_points, 2)` where - n_grid_points=N_x*N_y. + n_grid_points=N_x*N_y. - `stacked==False`: shape `(N_x, N_y, 2)` """ @@ -732,6 +767,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 19a81e30..13f32e57 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/gnn_layers.py b/neural_lam/gnn_layers.py index 9da9fd71..7a92ea76 100644 --- a/neural_lam/gnn_layers.py +++ b/neural_lam/gnn_layers.py @@ -1,3 +1,5 @@ +"""Interaction Network and PropagationNet GNN layers used by Neural-LAM.""" + # Standard library from typing import Optional, Type, Union @@ -36,26 +38,30 @@ def __init__( Parameters ---------- edge_index : torch.Tensor - Shape ``(2, M)``. Edges in PyG format; both sender and receiver - node indices start at 0. Dims: ``M`` is the number of edges. + Edge connectivity tensor in PyG format. + Shape ``(2, num_edges)``. input_dim : int - Dimensionality of input representations for both nodes and - edges. + Dimensionality of both node and edge input representations. update_edges : bool, optional - If True, compute and return updated edge representations. + If ``True``, compute and return updated edge representations in + addition to node representations. Default is ``True``. hidden_layers : int, optional - Number of hidden layers in each MLP. - hidden_dim : int, optional - Dimensionality of hidden layers. Defaults to ``input_dim``. - edge_chunk_sizes : list of int, optional - Chunk sizes to split edge representations into, each fed - through a separate MLP. ``None`` means a single shared MLP. - aggr_chunk_sizes : list of int, optional - Chunk sizes to split aggregated node representations into, - each fed through a separate MLP. ``None`` means a single - shared MLP. - aggr : str, optional - Message aggregation method (``'sum'`` or ``'mean'``). + Number of hidden layers in each MLP. Default is ``1``. + hidden_dim : int or None, optional + Width of hidden layers. If ``None``, defaults to ``input_dim``. + edge_chunk_sizes : list[int] or None, optional + Chunk sizes for splitting edge representations across separate + MLPs. ``None`` uses a single shared MLP. + aggr_chunk_sizes : list[int] or None, optional + Chunk sizes for splitting aggregated node representations across + separate MLPs. ``None`` uses a single shared MLP. + aggr : {"sum", "mean"}, optional + Message aggregation method. Default is ``"sum"``. + + Raises + ------ + ValueError + If ``aggr`` is not one of ``"sum"`` or ``"mean"``. """ if aggr not in ("sum", "mean"): raise ValueError(f"Unknown aggregation method: {aggr}") @@ -115,23 +121,23 @@ def forward( Parameters ---------- send_rep : torch.Tensor - Shape ``(B, N_send, d_h)``. Sender node representations. - Dims: ``B`` is batch size, ``N_send`` is the number of - sender nodes, and ``d_h`` is the hidden dimension. + Sender node representations. + Shape ``(num_send, input_dim)``. rec_rep : torch.Tensor - Shape ``(B, N_rec, d_h)``. Receiver node representations. - Dims: ``N_rec`` is the number of receiver nodes. + Receiver node representations. + Shape ``(num_rec, input_dim)``. edge_rep : torch.Tensor - Shape ``(B, M, d_h)``. Edge representations. Dims: ``M`` - is the number of edges. + Edge representations. + Shape ``(num_edges, input_dim)``. Returns ------- rec_rep : torch.Tensor - Shape ``(B, N_rec, d_h)``. Updated receiver node - representations. + Updated receiver node representations. + Shape ``(num_rec, input_dim)``. edge_rep : torch.Tensor - Shape ``(B, M, d_h)``. Updated edge representations. + Updated edge representations. + Shape ``(num_edges, input_dim)``. Only returned when ``update_edges=True``. """ # Always concatenate to [rec_nodes, send_nodes] for propagation, @@ -163,9 +169,7 @@ def node_residual_target( def message( self, x_j: torch.Tensor, x_i: torch.Tensor, edge_attr: torch.Tensor ) -> torch.Tensor: - """ - Compute messages from node j to node i. - """ + """Compute messages from node ``j`` to ``i``.""" return self.edge_mlp(torch.cat((edge_attr, x_j, x_i), dim=-1)) # pylint: disable-next=signature-differs @@ -178,8 +182,9 @@ def aggregate( ) -> tuple[torch.Tensor, torch.Tensor]: """ Overridden aggregation function to: - * return both aggregated and original messages, - * only aggregate to number of receiver nodes. + * return both aggregated and per-edge messages, + * only aggregate to the number of receiver nodes (``self.num_rec``) + rather than to ``dim_size``. """ aggr = super().aggregate(inputs, index, ptr, self.num_rec) return aggr, inputs @@ -205,6 +210,13 @@ def __init__( aggr_chunk_sizes: Optional[list[int]] = None, aggr: str = "sum", ) -> None: + """Initialise the :class:`PropagationNet` layer. + + Parameters share the meaning of :class:`InteractionNet.__init__`; see + that class for the full description. The propagation variant overrides + ``aggr`` defaults internally to favour stability of the propagation + residual. + """ # Use mean aggregation in propagation version to avoid instability super().__init__( edge_index, @@ -268,6 +280,21 @@ class SplitMLPs(nn.Module): """ def __init__(self, mlps: list[nn.Module], chunk_sizes: list[int]) -> None: + """ + Create a module that dispatches chunks of the input to separate MLPs. + + Parameters + ---------- + mlps : list of nn.Module + Sequence of MLPs to apply to each chunk. + chunk_sizes : list of int + Sizes used when splitting the input along ``dim=-2``. + + Raises + ------ + AssertionError + If the number of ``mlps`` and ``chunk_sizes`` differ. + """ super().__init__() assert len(mlps) == len( chunk_sizes diff --git a/neural_lam/loss_weighting.py b/neural_lam/loss_weighting.py index 90538bb8..8a37d07f 100644 --- a/neural_lam/loss_weighting.py +++ b/neural_lam/loss_weighting.py @@ -1,3 +1,5 @@ +"""Utility functions for configuring state-feature loss weighting.""" + # Local from .config import ( ManualStateFeatureWeighting, @@ -25,6 +27,12 @@ def get_manual_state_feature_weights( ------- list[float] List of floats containing the state feature weights. + + Raises + ------ + ValueError + If the set of feature names in ``weighting_config.weights`` does + not match the state feature names in the datastore. """ state_feature_names = datastore.get_vars_names(category="state") feature_weight_names = weighting_config.weights.keys() @@ -90,6 +98,12 @@ def get_state_feature_weighting( ------- list[float] List of floats containing the state feature weights. + + Raises + ------ + NotImplementedError + If ``config.training.state_feature_weighting`` is not a + recognised weighting configuration type. """ weighting_config = config.training.state_feature_weighting diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 7c7f29b6..1eb1d526 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -1,3 +1,5 @@ +"""Evaluation metrics shared across training and validation routines.""" + # Standard library from collections.abc import Callable from typing import Optional @@ -19,6 +21,12 @@ def get_metric(metric_name: str) -> Callable[..., torch.Tensor]: ------- callable Function implementing the requested metric. + + Raises + ------ + AssertionError + If ``metric_name`` (case-insensitive) is not a key in + ``DEFINED_METRICS``. """ metric_name_lower = metric_name.lower() assert ( @@ -39,40 +47,40 @@ def mask_and_reduce_metric( Parameters ---------- metric_entry_vals : torch.Tensor - Shape ``(..., N, d_state)``. Per-entry metric values. ``(...)`` + Shape ``(..., N, num_variables)``. Per-entry metric values. ``(...)`` denotes any number of broadcastable batch dimensions, ``N`` is - the number of grid nodes, and ``d_state`` is the number of state - variables. + the number of grid nodes, and ``num_variables`` is the number of + variables in the gridded representation (e.g. state features). mask : torch.Tensor or None Shape ``(N,)``. Boolean mask selecting which grid nodes to include. ``None`` means all nodes are used. average_grid : bool If True, average over the grid dimension ``N``. sum_vars : bool - If True, sum over the variable dimension ``d_state``. + If True, sum over the variable dimension ``num_variables``. Returns ------- torch.Tensor Reduced metric values. Shape is one of ``(...,)``, - ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)`` + ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ # Only keep grid nodes in mask if mask is not None: metric_entry_vals = metric_entry_vals[ ..., mask, : - ] # (..., N', d_state) + ] # (..., num_selected_nodes, num_variables) # Optionally reduce last two dimensions if average_grid: # Reduce grid first metric_entry_vals = torch.mean( metric_entry_vals, dim=-2 - ) # (..., d_state) + ) # (..., num_variables) if sum_vars: # Reduce vars second metric_entry_vals = torch.sum( metric_entry_vals, dim=-1 - ) # (..., N) or (...,) + ) # (..., num_grid_nodes) or (...,) return metric_entry_vals @@ -91,14 +99,14 @@ def wmse( Parameters ---------- pred : torch.Tensor - Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes + Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes any number of broadcastable batch dimensions, ``N`` is the number - of grid nodes, and ``d_state`` is the number of state variables. + of grid nodes, and ``num_variables`` is the number of state variables. target : torch.Tensor - Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as + Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor - Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted + Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation used as per-entry weight. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all @@ -112,13 +120,15 @@ def wmse( ------- torch.Tensor Reduced metric values. Shape is one of ``(...,)``, - ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)`` + ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ entry_mse = torch.nn.functional.mse_loss( pred, target, reduction="none" - ) # (..., N, d_state) - entry_mse_weighted = entry_mse / (pred_std**2) # (..., N, d_state) + ) # (..., num_grid_nodes, num_variables) + entry_mse_weighted = entry_mse / ( + pred_std**2 + ) # (..., num_grid_nodes, num_variables) return mask_and_reduce_metric( entry_mse_weighted, @@ -142,14 +152,14 @@ def mse( Parameters ---------- pred : torch.Tensor - Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes + Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes any number of broadcastable batch dimensions, ``N`` is the number - of grid nodes, and ``d_state`` is the number of state variables. + of grid nodes, and ``num_variables`` is the number of state variables. target : torch.Tensor - Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as + Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor - Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted + Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation (unused; ``pred_std`` is replaced by ones internally). mask : torch.Tensor or None, optional @@ -164,7 +174,7 @@ def mse( ------- torch.Tensor Reduced metric values. Shape is one of ``(...,)``, - ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)`` + ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ # Replace pred_std with constant ones @@ -187,14 +197,14 @@ def wmae( Parameters ---------- pred : torch.Tensor - Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes + Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes any number of broadcastable batch dimensions, ``N`` is the number - of grid nodes, and ``d_state`` is the number of state variables. + of grid nodes, and ``num_variables`` is the number of state variables. target : torch.Tensor - Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as + Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor - Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted + Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation used as per-entry weight. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all @@ -208,13 +218,15 @@ def wmae( ------- torch.Tensor Reduced metric values. Shape is one of ``(...,)``, - ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)`` + ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ entry_mae = torch.nn.functional.l1_loss( pred, target, reduction="none" - ) # (..., N, d_state) - entry_mae_weighted = entry_mae / pred_std # (..., N, d_state) + ) # (..., num_grid_nodes, num_variables) + entry_mae_weighted = ( + entry_mae / pred_std + ) # (..., num_grid_nodes, num_variables) return mask_and_reduce_metric( entry_mae_weighted, @@ -238,14 +250,14 @@ def mae( Parameters ---------- pred : torch.Tensor - Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes + Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes any number of broadcastable batch dimensions, ``N`` is the number - of grid nodes, and ``d_state`` is the number of state variables. + of grid nodes, and ``num_variables`` is the number of state variables. target : torch.Tensor - Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as + Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor - Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted + Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation (unused; ``pred_std`` is replaced by ones internally). mask : torch.Tensor or None, optional @@ -260,7 +272,7 @@ def mae( ------- torch.Tensor Reduced metric values. Shape is one of ``(...,)``, - ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)`` + ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ # Replace pred_std with constant ones @@ -283,14 +295,14 @@ def nll( Parameters ---------- pred : torch.Tensor - Shape ``(..., N, d_state)``. Predicted mean. ``(...)`` denotes + Shape ``(..., N, num_variables)``. Predicted mean. ``(...)`` denotes any number of broadcastable batch dimensions, ``N`` is the number - of grid nodes, and ``d_state`` is the number of state variables. + of grid nodes, and ``num_variables`` is the number of state variables. target : torch.Tensor - Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as + Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor - Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted + Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation of the Gaussian. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all @@ -304,12 +316,14 @@ def nll( ------- torch.Tensor Reduced metric values. Shape is one of ``(...,)``, - ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)`` + ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ - # Broadcast pred_std if shaped (d_state,), done internally in Normal class - dist = torch.distributions.Normal(pred, pred_std) # (..., N, d_state) - entry_nll = -dist.log_prob(target) # (..., N, d_state) + # Broadcast pred_std if shaped (num_variables,) via distribution internals + dist = torch.distributions.Normal( + pred, pred_std + ) # (..., num_grid_nodes, num_variables) + entry_nll = -dist.log_prob(target) # (..., num_grid_nodes, num_variables) return mask_and_reduce_metric( entry_nll, mask=mask, average_grid=average_grid, sum_vars=sum_vars @@ -331,14 +345,14 @@ def crps_gauss( Parameters ---------- pred : torch.Tensor - Shape ``(..., N, d_state)``. Predicted mean. ``(...)`` denotes + Shape ``(..., N, num_variables)``. Predicted mean. ``(...)`` denotes any number of broadcastable batch dimensions, ``N`` is the number - of grid nodes, and ``d_state`` is the number of state variables. + of grid nodes, and ``num_variables`` is the number of state variables. target : torch.Tensor - Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as + Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. pred_std : torch.Tensor - Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted + Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation of the Gaussian. mask : torch.Tensor or None, optional Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all @@ -352,19 +366,21 @@ def crps_gauss( ------- torch.Tensor Reduced metric values. Shape is one of ``(...,)``, - ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)`` + ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. """ std_normal = torch.distributions.Normal( torch.zeros((), device=pred.device), torch.ones((), device=pred.device) ) - target_standard = (target - pred) / pred_std # (..., N, d_state) + target_standard = ( + target - pred + ) / pred_std # (..., num_grid_nodes, num_variables) entry_crps = -pred_std * ( torch.pi ** (-0.5) - 2 * torch.exp(std_normal.log_prob(target_standard)) - target_standard * (2 * std_normal.cdf(target_standard) - 1) - ) # (..., N, d_state) + ) # (..., num_grid_nodes, num_variables) return mask_and_reduce_metric( entry_crps, mask=mask, average_grid=average_grid, sum_vars=sum_vars diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index 07738396..cb87d76d 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -1,3 +1,5 @@ +"""Neural-LAM model architectures including GraphLAM, HiLAM, and variants.""" + # Local from .forecasters.autoregressive import ARForecaster from .forecasters.base import Forecaster diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py index e070c397..7ea9f6fd 100644 --- a/neural_lam/models/forecasters/__init__.py +++ b/neural_lam/models/forecasters/__init__.py @@ -1,3 +1,7 @@ +""" +Forecasters for the Neural-LAM model. +""" + # Local from .autoregressive import ARForecaster from .base import Forecaster diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py index e033fb74..daa2081d 100644 --- a/neural_lam/models/forecasters/autoregressive.py +++ b/neural_lam/models/forecasters/autoregressive.py @@ -1,5 +1,6 @@ +"""Forecaster that uses an auto-regressive strategy to unroll a forecast.""" + # Standard library -from typing import Optional # Third-party import torch @@ -17,6 +18,16 @@ class ARForecaster(Forecaster): """ def __init__(self, predictor: StepPredictor, datastore: BaseDatastore): + """ + Initialize the ARForecaster. + + Parameters + ---------- + predictor : StepPredictor + The predictor to use for each step. + datastore : BaseDatastore + The datastore providing grid metadata and boundary masks. + """ super().__init__() self.predictor = predictor @@ -33,6 +44,15 @@ def __init__(self, predictor: StepPredictor, datastore: BaseDatastore): @property def predicts_std(self) -> bool: + """ + Whether the forecaster predicts standard deviation. + + Returns + ------- + bool + ``True`` if the forecaster predicts standard deviation, + ``False`` otherwise. + """ return self.predictor.predicts_std def forward( @@ -40,45 +60,46 @@ def forward( init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Unroll the autoregressive model: at each step ``i`` call - ``self.predictor`` to produce the next state, then overwrite - boundary nodes with the true value from ``boundary_states[:, i]``. + ``self.predictor`` to produce the next state, then overwrite boundary + nodes with the true value from ``boundary_states[:, i]``. Parameters ---------- init_states : torch.Tensor - Shape ``(B, 2, num_grid_nodes, d_f)``. The two initial states - ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: ``B`` is - batch size, ``2`` is the time index, ``num_grid_nodes`` is the - number of spatial nodes, and ``d_f`` is the state feature - dimension. + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: + ``B`` is batch size, ``2`` initial time steps (``X_{t-1}, X_t``), + ``num_grid_nodes`` is the number of spatial nodes, and + ``num_state_vars`` is the number of state variables. forcing_features : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. Forcing - features for each predicted step; ``pred_steps`` defines the - rollout length. Dims: ``B`` is batch size, ``pred_steps`` is + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + Forcing features for each predicted step; ``pred_steps`` defines + the rollout length. Dims: ``B`` is batch size, ``pred_steps`` is the number of predicted steps, ``num_grid_nodes`` is the - number of spatial nodes, and ``d_forcing`` is the forcing - feature dimension (already concatenated past/current/future + number of spatial nodes, and ``num_forcing_vars`` is the number + of forcing variables (already concatenated past/current/future windows). boundary_states : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, d_f)``. True state - values used ONLY to overwrite boundary nodes at each AR step. - The interior prediction at step ``i`` must not depend on - ``boundary_states[:, i]`` in any other way. Dims: ``B`` is - batch size, ``pred_steps`` is the number of predicted steps, - ``num_grid_nodes`` is the number of spatial nodes, and - ``d_f`` is the state feature dimension. + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. + True state values used ONLY to overwrite boundary nodes at + each AR step. The interior prediction at step ``i`` must not + depend on ``boundary_states[:, i]`` in any other way. Dims: + ``B`` is batch size, ``pred_steps`` is the number of + predicted steps, ``num_grid_nodes`` is the number of spatial + nodes, and + ``num_state_vars`` is the state feature dimension. Returns ------- prediction : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, d_f)``. Stacked + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. Stacked per-step forecasts (with boundary overwritten by the true value). Dims: same as ``boundary_states``. pred_std : torch.Tensor or None - Shape ``(B, pred_steps, num_grid_nodes, d_f)`` when the + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when the wrapped predictor outputs an std, otherwise ``None`` (in which case ``ForecasterModule`` substitutes the constant per-variable std). Dims: same as ``prediction``. diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py index eb5483a1..4d957916 100644 --- a/neural_lam/models/forecasters/base.py +++ b/neural_lam/models/forecasters/base.py @@ -1,6 +1,7 @@ +"""Base class for forecasters.""" + # Standard library from abc import ABC, abstractmethod -from typing import Optional # Third-party import torch @@ -17,7 +18,15 @@ class Forecaster(nn.Module, ABC): @property @abstractmethod def predicts_std(self) -> bool: - """Whether this forecaster outputs a predicted standard deviation.""" + """ + Whether this forecaster outputs a predicted standard deviation. + + Returns + ------- + bool + ``True`` if the forecaster predicts standard deviation, + ``False`` otherwise. + """ @abstractmethod def forward( @@ -25,7 +34,7 @@ def forward( init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Produce a forecast of length ``pred_steps`` from two initial states, the per-step forcing features, and the per-step true boundary states. @@ -33,42 +42,40 @@ def forward( Parameters ---------- init_states : torch.Tensor - Shape ``(B, 2, num_grid_nodes, d_f)``. The two initial states - ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: ``B`` is - batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), - ``num_grid_nodes`` is the number of spatial nodes, and ``d_f`` - is the state feature dimension. - forcing_features : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. External - forcings provided at each predicted step. Dims: ``B`` is batch - size, ``pred_steps`` is the autoregressive rollout length, + Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial + states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: + ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``), ``num_grid_nodes`` is the number of spatial nodes, and - ``d_forcing`` is the forcing feature dimension (already + ``num_state_vars`` is the state feature dimension. + forcing_features : torch.Tensor + Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. + External forcings provided at each predicted step. Dims: ``B`` + is batch size, ``pred_steps`` is the autoregressive rollout + length, ``num_grid_nodes`` is the number of spatial nodes, and + ``num_forcing_vars`` is the forcing feature dimension (already concatenated past/current/future windows). boundary_states : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, d_f)``. True state - values used ONLY to overwrite boundary nodes at each AR step - — interior predictions must not depend on ``boundary_states`` - in any other way. Dims: ``B`` is batch size, ``pred_steps`` - is the rollout length, ``num_grid_nodes`` is the number of - spatial nodes, and ``d_f`` is the state feature dimension. + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True + state values used ONLY to overwrite boundary nodes at each AR + step; interior predictions must not depend on ``boundary_states`` + in any other way. Dims: ``B`` is batch size, ``pred_steps`` is + the rollout length, ``num_grid_nodes`` is the number of spatial + nodes, and ``num_state_vars`` is the state feature dimension. This is a temporary mechanism that mirrors the pre-refactor ARModel behavior; it will be replaced by a dedicated - boundary-forcing input in #138 (training on interior + - boundary datastore), at which point this parameter will be - removed. + boundary-forcing input in #138 (training on interior + boundary + datastore), at which point this parameter will be removed. Returns ------- prediction : torch.Tensor - Shape ``(B, pred_steps, num_grid_nodes, d_f)``. Forecast of - state at each predicted step. Dims: same as + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. + Forecast of state at each predicted step. Dims: same as ``boundary_states``. pred_std : torch.Tensor or None - Shape ``(B, pred_steps, num_grid_nodes, d_f)`` when + Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when ``predicts_std`` is True, otherwise ``None``. Per-feature predicted standard deviation; when ``None``, the constant per-variable std is substituted upstream by ``ForecasterModule``. Dims: same as ``prediction``. """ - pass diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 215edab8..71ce7951 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -1,7 +1,9 @@ +"""Lightning module handling training, validation and testing loops.""" + # Standard library import os import warnings -from typing import Any, Dict, List, Optional +from typing import Any # Third-party import matplotlib.pyplot as plt @@ -41,13 +43,50 @@ def __init__( restore_opt: bool = False, n_example_pred: int = 1, create_gif: bool = False, - val_steps_to_log: Optional[List[int]] = None, - metrics_watch: Optional[List[str]] = None, - var_leads_metrics_watch: Optional[Dict[int, List[int]]] = None, + val_steps_to_log: list[int] | None = None, + metrics_watch: list[str] | None = None, + var_leads_metrics_watch: dict[int, list[int]] | None = None, args=None, ): + """ + Initialize the ForecasterModule. + + Parameters + ---------- + forecaster : Forecaster + The forecaster model to use for predictions. + config : NeuralLAMConfig + Configuration object for the neural LAM model. + datastore : BaseDatastore + Datastore providing grid metadata and data access. + loss : str, default "wmse" + The loss function to use. + lr : float, default 1e-3 + Learning rate for the optimizer. + restore_opt : bool, default False + Whether to restore optimizer state from checkpoint. + n_example_pred : int, default 1 + Number of example predictions to plot during testing. + create_gif : bool, default False + Whether to create GIFs of example predictions. + val_steps_to_log : list of int, optional + Specific rollout steps to log during validation/testing. + metrics_watch : list of str, optional + List of metrics to watch and log specifically. + var_leads_metrics_watch : dict of {int: list of int}, optional + Mapping from variable index to a list of rollout steps to log + individually for the configured metrics. + args : argparse.Namespace, optional + Pre-refactor ``ARModel`` checkpoint hyperparameters. When + provided, attributes on ``args`` take precedence over the + corresponding explicit kwargs (``loss``, ``lr``, ``restore_opt``, + ``n_example_pred``, ``create_gif``, ``val_steps_to_log``, + ``metrics_watch``, ``var_leads_metrics_watch``) so legacy + checkpoints round-trip through ``load_from_checkpoint`` + correctly. + """ super().__init__() - # Pre-refactor ARModel checkpoints saved every hyperparameter nested + # Pre-refactor ``ARModel`` checkpoints saved every hyperparameter nested # inside an argparse Namespace under the single key 'args'. When # Lightning calls __init__ during load_from_checkpoint it would # otherwise drop 'args' (not in the new signature) and silently fall @@ -77,11 +116,11 @@ def __init__( if var_leads_metrics_watch is None: var_leads_metrics_watch = {} - # datastore and forecaster are excluded from saved hparams and must be - # provided explicitly when calling load_from_checkpoint. Saving args - # makes the checkpoint self-describing: it carries model, graph_name, - # hidden_dim, etc. so the caller can reconstruct the exact forecaster - # architecture from the checkpoint alone. + # datastore and forecaster are excluded from saved hparams and must + # be provided explicitly when calling load_from_checkpoint. Saving + # args makes the checkpoint self-describing: it carries model, + # graph_name, hidden_dim, etc. so the caller can reconstruct the + # exact forecaster architecture from the checkpoint alone. self.save_hyperparameters(ignore=["datastore", "forecaster"]) self.datastore = datastore self.forecaster = forecaster @@ -171,10 +210,10 @@ def __init__( # Instantiate loss function self.loss = metrics.get_metric(loss) - self.val_metrics: Dict[str, List] = { + self.val_metrics: dict[str, list] = { "mse": [], } - self.test_metrics: Dict[str, List] = { + self.test_metrics: dict[str, list] = { "mse": [], "mae": [], } @@ -190,7 +229,7 @@ def __init__( self.plotted_examples = 0 # For storing spatial loss maps during evaluation - self.spatial_loss_maps: List[Any] = [] + self.spatial_loss_maps: list[Any] = [] # Warn once per phase if val_steps_to_log exceeds the actual rollout self._val_steps_warn_issued = False @@ -207,6 +246,25 @@ def _create_dataarray_from_tensor( split: str, category: str, ) -> xr.DataArray: + """ + Create an xarray DataArray from a torch tensor. + + Parameters + ---------- + tensor : torch.Tensor + The tensor to convert. + time : torch.Tensor + The time coordinates for the data. + split : str + The data split (e.g., 'train', 'val', 'test'). + category : str + The category of data (e.g., 'state', 'forcing'). + + Returns + ------- + xr.DataArray + The resulting xarray DataArray. + """ weather_dataset = WeatherDataset(datastore=self.datastore, split=split) time = np.array(time.cpu(), dtype="datetime64[ns]") da = weather_dataset.create_dataarray_from_tensor( @@ -215,6 +273,14 @@ def _create_dataarray_from_tensor( return da def configure_optimizers(self): + """ + Configure the optimizers and learning rate schedulers. + + Returns + ------- + torch.optim.Optimizer + The configured optimizer. + """ opt = torch.optim.AdamW( self.parameters(), lr=self.hparams.lr, betas=(0.9, 0.95) ) @@ -271,6 +337,21 @@ def on_after_batch_transfer(self, batch, dataloader_idx): return init_states, target_states, forcing, batch_times def common_step(self, batch): + """ + Perform a common prediction step for training, validation, and testing. + + Parameters + ---------- + batch : tuple + The batch of data containing initial states, target states, + forcing features, and batch times. + + Returns + ------- + tuple + A tuple containing prediction, target states, predicted standard + deviation, and batch times. + """ init_states, target_states, forcing_features, batch_times = batch prediction, pred_std = self.forecaster( init_states, forcing_features, target_states @@ -278,6 +359,19 @@ def common_step(self, batch): return prediction, target_states, pred_std, batch_times def training_step(self, batch): + """ + Perform a single training step. + + Parameters + ---------- + batch : tuple + The batch of data. + + Returns + ------- + torch.Tensor + The computed loss for the training step. + """ prediction, target_states, pred_std, _ = self.common_step(batch) if pred_std is None: pred_std = self.per_var_std @@ -303,6 +397,19 @@ def training_step(self, batch): return batch_loss def all_gather_cat(self, tensor_to_gather): + """ + Gather tensors from all GPUs and concatenate them. + + Parameters + ---------- + tensor_to_gather : torch.Tensor + The tensor to gather from all processes. + + Returns + ------- + torch.Tensor + The concatenated tensor from all processes. + """ gathered = self.all_gather(tensor_to_gather) # all_gather adds dim 0 only on multi-device; on single # device it returns the same tensor unchanged. @@ -329,6 +436,16 @@ def _warn_skipped_val_steps(self, pred_steps: int, phase: str) -> None: setattr(self, flag, True) def validation_step(self, batch, batch_idx): + """ + Perform a single validation step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ prediction, target_states, pred_std, _ = self.common_step(batch) if pred_std is None: pred_std = self.per_var_std @@ -369,6 +486,10 @@ def validation_step(self, batch, batch_idx): self.val_metrics["mse"].append(entry_mses) def on_validation_epoch_end(self): + """ + Perform actions at the end of the validation epoch. + Aggregates and plots validation metrics. + """ self.aggregate_and_plot_metrics(self.val_metrics, prefix="val") if self.trainer.is_global_zero and self.hparams.metrics_watch: @@ -388,6 +509,16 @@ def on_validation_epoch_end(self): # pylint: disable-next=unused-argument def test_step(self, batch, batch_idx): + """ + Perform a single test step. + + Parameters + ---------- + batch : tuple + The batch of data. + batch_idx : int + The index of the batch. + """ prediction, target_states, pred_std, _ = self.common_step(batch) if pred_std is not None: @@ -467,6 +598,20 @@ def test_step(self, batch, batch_idx): ) def plot_examples(self, batch, n_examples, split, prediction): + """ + Plot example predictions. + + Parameters + ---------- + batch : tuple + The batch of data. + n_examples : int + Number of examples to plot. + split : str + The data split. + prediction : torch.Tensor + The model predictions. + """ target = batch[1] time = batch[3] @@ -532,7 +677,7 @@ def plot_examples(self, batch, n_examples, split, prediction): f"example_plots_{example_i}", ) os.makedirs(plot_dir_path, exist_ok=True) - png_frames: Dict[str, List[str]] = { + png_frames: dict[str, list[str]] = { var_name: [] for var_name in self.datastore.get_vars_names("state") } @@ -623,6 +768,23 @@ def plot_examples(self, batch, n_examples, split, prediction): ) def create_metric_log_dict(self, metric_tensor, prefix, metric_name): + """ + Create a dictionary of metrics to log. + + Parameters + ---------- + metric_tensor : torch.Tensor + The metric values. + prefix : str + Prefix for the log names (e.g., 'val', 'test'). + metric_name : str + The name of the metric. + + Returns + ------- + dict + Dictionary of logged metrics and figures. + """ log_dict = {} metric_fig = vis.plot_error_heatmap( errors=metric_tensor, @@ -656,6 +818,16 @@ def create_metric_log_dict(self, metric_tensor, prefix, metric_name): return log_dict def aggregate_and_plot_metrics(self, metrics_dict, prefix): + """ + Aggregate metrics from all GPUs and plot them. + + Parameters + ---------- + metrics_dict : dict + Dictionary of metrics lists. + prefix : str + Prefix for the log names. + """ log_dict = {} for metric_name, metric_val_list in metrics_dict.items(): metric_tensor = self.all_gather_cat( @@ -713,6 +885,10 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix): plt.close("all") def on_test_epoch_end(self): + """ + Perform actions at the end of the test epoch. + Aggregates and plots test metrics and spatial loss maps. + """ self.aggregate_and_plot_metrics(self.test_metrics, prefix="test") spatial_loss_tensor = self.all_gather_cat( @@ -786,11 +962,21 @@ def on_test_epoch_end(self): self.plotted_examples = 0 def on_load_checkpoint(self, checkpoint): + """ + Perform actions when loading a checkpoint. + Handles backward compatibility for older checkpoints. + + Parameters + ---------- + checkpoint : dict + The loaded checkpoint dictionary. + """ loaded_state_dict = checkpoint["state_dict"] # 1. Broad namespace remap: for pre-refactor checkpoints - # The old ARModel was a flat LightningModule. Everything that belonged - # to the predictor needs to be moved to 'forecaster.predictor.' + # The old ``ARModel`` was a flat LightningModule. Everything that + # belonged to the predictor needs to be moved to + # 'forecaster.predictor.' old_keys = list(loaded_state_dict.keys()) for key in old_keys: if not key.startswith("forecaster.") and key not in ( diff --git a/neural_lam/models/step_predictors/__init__.py b/neural_lam/models/step_predictors/__init__.py index 7acaf62f..4ddc8b98 100644 --- a/neural_lam/models/step_predictors/__init__.py +++ b/neural_lam/models/step_predictors/__init__.py @@ -1,2 +1,6 @@ +""" +Step predictors for the Neural-LAM model. +""" + # Local from .base import StepPredictor diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py index 8e9b4fc4..e5686958 100644 --- a/neural_lam/models/step_predictors/base.py +++ b/neural_lam/models/step_predictors/base.py @@ -1,6 +1,7 @@ +"""Base class for step predictors.""" + # Standard library from abc import ABC, abstractmethod -from typing import Dict, Optional # Third-party import torch @@ -21,14 +22,28 @@ def __init__( self, datastore: BaseDatastore, output_std: bool = False, - output_clamping_lower: Optional[Dict[str, float]] = None, - output_clamping_upper: Optional[Dict[str, float]] = None, + output_clamping_lower: dict[str, float] | None = None, + output_clamping_upper: dict[str, float] | None = None, ): + """ + Initialize the StepPredictor. + + Parameters + ---------- + datastore : BaseDatastore + The datastore providing grid metadata and data access. + output_std : bool, default False + Whether to output a predicted standard deviation. + output_clamping_lower : dict, optional + Lower clamping limits for state variables. + output_clamping_upper : dict, optional + Upper clamping limits for state variables. + """ super().__init__() - self._output_clamping_lower: Dict[str, float] = ( + self._output_clamping_lower: dict[str, float] = ( dict(output_clamping_lower) if output_clamping_lower else {} ) - self._output_clamping_upper: Dict[str, float] = ( + self._output_clamping_upper: dict[str, float] = ( dict(output_clamping_upper) if output_clamping_upper else {} ) @@ -80,7 +95,15 @@ def __init__( @property def predicts_std(self) -> bool: - """Whether this predictor outputs a predicted standard deviation.""" + """ + Whether this predictor outputs a predicted standard deviation. + + Returns + ------- + bool + ``True`` if the predictor predicts standard deviation, + ``False`` otherwise. + """ return self.output_std def expand_to_batch(self, x: torch.Tensor, batch_size: int) -> torch.Tensor: @@ -90,8 +113,8 @@ def expand_to_batch(self, x: torch.Tensor, batch_size: int) -> torch.Tensor: Parameters ---------- x : torch.Tensor - Shape ``(N, d)``. Tensor to expand. Dims: ``N`` is the number - of nodes and ``d`` is the feature dimension. + Shape ``(N, d)``. Tensor to expand. Dims: ``N`` is the number of + nodes and ``d`` is the feature dimension. batch_size : int Target batch size ``B``. @@ -108,7 +131,7 @@ def forward( prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Advance the state by one step: ``(X_{t-1}, X_t, forcing_t) -> X_{t+1}``. @@ -116,32 +139,31 @@ def forward( Parameters ---------- prev_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. The current state + Shape ``(B, num_grid_nodes, num_state_vars)``. The current state ``X_t``. Dims: ``B`` is batch size, ``num_grid_nodes`` is the - number of spatial nodes, and ``d_f`` is the number of state - variables. + number of spatial nodes, and ``num_state_vars`` is the number of + state variables. prev_prev_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. The previous state + Shape ``(B, num_grid_nodes, num_state_vars)``. The previous state ``X_{t-1}``, used as additional conditioning. Dims: same as ``prev_state``. forcing : torch.Tensor - Shape ``(B, num_grid_nodes, d_forcing)``. External forcings + Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings for this step (already concatenated past/current/future windows). Dims: ``B`` is batch size, ``num_grid_nodes`` is - the number of spatial nodes, and ``d_forcing`` is the + the number of spatial nodes, and ``num_forcing_vars`` is the forcing feature dimension. Returns ------- pred_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. The predicted next + Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted next state ``X_{t+1}``. Dims: same as ``prev_state``. pred_std : torch.Tensor or None - Shape ``(B, num_grid_nodes, d_f)`` when ``output_std`` + Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` is True, otherwise ``None``. Per-feature predicted standard deviation. Dims: same as ``prev_state``. """ - pass def prepare_clamping_params(self, datastore: BaseDatastore): """ @@ -151,6 +173,11 @@ def prepare_clamping_params(self, datastore: BaseDatastore): ``self._output_clamping_lower`` and ``self._output_clamping_upper`` (set in ``__init__``) and registers the buffers and clamping functions used by ``get_clamped_new_state``. + + Parameters + ---------- + datastore : BaseDatastore + The datastore providing variable names. """ # Read clamping limits stored on self @@ -177,10 +204,21 @@ def prepare_clamping_params(self, datastore: BaseDatastore): sigmoid_center = 0 softplus_center = 0 - normalize_clamping_lim = ( - lambda x, feature_idx: (x - self.state_mean[feature_idx]) - / self.state_std[feature_idx] - ) + def normalize_clamping_lim(x, feature_idx): + """Normalize a clamping limit from the original feature space to the + standardized space of the model's output. + + Parameters + ---------- + x : float + The clamping limit in the original feature space. + feature_idx : int + The index of the feature this limit applies to, used to look up + the mean and std for normalization. + """ + return (x - self.state_mean[feature_idx]) / self.state_std[ + feature_idx + ] # Check which clamping functions to use for each feature sigmoid_lower_upper_idx = [] @@ -294,19 +332,19 @@ def get_clamped_new_state(self, state_delta, prev_state): Parameters ---------- state_delta : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. Raw predicted state + Shape ``(B, num_grid_nodes, num_state_vars)``. Raw predicted state increment (network output, already rescaled). Dims: ``B`` is batch size, ``num_grid_nodes`` is the number of spatial nodes, - and ``d_f`` is the number of state variables. + and ``num_state_vars`` is the number of state variables. prev_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. Current state ``X_t`` + Shape ``(B, num_grid_nodes, num_state_vars)``. Current state ``X_t`` used as the base for the clamped update. Dims: same as ``state_delta``. Returns ------- torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. Clamped next state. + Shape ``(B, num_grid_nodes, num_state_vars)``. Clamped next state. Dims: same as ``state_delta``. """ diff --git a/neural_lam/models/step_predictors/graph/__init__.py b/neural_lam/models/step_predictors/graph/__init__.py index 3145d3ff..c8bf90dc 100644 --- a/neural_lam/models/step_predictors/graph/__init__.py +++ b/neural_lam/models/step_predictors/graph/__init__.py @@ -1,3 +1,7 @@ +""" +Graph-based step predictors. +""" + # Local from .base import BaseGraphModel from .graph_lam import GraphLAM diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 4522d307..f7fb53a7 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/neural_lam/models/step_predictors/graph/base.py @@ -1,5 +1,6 @@ +"""Base class for graph-based step predictors.""" + # Standard library -from typing import Dict, Optional # Third-party import torch @@ -28,11 +29,40 @@ def __init__( num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, - output_clamping_lower: Optional[Dict[str, float]] = None, - output_clamping_upper: Optional[Dict[str, float]] = None, + output_clamping_lower: dict[str, float] | None = None, + output_clamping_upper: dict[str, float] | None = None, g2m_gnn_type: str = "InteractionNet", m2g_gnn_type: str = "InteractionNet", ): + """ + Initialize the BaseGraphModel. + + Parameters + ---------- + datastore : BaseDatastore + Datastore supplying data and information about dataset and + forecast region. + graph_name : str, default "multiscale" + The name of the graph to load. + hidden_dim : int, default 64 + The dimension of the hidden representations. + hidden_layers : int, default 1 + The number of hidden layers in the MLPs. + processor_layers : int, default 4 + The number of processor layers in the GNN. + mesh_aggr : str, default "sum" + The aggregation method for mesh nodes. + num_past_forcing_steps : int, default 1 + The number of past forcing steps to include. + num_future_forcing_steps : int, default 1 + The number of future forcing steps to include. + output_std : bool, default False + Whether to output a predicted standard deviation. + output_clamping_lower : dict, optional + Lower clamping limits for state variables. + output_clamping_upper : dict, optional + Upper clamping limits for state variables. + """ super().__init__( datastore=datastore, output_std=output_std, @@ -143,7 +173,14 @@ def __init__( def get_num_mesh(self): """ Compute number of mesh nodes from loaded features, - and number of mesh nodes that should be ignored in encoding/decoding + and number of mesh nodes that should be ignored in encoding/decoding. + + Returns + ------- + num_mesh_nodes : int + The number of mesh nodes. + num_ignore_mesh_nodes : int + The number of mesh nodes to ignore. """ raise NotImplementedError("get_num_mesh not implemented") @@ -154,9 +191,9 @@ def embedd_mesh_nodes(self): Returns ------- torch.Tensor - Shape ``(num_mesh_nodes, d_h)``. Embedded mesh node + Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node representations. Dims: ``num_mesh_nodes`` is the number of - mesh nodes and ``d_h`` is the hidden dimension. + mesh nodes and ``hidden_dim`` is the hidden dimension. """ raise NotImplementedError("embedd_mesh_nodes not implemented") @@ -168,15 +205,15 @@ def process_step(self, mesh_rep): Parameters ---------- mesh_rep : torch.Tensor - Shape ``(B, num_mesh_nodes, d_h)``. Current mesh node + Shape ``(B, num_mesh_nodes, hidden_dim)``. Current mesh node representations. Dims: ``B`` is batch size, - ``num_mesh_nodes`` is the number of mesh nodes, and ``d_h`` + ``num_mesh_nodes`` is the number of mesh nodes, and ``hidden_dim`` is the hidden dimension. Returns ------- torch.Tensor - Shape ``(B, num_mesh_nodes, d_h)``. Updated mesh node + Shape ``(B, num_mesh_nodes, hidden_dim)``. Updated mesh node representations. Dims: same as ``mesh_rep``. """ raise NotImplementedError("process_step not implemented") @@ -193,29 +230,29 @@ def forward(self, prev_state, prev_prev_state, forcing): Parameters ---------- prev_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. The current state - ``X_t``. Dims: ``B`` is batch size, ``num_grid_nodes`` is the - number of spatial grid nodes, and ``d_f`` is the number of - state variables. + Shape ``(B, num_grid_nodes, num_state_vars)``. The current + state ``X_t``. Dims: ``B`` is batch size, + ``num_grid_nodes`` is the number of spatial grid nodes, and + ``num_state_vars`` is the number of state variables. prev_prev_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. The previous state + Shape ``(B, num_grid_nodes, num_state_vars)``. The previous state ``X_{t-1}``, used as additional conditioning. Dims: same as ``prev_state``. forcing : torch.Tensor - Shape ``(B, num_grid_nodes, d_forcing)``. External forcings + Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings for this step (already concatenated past/current/future windows). Dims: ``B`` is batch size, ``num_grid_nodes`` is - the number of spatial grid nodes, and ``d_forcing`` is the + the number of spatial grid nodes, and ``num_forcing_vars`` is the forcing feature dimension. Returns ------- new_state : torch.Tensor - Shape ``(B, num_grid_nodes, d_f)``. The predicted next state - ``X_{t+1}`` after delta-add and clamping. Dims: same as - ``prev_state``. + Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted + next state ``X_{t+1}`` after delta-add and clamping. Dims: + same as ``prev_state``. pred_std : torch.Tensor or None - Shape ``(B, num_grid_nodes, d_f)`` when ``output_std`` is + Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` is True, otherwise ``None``. Per-feature predicted standard deviation (raw softplus output, not rescaled by diff statistics). Dims: same as ``prev_state``. @@ -234,25 +271,31 @@ def forward(self, prev_state, prev_prev_state, forcing): ) # Embed all features - grid_emb = self.grid_embedder(grid_features) # (B, num_grid_nodes, d_h) - g2m_emb = self.g2m_embedder(self.g2m_features) # (M_g2m, d_h) - m2g_emb = self.m2g_embedder(self.m2g_features) # (M_m2g, d_h) + grid_emb = self.grid_embedder( + grid_features + ) # (B, num_grid_nodes, hidden_dim) + g2m_emb = self.g2m_embedder( + self.g2m_features + ) # (num_edges, hidden_dim) + m2g_emb = self.m2g_embedder( + self.m2g_features + ) # (num_edges, hidden_dim) mesh_emb = self.embedd_mesh_nodes() # Map from grid to mesh mesh_emb_expanded = self.expand_to_batch( mesh_emb, batch_size - ) # (B, num_mesh_nodes, d_h) + ) # (B, num_mesh_nodes, hidden_dim) g2m_emb_expanded = self.expand_to_batch(g2m_emb, batch_size) # This also splits representation into grid and mesh mesh_rep = self.g2m_gnn( grid_emb, mesh_emb_expanded, g2m_emb_expanded - ) # (B, num_mesh_nodes, d_h) + ) # (B, num_mesh_nodes, hidden_dim) # Also MLP with residual for grid representation grid_rep = grid_emb + self.encoding_grid_mlp( grid_emb - ) # (B, num_grid_nodes, d_h) + ) # (B, num_grid_nodes, hidden_dim) # Run processor step mesh_rep = self.process_step(mesh_rep) @@ -261,7 +304,7 @@ def forward(self, prev_state, prev_prev_state, forcing): m2g_emb_expanded = self.expand_to_batch(m2g_emb, batch_size) grid_rep = self.m2g_gnn( mesh_rep, grid_rep, m2g_emb_expanded - ) # (B, num_grid_nodes, d_h) + ) # (B, num_grid_nodes, hidden_dim) # Map to output dimension, only for grid net_output = self.output_map( @@ -271,7 +314,7 @@ def forward(self, prev_state, prev_prev_state, forcing): if self.output_std: pred_delta_mean, pred_std_raw = net_output.chunk( 2, dim=-1 - ) # both (B, num_grid_nodes, d_f) + ) # both (B, num_grid_nodes, num_state_vars) # NOTE: The predicted std. is not scaled in any way here # linter for some reason does not think softplus is callable # pylint: disable-next=not-callable diff --git a/neural_lam/models/step_predictors/graph/graph_lam.py b/neural_lam/models/step_predictors/graph/graph_lam.py index 48fa4296..b8ba09b6 100644 --- a/neural_lam/models/step_predictors/graph/graph_lam.py +++ b/neural_lam/models/step_predictors/graph/graph_lam.py @@ -1,5 +1,6 @@ +"""Graph-based LAM model with a flat mesh.""" + # Standard library -from typing import Dict, Optional # Third-party import torch_geometric as pyg @@ -30,11 +31,39 @@ def __init__( num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, - output_clamping_lower: Optional[Dict[str, float]] = None, - output_clamping_upper: Optional[Dict[str, float]] = None, + output_clamping_lower: dict[str, float] | None = None, + output_clamping_upper: dict[str, float] | None = None, g2m_gnn_type: str = "InteractionNet", m2g_gnn_type: str = "InteractionNet", ): + """ + Initialize the GraphLAM model. + + Parameters + ---------- + datastore : BaseDatastore + The datastore providing grid metadata and data access. + graph_name : str, default "multiscale" + The name of the graph to load. + hidden_dim : int, default 64 + The dimension of the hidden representations. + hidden_layers : int, default 1 + The number of hidden layers in the MLPs. + processor_layers : int, default 4 + The number of processor layers in the GNN. + mesh_aggr : str, default "sum" + The aggregation method for mesh nodes. + num_past_forcing_steps : int, default 1 + The number of past forcing steps to include. + num_future_forcing_steps : int, default 1 + The number of future forcing steps to include. + output_std : bool, default False + Whether to output a predicted standard deviation. + output_clamping_lower : dict, optional + Lower clamping limits for state variables. + output_clamping_upper : dict, optional + Upper clamping limits for state variables. + """ super().__init__( datastore=datastore, graph_name=graph_name, @@ -82,7 +111,10 @@ def __init__( self.processor = pyg.nn.Sequential( "mesh_rep, edge_rep", [ - (net, "mesh_rep, mesh_rep, edge_rep -> mesh_rep, edge_rep") + ( + net, + "mesh_rep, mesh_rep, edge_rep -> mesh_rep, edge_rep", + ) for net in processor_nets ], ) @@ -90,7 +122,14 @@ def __init__( def get_num_mesh(self): """ Compute number of mesh nodes from loaded features, - and number of mesh nodes that should be ignored in encoding/decoding + and number of mesh nodes that should be ignored in encoding/decoding. + + Returns + ------- + num_mesh_nodes : int + The number of mesh nodes. + num_ignore_mesh_nodes : int + The number of mesh nodes to ignore. """ return self.mesh_static_features.shape[0], 0 @@ -101,13 +140,13 @@ def embedd_mesh_nodes(self): Returns ------- torch.Tensor - Shape ``(num_mesh_nodes, d_h)``. Embedded mesh node representations. - Dims: ``num_mesh_nodes`` is the number of mesh nodes and ``d_h`` is - the hidden dimension. + Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node + representations. Dims: ``num_mesh_nodes`` is the number of + mesh nodes and ``hidden_dim`` is the hidden dimension. """ return self.mesh_embedder( self.mesh_static_features - ) # (num_mesh_nodes, d_h) + ) # (num_mesh_nodes, hidden_dim) def process_step(self, mesh_rep): """ @@ -117,25 +156,27 @@ def process_step(self, mesh_rep): Parameters ---------- mesh_rep : torch.Tensor - Shape ``(B, num_mesh_nodes, d_h)``. Current mesh node + Shape ``(B, num_mesh_nodes, hidden_dim)``. Current mesh node representations. Dims: ``B`` is batch size, ``num_mesh_nodes`` is - the number of mesh nodes, and ``d_h`` is the hidden + the number of mesh nodes, and ``hidden_dim`` is the hidden dimension. Returns ------- torch.Tensor - Shape ``(B, num_mesh_nodes, d_h)``. Updated mesh node + Shape ``(B, num_mesh_nodes, hidden_dim)``. Updated mesh node representations. Dims: same as ``mesh_rep``. """ # Embed m2m here first batch_size = mesh_rep.shape[0] - m2m_emb = self.m2m_embedder(self.m2m_features) # (M_mesh, d_h) + m2m_emb = self.m2m_embedder( + self.m2m_features + ) # (num_edges, hidden_dim) m2m_emb_expanded = self.expand_to_batch( m2m_emb, batch_size - ) # (B, M_mesh, d_h) + ) # (B, num_edges, hidden_dim) mesh_rep, _ = self.processor( mesh_rep, m2m_emb_expanded - ) # (B, num_mesh_nodes, d_h) + ) # (B, num_mesh_nodes, hidden_dim) return mesh_rep diff --git a/neural_lam/models/step_predictors/graph/hi_lam.py b/neural_lam/models/step_predictors/graph/hi_lam.py index cc46e67d..ee3eda8b 100644 --- a/neural_lam/models/step_predictors/graph/hi_lam.py +++ b/neural_lam/models/step_predictors/graph/hi_lam.py @@ -1,5 +1,8 @@ +""" +Hierarchical graph-based LAM model. +""" + # Standard library -from typing import Dict, Optional # Third-party from torch import nn @@ -28,13 +31,41 @@ def __init__( num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, - output_clamping_lower: Optional[Dict[str, float]] = None, - output_clamping_upper: Optional[Dict[str, float]] = None, + output_clamping_lower: dict[str, float] | None = None, + output_clamping_upper: dict[str, float] | None = None, g2m_gnn_type: str = "InteractionNet", m2g_gnn_type: str = "InteractionNet", mesh_up_gnn_type: str = "InteractionNet", mesh_down_gnn_type: str = "InteractionNet", ): + """ + Initialize the HiLAM model. + + Parameters + ---------- + datastore : BaseDatastore + The datastore providing grid metadata and data access. + graph_name : str, default "multiscale" + The name of the graph to load. + hidden_dim : int, default 64 + The dimension of the hidden representations. + hidden_layers : int, default 1 + The number of hidden layers in the MLPs. + processor_layers : int, default 4 + The number of processor layers in the GNN. + mesh_aggr : str, default "sum" + The aggregation method for mesh nodes. + num_past_forcing_steps : int, default 1 + The number of past forcing steps to include. + num_future_forcing_steps : int, default 1 + The number of future forcing steps to include. + output_std : bool, default False + Whether to output a predicted standard deviation. + output_clamping_lower : dict, optional + Lower clamping limits for state variables. + output_clamping_upper : dict, optional + Upper clamping limits for state variables. + """ super().__init__( datastore=datastore, graph_name=graph_name, @@ -72,6 +103,11 @@ def __init__( def make_same_gnns(self): """ Make intra-level GNNs. + + Returns + ------- + nn.ModuleList + List of GNNs for each level. """ return nn.ModuleList( [ @@ -87,6 +123,11 @@ def make_same_gnns(self): def make_up_gnns(self): """ Make GNNs for processing steps up through the hierarchy. + + Returns + ------- + nn.ModuleList + List of GNNs for each inter-level gap (upwards). """ gnn_class = get_gnn_class(self.mesh_up_gnn_type) return nn.ModuleList( @@ -103,6 +144,11 @@ def make_up_gnns(self): def make_down_gnns(self): """ Make GNNs for processing steps down through the hierarchy. + + Returns + ------- + nn.ModuleList + List of GNNs for each inter-level gap (downwards). """ gnn_class = get_gnn_class(self.mesh_down_gnn_type) return nn.ModuleList( @@ -131,17 +177,19 @@ def mesh_down_step( Parameters ---------- mesh_rep_levels : list of torch.Tensor - One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``. - Node representations at each hierarchy level. Dims: ``B`` is - batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, - and ``d_h`` is the hidden dimension. + One tensor per level, each of shape + ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations + at each hierarchy level. Dims: ``B`` is batch size, + ``num_mesh_nodes[l]`` is the node count at level ``l``, and + ``hidden_dim`` is the hidden dimension. mesh_same_rep : list of torch.Tensor - One tensor per level, each of shape ``(B, M_same[l], d_h)``. - Same-level edge representations. + One tensor per level, each of shape + ``(B, num_edges[l], hidden_dim)``. Same-level edge + representations. mesh_down_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_down[l], d_h)``. Downward edge representations from - level ``l+1`` to ``l``. + ``(B, num_edges[l], hidden_dim)``. Downward edge + representations from level ``l+1`` to ``l``. down_gnns : nn.ModuleList GNNs for downward edges, one per inter-level gap. same_gnns : nn.ModuleList @@ -154,7 +202,9 @@ def mesh_down_step( """ # Run same level processing on level L mesh_rep_levels[-1], mesh_same_rep[-1] = same_gnns[-1]( - mesh_rep_levels[-1], mesh_rep_levels[-1], mesh_same_rep[-1] + mesh_rep_levels[-1], + mesh_rep_levels[-1], + mesh_same_rep[-1], ) # Let level_l go from L-1 to 0 @@ -164,25 +214,22 @@ def mesh_down_step( reversed(same_gnns[:-1]), ): # Extract representations - send_node_rep = mesh_rep_levels[ - level_l + 1 - ] # (B, num_mesh_nodes[l+1], d_h) - rec_node_rep = mesh_rep_levels[ - level_l - ] # (B, num_mesh_nodes[l], d_h) + send_node_rep = mesh_rep_levels[level_l + 1] + rec_node_rep = mesh_rep_levels[level_l] down_edge_rep = mesh_down_rep[level_l] same_edge_rep = mesh_same_rep[level_l] # Apply down GNN new_node_rep, mesh_down_rep[level_l] = down_gnn( - send_node_rep, rec_node_rep, down_edge_rep + send_node_rep, + rec_node_rep, + down_edge_rep, ) # Run same level processing on level l mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn( new_node_rep, new_node_rep, same_edge_rep ) - # (B, num_mesh_nodes[l], d_h) and (B, M_same[l], d_h) return mesh_rep_levels, mesh_same_rep, mesh_down_rep @@ -196,16 +243,18 @@ def mesh_up_step( Parameters ---------- mesh_rep_levels : list of torch.Tensor - One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``. - Node representations at each hierarchy level. Dims: ``B`` is - batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, - and ``d_h`` is the hidden dimension. + One tensor per level, each of shape + ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations + at each hierarchy level. Dims: ``B`` is batch size, + ``num_mesh_nodes[l]`` is the node count at level ``l``, and + ``hidden_dim`` is the hidden dimension. mesh_same_rep : list of torch.Tensor - One tensor per level, each of shape ``(B, M_same[l], d_h)``. - Same-level edge representations. + One tensor per level, each of shape + ``(B, num_edges[l], hidden_dim)``. Same-level edge + representations. mesh_up_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_up[l], d_h)``. Upward edge representations from + ``(B, num_edges[l], hidden_dim)``. Upward edge representations from level ``l`` to ``l+1``. up_gnns : nn.ModuleList GNNs for upward edges, one per inter-level gap. @@ -220,34 +269,33 @@ def mesh_up_step( # Run same level processing on level 0 mesh_rep_levels[0], mesh_same_rep[0] = same_gnns[0]( - mesh_rep_levels[0], mesh_rep_levels[0], mesh_same_rep[0] + mesh_rep_levels[0], + mesh_rep_levels[0], + mesh_same_rep[0], ) # Let level_l go from 1 to L for level_l, (up_gnn, same_gnn) in enumerate( - zip(up_gnns, same_gnns[1:]), start=1 + zip(up_gnns, same_gnns[1:]), + start=1, ): # Extract representations - send_node_rep = mesh_rep_levels[ - level_l - 1 - ] # (B, num_mesh_nodes[l-1], d_h) - rec_node_rep = mesh_rep_levels[ - level_l - ] # (B, num_mesh_nodes[l], d_h) + send_node_rep = mesh_rep_levels[level_l - 1] + rec_node_rep = mesh_rep_levels[level_l] up_edge_rep = mesh_up_rep[level_l - 1] same_edge_rep = mesh_same_rep[level_l] # Apply up GNN new_node_rep, mesh_up_rep[level_l - 1] = up_gnn( - send_node_rep, rec_node_rep, up_edge_rep + send_node_rep, + rec_node_rep, + up_edge_rep, ) - # (B, num_mesh_nodes[l], d_h) and (B, M_up[l-1], d_h) # Run same level processing on level l mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn( new_node_rep, new_node_rep, same_edge_rep ) - # (B, num_mesh_nodes[l], d_h) and (B, M_same[l], d_h) return mesh_rep_levels, mesh_same_rep, mesh_up_rep @@ -261,19 +309,21 @@ def hi_processor_step( Parameters ---------- mesh_rep_levels : list of torch.Tensor - One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``. - Node representations at each hierarchy level. Dims: ``B`` is - batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, - and ``d_h`` is the hidden dimension. + One tensor per level, each of shape + ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations + at each hierarchy level. Dims: ``B`` is batch size, + ``num_mesh_nodes[l]`` is the node count at level ``l``, and + ``hidden_dim`` is the hidden dimension. mesh_same_rep : list of torch.Tensor - One tensor per level, each of shape ``(B, M_same[l], d_h)``. - Same-level edge representations. + One tensor per level, each of shape + ``(B, num_edges[l], hidden_dim)``. Same-level edge + representations. mesh_up_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_up[l], d_h)``. Upward edge representations. + ``(B, num_edges[l], hidden_dim)``. Upward edge representations. mesh_down_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_down[l], d_h)``. Downward edge representations. + ``(B, num_edges[l], hidden_dim)``. Downward edge representations. Returns ------- diff --git a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py index 758a6c28..fbab93ef 100644 --- a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py +++ b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py @@ -1,5 +1,8 @@ +""" +Parallel hierarchical graph-based LAM model. +""" + # Standard library -from typing import Dict, Optional # Third-party import torch @@ -31,13 +34,41 @@ def __init__( num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, - output_clamping_lower: Optional[Dict[str, float]] = None, - output_clamping_upper: Optional[Dict[str, float]] = None, + output_clamping_lower: dict[str, float] | None = None, + output_clamping_upper: dict[str, float] | None = None, g2m_gnn_type: str = "InteractionNet", m2g_gnn_type: str = "InteractionNet", mesh_up_gnn_type: str = "InteractionNet", mesh_down_gnn_type: str = "InteractionNet", ): + """ + Initialize the HiLAMParallel model. + + Parameters + ---------- + datastore : BaseDatastore + The datastore providing grid metadata and data access. + graph_name : str, default "multiscale" + The name of the graph to load. + hidden_dim : int, default 64 + The dimension of the hidden representations. + hidden_layers : int, default 1 + The number of hidden layers in the MLPs. + processor_layers : int, default 4 + The number of processor layers in the GNN. + mesh_aggr : str, default "sum" + The aggregation method for mesh nodes. + num_past_forcing_steps : int, default 1 + The number of past forcing steps to include. + num_future_forcing_steps : int, default 1 + The number of future forcing steps to include. + output_std : bool, default False + Whether to output a predicted standard deviation. + output_clamping_lower : dict, optional + Lower clamping limits for state variables. + output_clamping_upper : dict, optional + Upper clamping limits for state variables. + """ super().__init__( datastore=datastore, graph_name=graph_name, @@ -97,19 +128,21 @@ def hi_processor_step( Parameters ---------- mesh_rep_levels : list of torch.Tensor - One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``. - Node representations at each hierarchy level. Dims: ``B`` is - batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, - and ``d_h`` is the hidden dimension. + One tensor per level, each of shape + ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at + each hierarchy level. Dims: ``B`` is batch size, + ``num_mesh_nodes[l]`` is the node count at level ``l``, and + ``hidden_dim`` is the hidden dimension. mesh_same_rep : list of torch.Tensor - One tensor per level, each of shape ``(B, M_same[l], d_h)``. - Same-level edge representations. + One tensor per level, each of shape + ``(B, num_edges[l], hidden_dim)``. Same-level edge + representations. mesh_up_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_up[l], d_h)``. Upward edge representations. + ``(B, num_edges[l], hidden_dim)``. Upward edge representations. mesh_down_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_down[l], d_h)``. Downward edge representations. + ``(B, num_edges[l], hidden_dim)``. Downward edge representations. Returns ------- @@ -119,10 +152,12 @@ def hi_processor_step( """ # First join all node and edge representations to single tensors - mesh_rep = torch.cat(mesh_rep_levels, dim=1) # (B, num_mesh_nodes, d_h) + mesh_rep = torch.cat( + mesh_rep_levels, dim=1 + ) # (B, num_mesh_nodes, hidden_dim) mesh_edge_rep = torch.cat( mesh_same_rep + mesh_up_rep + mesh_down_rep, axis=1 - ) # (B, M_mesh, d_h) + ) # (B, num_edges, hidden_dim) # Here, update mesh_*_rep and mesh_rep mesh_rep, mesh_edge_rep = self.processor(mesh_rep, mesh_edge_rep) diff --git a/neural_lam/models/step_predictors/graph/hierarchical.py b/neural_lam/models/step_predictors/graph/hierarchical.py index 592a266d..8baa99b5 100644 --- a/neural_lam/models/step_predictors/graph/hierarchical.py +++ b/neural_lam/models/step_predictors/graph/hierarchical.py @@ -1,5 +1,6 @@ +"""Base implementations for hierarchical (multi-level) graph models.""" + # Standard library -from typing import Dict, Optional # Third-party from torch import nn @@ -27,13 +28,14 @@ def __init__( num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, - output_clamping_lower: Optional[Dict[str, float]] = None, - output_clamping_upper: Optional[Dict[str, float]] = None, + output_clamping_lower: dict[str, float] | None = None, + output_clamping_upper: dict[str, float] | None = None, g2m_gnn_type: str = "InteractionNet", m2g_gnn_type: str = "InteractionNet", mesh_up_gnn_type: str = "InteractionNet", mesh_down_gnn_type: str = "InteractionNet", ): + """Extend :class:`BaseGraphModel` with hierarchical mesh structures.""" super().__init__( datastore=datastore, graph_name=graph_name, @@ -140,8 +142,13 @@ def __init__( def get_num_mesh(self): """ - Compute number of mesh nodes from loaded features, - and number of mesh nodes that should be ignored in encoding/decoding + Compute mesh node counts used for encoding and decoding. + + Returns + ------- + tuple[int, int] + Total number of mesh nodes and the number to ignore during + encoding/decoding. """ num_mesh_nodes = sum( node_feat.shape[0] for node_feat in self.mesh_static_features @@ -159,9 +166,9 @@ def embedd_mesh_nodes(self): Returns ------- torch.Tensor - Shape ``(num_mesh_nodes[0], d_h)``. Embedded bottom-level + Shape ``(num_mesh_nodes[0], hidden_dim)``. Embedded bottom-level mesh node representations. Dims: ``num_mesh_nodes[0]`` is - the number of nodes at level 0 and ``d_h`` is the hidden + the number of nodes at level 0 and ``hidden_dim`` is the hidden dimension. """ return self.mesh_embedders[0](self.mesh_static_features[0]) @@ -174,22 +181,22 @@ def process_step(self, mesh_rep): Parameters ---------- mesh_rep : torch.Tensor - Shape ``(B, num_mesh_nodes[0], d_h)``. Bottom-level mesh + Shape ``(B, num_mesh_nodes[0], hidden_dim)``. Bottom-level mesh node representations from the encoder. Dims: ``B`` is batch size, ``num_mesh_nodes[0]`` is the number of nodes at - level 0, and ``d_h`` is the hidden dimension. + level 0, and ``hidden_dim`` is the hidden dimension. Returns ------- torch.Tensor - Shape ``(B, num_mesh_nodes[0], d_h)``. Updated bottom-level + Shape ``(B, num_mesh_nodes[0], hidden_dim)``. Updated bottom-level mesh node representations. Dims: same as ``mesh_rep``. """ batch_size = mesh_rep.shape[0] # EMBED REMAINING MESH NODES (levels >= 1) - # Create list of mesh node representations for each level, - # each of size (B, num_mesh_nodes[l], d_h) + # each of size (B, num_mesh_nodes[l], hidden_dim) mesh_rep_levels = [mesh_rep] + [ self.expand_to_batch(emb(node_static_features), batch_size) for emb, node_static_features in zip( @@ -225,10 +232,10 @@ def process_step(self, mesh_rep): # Extract representations send_node_rep = mesh_rep_levels[ level_l - 1 - ] # (B, num_mesh_nodes[l-1], d_h) + ] # (B, num_mesh_nodes[l-1], hidden_dim) rec_node_rep = mesh_rep_levels[ level_l - ] # (B, num_mesh_nodes[l], d_h) + ] # (B, num_mesh_nodes[l], hidden_dim) edge_rep = mesh_up_rep[level_l - 1] # Apply GNN @@ -238,9 +245,11 @@ def process_step(self, mesh_rep): # Update node and edge vectors in lists mesh_rep_levels[level_l] = ( - new_node_rep # (B, num_mesh_nodes[l], d_h) + new_node_rep # (B, num_mesh_nodes[l], hidden_dim) + ) + mesh_up_rep[level_l - 1] = ( + new_edge_rep # (B, num_edges[l-1], hidden_dim) ) - mesh_up_rep[level_l - 1] = new_edge_rep # (B, M_up[l-1], d_h) # - PROCESSOR - mesh_rep_levels, _, _, mesh_down_rep = self.hi_processor_step( @@ -255,10 +264,10 @@ def process_step(self, mesh_rep): # Extract representations send_node_rep = mesh_rep_levels[ level_l + 1 - ] # (B, num_mesh_nodes[l+1], d_h) + ] # (B, num_mesh_nodes[l+1], hidden_dim) rec_node_rep = mesh_rep_levels[ level_l - ] # (B, num_mesh_nodes[l], d_h) + ] # (B, num_mesh_nodes[l], hidden_dim) edge_rep = mesh_down_rep[level_l] # Apply GNN @@ -266,11 +275,11 @@ def process_step(self, mesh_rep): # Update node and edge vectors in lists mesh_rep_levels[level_l] = ( - new_node_rep # (B, num_mesh_nodes[l], d_h) + new_node_rep # (B, num_mesh_nodes[l], hidden_dim) ) # Return only bottom level representation - return mesh_rep_levels[0] # (B, num_mesh_nodes[0], d_h) + return mesh_rep_levels[0] # (B, num_mesh_nodes[0], hidden_dim) def hi_processor_step( self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep @@ -282,22 +291,23 @@ def hi_processor_step( ---------- mesh_rep_levels : list of torch.Tensor One tensor per level, each of shape - ``(B, num_mesh_nodes[l], d_h)``. Node representations at + ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at each hierarchy level. Dims: ``B`` is batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``, and - ``d_h`` is the hidden dimension. + ``hidden_dim`` is the hidden dimension. mesh_same_rep : list of torch.Tensor - One tensor per level, each of shape ``(B, M_same[l], d_h)``. - Same-level edge representations. ``M_same[l]`` is the edge - count at level ``l``. + One tensor per level, each of shape + ``(B, num_edges[l], hidden_dim)``. Same-level edge + representations. ``num_edges[l]`` is the edge count at + level ``l``. mesh_up_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_up[l], d_h)``. Upward edge representations from + ``(B, num_edges[l], hidden_dim)``. Upward edge representations from level ``l`` to ``l+1``. mesh_down_rep : list of torch.Tensor One tensor per inter-level gap, each of shape - ``(B, M_down[l], d_h)``. Downward edge representations from - level ``l+1`` to ``l``. + ``(B, num_edges[l], hidden_dim)``. Downward edge + representations from level ``l+1`` to ``l``. Returns ------- diff --git a/neural_lam/plot_graph.py b/neural_lam/plot_graph.py index f79db4ae..77c9b3cf 100644 --- a/neural_lam/plot_graph.py +++ b/neural_lam/plot_graph.py @@ -1,3 +1,5 @@ +"""Command-line utility for plotting saved Neural-LAM graphs.""" + # Standard library import os from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser @@ -29,7 +31,7 @@ def plot_graph( Parameters ---------- grid_pos : np.ndarray - Grid node positions, shape (N_grid, 2). + Grid node positions, shape (num_grid_nodes, 2). hierarchical : bool Whether the loaded graph is hierarchical. graph_ldict : dict @@ -264,7 +266,7 @@ def main() -> None: config_path=args.datastore_config_path ) - xy = datastore.get_xy("state", stacked=True) # (N_grid, 2) + xy = datastore.get_xy("state", stacked=True) # (num_grid_nodes, 2) pos_max = np.max(np.abs(xy)) grid_pos = xy / pos_max # Divide by maximum coordinate diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index da9b3993..f98065c4 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -1,3 +1,5 @@ +"""CLI entry point for training Neural-LAM models.""" + # Standard library import json import os @@ -22,7 +24,10 @@ class AdaptiveHelpFormatter(ArgumentDefaultsHelpFormatter): + """``--help`` formatter that scales the column width to the terminal.""" + def __init__(self, prog): + """Pick a help-column width based on the current terminal size.""" terminal_width = shutil.get_terminal_size(fallback=(100, 20)).columns width = max(80, min(terminal_width, 120)) help_position = min(44, width // 3) diff --git a/neural_lam/utils.py b/neural_lam/utils.py index efa52c93..942eb206 100644 --- a/neural_lam/utils.py +++ b/neural_lam/utils.py @@ -1,3 +1,5 @@ +"""Utility helpers shared across Neural-LAM training and evaluation.""" + # Standard library import datetime import os @@ -34,6 +36,16 @@ class BufferList(nn.Module): def __init__( self, buffer_tensors: list[torch.Tensor], persistent: bool = True ) -> None: + """ + Register a collection of tensors as buffers inside a module. + + Parameters + ---------- + buffer_tensors : Sequence[torch.Tensor] + Buffers to register in the order they should be indexed. + persistent : bool, optional + If ``True``, buffers are saved in checkpoints. Default ``True``. + """ super().__init__() self.n_buffers = len(buffer_tensors) for buffer_i, tensor in enumerate(buffer_tensors): @@ -41,15 +53,25 @@ def __init__( @overload def __getitem__(self, key: int) -> torch.Tensor: - pass + """Integer-indexed access overload; see the implementation below.""" @overload def __getitem__(self, key: slice) -> list[torch.Tensor]: - pass + """Slice-indexed access overload; see the implementation below.""" def __getitem__( self, key: Union[int, slice] ) -> Union[torch.Tensor, list[torch.Tensor]]: + """Return the buffer(s) at ``key``. + + Supports integer indexing (with Python-style negative indices) + and slice indexing (which returns a list of tensors). + + Raises + ------ + IndexError + If ``key`` is an out-of-range integer. + """ # Unpack slice indices and call recursively for each position if isinstance(key, slice): return [self[i] for i in range(*key.indices(len(self)))] @@ -63,17 +85,43 @@ def __getitem__( return getattr(self, f"b{key}") def __len__(self) -> int: + """Return the number of registered buffers.""" return self.n_buffers def __iter__(self) -> Iterator[torch.Tensor]: + """Iterate over the registered buffers in ascending index order.""" return (self[i] for i in range(len(self))) def __itruediv__(self, other: float) -> "BufferList": - """Divide each element in list with other""" + """ + Divide each element in list with other. + + Parameters + ---------- + other : float + The value to divide by. + + Returns + ------- + BufferList + The modified BufferList. + """ return self.__imul__(1.0 / other) def __imul__(self, other: float) -> "BufferList": - """Multiply each element in list with other""" + """ + Multiply each element in list with other. + + Parameters + ---------- + other : float + The value to multiply by. + + Returns + ------- + BufferList + The modified BufferList. + """ for buffer_tensor in self: buffer_tensor *= other @@ -82,7 +130,17 @@ def __imul__(self, other: float) -> "BufferList": def zero_index_edge_index(edge_index: torch.Tensor) -> torch.Tensor: """ - Make both sender and receiver indices of edge_index start at 0 + Make both sender and receiver indices of edge_index start at 0. + + Parameters + ---------- + edge_index : torch.Tensor + Edge index tensor of shape (2, num_edges). + + Returns + ------- + torch.Tensor + Edge index tensor with indices starting at 0. """ return edge_index - edge_index.min(dim=1, keepdim=True)[0] @@ -101,7 +159,7 @@ def zero_index_m2g( Parameters ---------- m2g_edge_index : torch.Tensor - Edge index tensor of shape (2, N_edges). + Edge index tensor of shape (2, num_edges). mesh_static_features : list of torch.Tensor Mesh node feature tensors. mesh_first : bool @@ -155,7 +213,7 @@ def zero_index_g2m( Parameters ---------- g2m_edge_index : torch.Tensor - Edge index tensor of shape (2, N_edges). + Edge index tensor of shape (2, num_edges). mesh_static_features : list of torch.Tensor Mesh node feature tensors. mesh_first : bool @@ -243,6 +301,21 @@ def load_graph( """ def loads_file(fn: str) -> Any: + """ + Load ``torch.load`` data from ``graph_dir_path``. + + Applies ``map_location`` so tensors land on the requested device. + + Parameters + ---------- + fn : str + The filename to load. + + Returns + ------- + Any + The loaded data. + """ return torch.load( os.path.join(graph_dir_path, fn), map_location=device, @@ -259,8 +332,8 @@ def loads_file(fn: str) -> Any: [zero_index_edge_index(ei) for ei in loads_file("m2m_edge_index.pt")], persistent=False, ) # List of (2, M_m2m[l]) - g2m_edge_index = loads_file("g2m_edge_index.pt") # (2, M_g2m) - m2g_edge_index = loads_file("m2g_edge_index.pt") # (2, M_m2g) + g2m_edge_index = loads_file("g2m_edge_index.pt") # (2, num_edges) + m2g_edge_index = loads_file("m2g_edge_index.pt") # (2, num_edges) # Change first indices to 0 # m2g and g2m has to be handled specially as not all mesh nodes @@ -278,13 +351,13 @@ def loads_file(fn: str) -> Any: assert g2m_edge_index.min() >= 0, "Negative node index in g2m" n_levels = len(m2m_edge_index) - hierarchical = n_levels > 1 # Nor just single level mesh graph + hierarchical = n_levels > 1 # Not just single level mesh graph # Load static edge features - # List of (M_m2m[l], d_edge_f) + # List of (M_m2m[l], input_dim) m2m_features = loads_file("m2m_features.pt") - g2m_features = loads_file("g2m_features.pt") # (M_g2m, d_edge_f) - m2g_features = loads_file("m2g_features.pt") # (M_m2g, d_edge_f) + g2m_features = loads_file("g2m_features.pt") # (num_edges, input_dim) + m2g_features = loads_file("m2g_features.pt") # (num_edges, input_dim) # Normalize by dividing with longest edge (found in m2m) longest_edge = max( @@ -312,21 +385,21 @@ def loads_file(fn: str) -> Any: for ei in loads_file("mesh_up_edge_index.pt") ], persistent=False, - ) # List of (2, M_up[l]) + ) # List of (2, num_edges[l]) mesh_down_edge_index = BufferList( [ zero_index_edge_index(ei) for ei in loads_file("mesh_down_edge_index.pt") ], persistent=False, - ) # List of (2, M_down[l]) + ) # List of (2, num_edges[l]) mesh_up_features = loads_file( "mesh_up_features.pt" - ) # List of (M_up[l], d_edge_f) + ) # List of (num_edges[l], input_dim) mesh_down_features = loads_file( "mesh_down_features.pt" - ) # List of (M_down[l], d_edge_f) + ) # List of (num_edges[l], input_dim) # Rescale mesh_up_features = BufferList(mesh_up_features, persistent=False) @@ -365,13 +438,22 @@ def loads_file(fn: str) -> Any: def make_mlp(blueprint: list[int], layer_norm: bool = True) -> nn.Sequential: """ - Create MLP from list blueprint, with - input dimensionality: blueprint[0] - output dimensionality: blueprint[-1] and - hidden layers of dimensions: blueprint[1], ..., blueprint[-2] + Construct a multilayer perceptron from a blueprint of layer widths. - if layer_norm is True, includes a LayerNorm layer at - the output (as used in GraphCast) + Parameters + ---------- + blueprint : list[int] + Sequence of layer dimensions where ``blueprint[0]`` is the input size, + ``blueprint[-1]`` is the output size, the intermediate entries specify + the hidden layer widths, and ``len(blueprint) - 2`` is the number of + hidden layers. + layer_norm : bool, optional + If ``True``, append a ``LayerNorm`` to the output as in GraphCast. + + Returns + ------- + torch.nn.Sequential + Sequential module implementing the specified MLP. """ hidden_layers = len(blueprint) - 2 assert hidden_layers >= 0, "Invalid MLP blueprint" @@ -392,7 +474,12 @@ def make_mlp(blueprint: list[int], layer_norm: bool = True) -> nn.Sequential: @cache def has_working_latex() -> bool: """ - Check if LaTeX is available or its toolchain + Check whether a LaTeX toolchain is available on the system. + + Returns + ------- + bool + ``True`` if ``latex`` and the required auxiliary tools are callable. """ # If latex/toolchain is not available, some visualizations might not render # correctly, but will at least not raise an error. Alternatively, use @@ -454,10 +541,18 @@ def has_working_latex() -> bool: def fractional_plot_bundle(fraction: float) -> dict[str, Any]: """ - Get the tueplots bundle, but with figure width as a fraction of - the page width. - """ + Return a ``tueplots`` bundle scaled to a fraction of the page width. + Parameters + ---------- + fraction : float + Denominator applied to the default NeurIPS figure width. + + Returns + ------- + dict + Matplotlib rcParams bundle with updated ``figure.figsize``. + """ usetex = has_working_latex() bundle = bundles.neurips2023(usetex=usetex, family="serif") bundle.update(figsizes.neurips2023()) @@ -473,14 +568,19 @@ def fractional_plot_bundle(fraction: float) -> dict[str, Any]: def log_on_rank_zero( msg: str, level: str = "info", *args: Any, **kwargs: Any ) -> None: - """Log a message only on rank zero using loguru logger. + """ + Log a message only on rank zero using loguru logger. Parameters ---------- msg : str The message to log. - level : str, optional - The logging level (e.g. "info", "warning", "error"). Default is "info". + level : str, default "info" + The logging level (e.g. "info", "warning", "error"). + *args : Any + Positional arguments passed to the logger. + **kwargs : Any + Keyword arguments passed to the logger. """ if rank_zero_only.rank == 0: log_fn = getattr(logger, level, logger.info) @@ -491,7 +591,14 @@ def init_training_logger_metrics( training_logger: Any, val_steps: list[int] ) -> None: """ - Set up logger metrics to track + Configure validation metric aggregation for the active training logger. + + Parameters + ---------- + training_logger : Any + Logger instance used during training. + val_steps : list of int + Autoregressive rollout lengths to log as separate metrics. """ experiment = training_logger.experiment if isinstance(training_logger, WandbLogger): @@ -511,16 +618,15 @@ def init_training_logger_metrics( def setup_training_logger( datastore: Any, args: Any, run_name: str, run_dir: str ) -> Any: - """Set up the training logger (WandB or MLFlow). + """ + Set up the training logger (WandB or MLFlow). Parameters ---------- - datastore : Datastore - Datastore object. - + datastore : Any + Datastore providing metadata for logging configuration. args : argparse.Namespace - Arguments from command line. - + Parsed training arguments controlling the logger backend. run_name : str Name of the run. @@ -531,8 +637,8 @@ def setup_training_logger( Returns ------- - training_logger : pytorch_lightning.loggers.base - Logger object. + Any + The initialized logger object. Raises ------ @@ -546,7 +652,6 @@ def setup_training_logger( This allows the same job script to be safely resubmitted on HPC systems. The run name is set to ``None`` when resuming to preserve the existing name. """ - if args.wandb_id and args.logger != "wandb": logger.warning( f"--wandb_id is set but logger is {args.logger!r}; " @@ -598,13 +703,34 @@ def inverse_softplus( x: torch.Tensor, beta: float = 1.0, threshold: float = 20.0 ) -> torch.Tensor: """ - Inverse of torch.nn.functional.softplus + Inverse of :func:`torch.nn.functional.softplus`. + + For most inputs this function is exact up to numerical precision. The + input is clamped to ensure numerical stability: values above + ``threshold / beta`` are treated as linear (which is exact in that + regime), and values near zero are clamped to avoid ``log`` of + non-positive numbers. Only near the lower clamping bound does the + result deviate from the true inverse. - Input is clamped to approximately positive values of x, and the function is - linear for inputs above x*beta for numerical stability. + Parameters + ---------- + x : torch.Tensor + Input tensor whose softplus inverse should be computed. + beta : float, optional + Softplus ``beta`` parameter that controls the sharpness. Default ``1``. + threshold : float, optional + Threshold above which the function is treated as linear for numerical + stability. Default ``20``. - Note that this torch.clamp will make gradients 0, but this is not a - problem as values of x that are this close to 0 have gradients of 0 anyhow. + Returns + ------- + torch.Tensor + Tensor containing the inverse-softplus values. + + Notes + ----- + ``torch.clamp`` will zero the gradients near the bounds, but values this + close to zero or ``threshold / beta`` already have negligible gradients. """ x_clamped = torch.clamp( x, min=torch.log(torch.tensor(1e-6 + 1)) / beta, max=threshold / beta @@ -621,12 +747,30 @@ def inverse_softplus( def inverse_sigmoid(x: torch.Tensor) -> torch.Tensor: """ - Inverse of torch.sigmoid + Inverse of ``torch.sigmoid`` with clamping for numerical stability. + + Sigmoid output takes values in ``[0, 1]``; we clamp the input slightly + within that open interval before applying ``log(x / (1 - x))``. + + Note that ``torch.clamp`` will make gradients 0 near the bounds, but + this is not a problem as values of x that are this close to 0 or 1 + have gradients of 0 anyhow. + + Parameters + ---------- + x : torch.Tensor + Input tensor assumed to contain logits after a sigmoid. - Sigmoid output takes values in [0,1], this makes sure input is just within - this interval. - Note that this torch.clamp will make gradients 0, but this is not a problem - as values of x that are this close to 0 or 1 have gradients of 0 anyhow. + Returns + ------- + torch.Tensor + Tensor containing ``log(x / (1 - x))`` after clamping away from the + saturation limits. + + Notes + ----- + ``torch.clamp`` zeroes gradients for values at the bounds, but values this + close to 0 or 1 already have negligible gradients. """ x_clamped = torch.clamp(x, min=1e-6, max=1 - 1e-6) return torch.log(x_clamped / (1 - x_clamped)) @@ -634,8 +778,7 @@ def inverse_sigmoid(x: torch.Tensor) -> torch.Tensor: def get_integer_time(tdelta: datetime.timedelta) -> tuple[int, str]: """ - Get the largest time unit that can represent the given timedelta as an - integer. + Express a :class:`datetime.timedelta` as an integer number of time units. Parameters ---------- diff --git a/neural_lam/vis.py b/neural_lam/vis.py index 757a9c1f..1e3dd415 100644 --- a/neural_lam/vis.py +++ b/neural_lam/vis.py @@ -1,6 +1,8 @@ +"""Visualization helpers for analysing Neural-LAM predictions and errors.""" + # Standard library import warnings -from typing import Optional, Union +from typing import Optional # Third-party import cartopy.crs as ccrs @@ -35,10 +37,21 @@ def _tex_safe(s: str) -> str: - """Escape TeX special characters in s if TeX rendering is currently active. + """ + Escape TeX special characters in s if TeX rendering is active. Needed because % is a TeX comment character; without escaping it would silently truncate any text that follows it (e.g. the title for r2m (%)). + + Parameters + ---------- + s : str + The string to escape. + + Returns + ------- + str + The escaped string. """ if plt.rcParams.get("text.usetex", False): s = s.replace("%", r"\%") @@ -46,7 +59,21 @@ def _tex_safe(s: str) -> str: def _compute_heatmap_layout(n_rows: int, n_cols: int) -> dict[str, float]: - """Choose figure and font sizes from the heatmap dimensions.""" + """ + Choose figure and font sizes from the heatmap dimensions. + + Parameters + ---------- + n_rows : int + Number of rows in the heatmap. + n_cols : int + Number of columns in the heatmap. + + Returns + ------- + dict[str, float] + Dictionary containing figure width, figure height, font sizes, etc. + """ max_dim = max(n_rows, n_cols) # Size the figure so each cell gets ~0.8 x 0.5 inches; floor at 8 x 4.5. @@ -74,7 +101,19 @@ def _compute_heatmap_layout(n_rows: int, n_cols: int) -> dict[str, float]: def _get_heatmap_var_labels(datastore: BaseRegularGridDatastore) -> list[str]: - """Build state-variable labels from datastore metadata.""" + """ + Build state-variable labels from datastore metadata. + + Parameters + ---------- + datastore : BaseRegularGridDatastore + The datastore containing metadata about the grid. + + Returns + ------- + list[str] + List of formatted variable labels. + """ var_names = datastore.get_vars_names(category="state") var_units = datastore.get_vars_units(category="state") return [ @@ -85,11 +124,24 @@ def _get_heatmap_var_labels(datastore: BaseRegularGridDatastore) -> list[str]: def _to_heatmap_matrix(values) -> np.ndarray: """ - Convert heatmap inputs to a `(d_f, pred_steps)` matrix. + Convert heatmap inputs to a ``(num_state_vars, pred_steps)`` matrix. + + A single-step tensor may arrive as one-dimensional ``(num_state_vars,)``, + especially in single-GPU or focused metric logging paths. In that case we + first treat it as one row of ``(pred_steps=1, num_state_vars)`` before + transposing. - A single-step tensor may arrive as one-dimensional `(d_f,)`, especially in - single-GPU or focused metric logging paths. In that case we first treat it - as one row of `(pred_steps=1, d_f)` before transposing. + Parameters + ---------- + values : array-like + The input values to convert. + Shape ``(num_state_vars,)`` or ``(pred_steps, num_state_vars)``. + + Returns + ------- + np.ndarray + The converted heatmap matrix with shape + ``(num_state_vars, pred_steps)``. """ if hasattr(values, "detach"): values = values.detach().cpu().numpy() @@ -102,7 +154,23 @@ def _to_heatmap_matrix(values) -> np.ndarray: def _get_feature_scale( ds_stats: xr.Dataset, var_name: str, n_vars: int ) -> np.ndarray | None: - """Extract a 1D per-feature scale, averaging over any extra dims.""" + """ + Extract a 1D per-feature scale, averaging over any extra dims. + + Parameters + ---------- + ds_stats : xr.Dataset + The standardization statistics dataset. + var_name : str + The name of the variable to extract scale for. + n_vars : int + The number of variables expected. + + Returns + ------- + np.ndarray or None + The extracted scale as a 1D array, or None if unavailable. + """ if var_name not in ds_stats: return None @@ -136,9 +204,39 @@ def _get_heatmap_color_values( Both modes fall back to per-variable max normalization when their required stat is unavailable, appending "[fallback]" to the colorbar label. + + Parameters + ---------- + errors_np : np.ndarray + The error values to normalize. + datastore : BaseRegularGridDatastore + The datastore containing standardization stats. + normalization : str + The normalization mode to use ('state_std' or 'diff_std'). + + Returns + ------- + tuple[np.ndarray, str, matplotlib.colors.Colormap] + A 3-tuple containing: + - color_values: The normalized values for the colormap. + - colorbar_label: The label for the colorbar. + - cmap: The colormap to use. + + Raises + ------ + ValueError + If ``normalization`` is not one of ``'state_std'`` or ``'diff_std'``. """ def _per_var_fallback(): + """ + Normalize errors by per-variable maximum value. + + Returns + ------- + tuple[np.ndarray, str, matplotlib.colors.Colormap] + Normalized errors, fallback label, and colormap. + """ max_err = errors_np.max(axis=1, keepdims=True) safe = np.where(max_err > np.finfo(float).eps, max_err, 1.0) return ( @@ -218,7 +316,21 @@ def _per_var_fallback(): def _get_annotation_text_color( value: float, image: matplotlib.image.AxesImage ) -> str: - """Choose a readable annotation color from the rendered background.""" + """ + Choose a readable annotation color from the rendered background. + + Parameters + ---------- + value : float + The numeric value at the cell. + image : matplotlib.image.AxesImage + The rendered image object to determine background color. + + Returns + ------- + str + 'white' or 'black' depending on the background luminance. + """ if not np.isfinite(value): return "black" @@ -231,21 +343,22 @@ def plot_on_axis( ax: matplotlib.axes.Axes, da: xr.DataArray, datastore: BaseRegularGridDatastore, - vmin: Optional[float] = None, - vmax: Optional[float] = None, - ax_title: Optional[str] = None, - cmap: Union[str, matplotlib.colors.Colormap] = "plasma", - boundary_alpha: Optional[float] = None, + vmin: float | None = None, + vmax: float | None = None, + ax_title: str | None = None, + cmap: str | matplotlib.colors.Colormap = "plasma", + boundary_alpha: float | None = None, crop_to_interior: bool = False, ) -> matplotlib.collections.QuadMesh: - """Plot weather state on a projection-aware axis using datastore metadata. + """ + Plot weather state on a projection-aware axis using datastore metadata. Parameters ---------- ax : matplotlib.axes.Axes The axis to plot on. Should have a cartopy projection. da : xarray.DataArray - The data to plot. Should have shape (N_grid,). + The data to plot. Should have shape (num_grid_nodes,). datastore : BaseRegularGridDatastore The datastore containing metadata about the grid. vmin : float, optional @@ -365,8 +478,9 @@ def plot_error_heatmap( Parameters ---------- errors : torch.Tensor - Shape ``(pred_steps, d_f)``. Per-step, per-variable errors. These - values are used for the numeric annotations in each cell. + Shape ``(pred_steps, num_state_vars)``. Per-step, per-variable + errors. These values are used for the numeric annotations in each + cell. datastore : BaseRegularGridDatastore Datastore providing variable names, units, and step length. title : str, optional @@ -474,7 +588,23 @@ def plot_error_map( datastore: BaseRegularGridDatastore, title: Optional[str] = None, ) -> matplotlib.figure.Figure: - """Deprecated: use :func:`plot_error_heatmap` instead.""" + """ + Deprecated: use :func:`plot_error_heatmap` instead. + + Parameters + ---------- + errors : torch.Tensor + The error values to plot. + datastore : BaseRegularGridDatastore + The datastore containing grid metadata. + title : str, optional + The title for the plot. + + Returns + ------- + matplotlib.figure.Figure + The completed heatmap figure. + """ warnings.warn( "plot_error_map is deprecated, use plot_error_heatmap instead", DeprecationWarning, @@ -500,11 +630,11 @@ def plot_prediction( Parameters ---------- datastore : BaseRegularGridDatastore - Datastore providing grid metadata and projection. + Datastore providing grid metadata and projection. da_prediction : xarray.DataArray - Shape ``(N_grid,)``. Predicted field values. + Shape ``(num_grid_nodes,)``. Predicted field values. da_target : xarray.DataArray - Shape ``(N_grid,)``. Ground-truth field values. + Shape ``(num_grid_nodes,)``. Ground-truth field values. title : str, optional Overall figure title. vrange : tuple of (float, float), optional @@ -584,15 +714,14 @@ def plot_spatial_error( Parameters ---------- error : torch.Tensor - Shape ``(N_grid,)``. Per-node error values. Dims: ``N_grid`` is - the number of grid nodes. + Error magnitudes on the flattened grid. + * **Shape**: ``(num_grid_nodes,)`` datastore : BaseRegularGridDatastore - Datastore providing grid metadata and projection. - title : str, optional - Figure title. - vrange : tuple of (float, float), optional - ``(vmin, vmax)`` for the colour scale. Inferred from data if not - given. + Datastore providing coordinate metadata and boundary masks. + title : str or None, optional + Optional figure title. + vrange : tuple[float, float] or None, optional + Explicit value range ``(vmin, vmax)`` for the color scale. boundary_alpha : float, optional Alpha transparency for the boundary overlay (default 0.7). crop_to_interior : bool, optional @@ -603,9 +732,8 @@ def plot_spatial_error( Returns ------- matplotlib.figure.Figure - The completed spatial error figure. + Figure handle containing the plotted map. """ - error_np = error.detach().cpu().numpy() if vrange is None: diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py index bfc95625..4168396a 100644 --- a/neural_lam/weather_dataset.py +++ b/neural_lam/weather_dataset.py @@ -1,3 +1,5 @@ +"""Dataset helpers wrapping Neural-LAM datastores for PyTorch Lightning.""" + # Standard library import datetime import warnings @@ -16,31 +18,8 @@ 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. + Loads and processes weather data from a given datastore. See + :meth:`__init__` for the full parameter list. """ def __init__( @@ -52,6 +31,36 @@ def __init__( num_future_forcing_steps: int = 1, load_single_member: bool = False, ) -> None: + """ + 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``. + 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. + """ super().__init__() self.split = split @@ -113,6 +122,14 @@ def __init__( ) def __len__(self) -> int: + """ + Return the number of autoregressive training samples available. + + Returns + ------- + int + Number of (init, target) pairs derivable from the datastore. + """ assert self.da_state is not None if self.datastore.is_forecast: # for now we simply create a single sample for each analysis time @@ -186,11 +203,11 @@ def _slice_state_time( ) -> 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`). + 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 ---------- @@ -277,8 +294,8 @@ def _slice_forcing_time( """ # 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. + # current autoregressive time step. The two `init_steps` can also be + # used as past forcings. init_steps = 2 da_list = [] @@ -466,13 +483,15 @@ def __getitem__( Returns ------- init_states : torch.Tensor - Initial states, shape (2, N_grid, d_features). + Initial states, shape ``(2, num_grid_nodes, num_state_vars)``. target_states : torch.Tensor - Target states, shape (ar_steps, N_grid, d_features). + Target states, shape ``(ar_steps, num_grid_nodes, num_state_vars)``. forcing : torch.Tensor - Windowed forcing, shape (ar_steps, N_grid, d_windowed_forcing). + Windowed forcing, shape ``(ar_steps, num_grid_nodes, F)`` where + ``F = num_forcing_vars * (num_past_forcing_steps`` + ``+ num_future_forcing_steps + 1)``. target_times : torch.Tensor - Times of the target steps, shape (ar_steps,). + Times of the target steps, shape ``(ar_steps,)``. """ n_samples = len(self) @@ -505,9 +524,9 @@ def __getitem__( 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) + # init_states: (2, num_grid_nodes, num_state_vars) + # target_states: (ar_steps, num_grid_nodes, num_state_vars) + # forcing: (ar_steps, num_grid_nodes, num_forcing_vars * window) # target_times: (ar_steps,) return init_states, target_states, forcing, target_times @@ -562,7 +581,7 @@ def create_dataarray_from_tensor( """ def _is_listlike(obj): - # match list, tuple, numpy array + """Return ``True`` for list/tuple/ndarray-like containers.""" return hasattr(obj, "__iter__") and not isinstance(obj, str) add_time_as_dim = False @@ -634,6 +653,30 @@ def __init__( num_workers: int = 16, eval_split: str = "test", ) -> None: + """ + 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``. + 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"``. + """ super().__init__() self._datastore = datastore self.num_past_forcing_steps = num_past_forcing_steps @@ -654,6 +697,16 @@ def __init__( self.multiprocessing_context = "spawn" def setup(self, stage: Optional[str] = None) -> None: + """ + 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. + """ if stage == "fit" or stage is None: self.train_dataset = WeatherDataset( datastore=self._datastore, diff --git a/pyproject.toml b/pyproject.toml index f288ea38..113329e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -146,6 +146,18 @@ allow-any-import-level = "neural_lam" [tool.pylint.SIMILARITIES] min-similarity-lines = 10 +[tool.interrogate] +fail-under = 100 +verbose = 1 +quiet = false +ignore-init-method = false +ignore-init-module = false +ignore-magic = false +ignore-private = false +ignore-nested-functions = false +ignore-property-decorators = false +exclude = ["tests", "docs", "build"] + [tool.pytest.ini_options] markers = [ "slow: marks tests as slow (deselected by default, run with -m slow)", diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py index f081bbf9..73e2f905 100644 --- a/tests/test_prediction_model_classes.py +++ b/tests/test_prediction_model_classes.py @@ -315,7 +315,7 @@ def test_forecaster_module_old_checkpoint(tmp_path): def test_graph_lam_no_static_features(): """GraphLAM (real GNN) should run a forward pass when the datastore has - no static features — verifying that the empty static tensor flows through + no static features - verifying that the empty static tensor flows through the graph encoder/processor/decoder without error.""" base_datastore = init_datastore_example("mdp")