Enforce 100% type-hint coverage and align typing to PEP 585/604

.github/workflows/install-and-test.yml

@@ -67,6 +67,11 @@ jobs:
         run: |
           python -c "import torch; assert torch.__version__.endswith('+cu130')"
 
+      - name: Type check with ty
+        if: matrix.device == 'cpu'
+        run: |
+          uv run ty check neural_lam
+
       - name: Load cache data
         uses: actions/cache/restore@v4
         with:

.github/workflows/pre-commit.yml

@@ -21,3 +21,5 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - uses: pre-commit/action@v3.0.1
+      env:
+        SKIP: ty

.pre-commit-config.yaml

@@ -43,12 +43,16 @@ repos:
         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
+  - repo: local
     hooks:
-      - id: mypy
-        additional_dependencies: [types-PyYAML, types-Pillow, types-tqdm]
+      - id: ty
+        name: ty
         description: Check for type errors
+        entry: uv run ty check neural_lam
+        language: system
+        pass_filenames: false
+        files: ^neural_lam/
+        types: [python]
 
   - repo: local
     hooks:

CHANGELOG.md

@@ -48,6 +48,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Allow `graph_lam` training and checkpoint reloads to accept the full set of
+  GNN type CLI options without passing hierarchical-only options to unsupported
+  constructors ([#686](https://github.com/mllam/neural-lam/issues/686)).
+
 - Fix `RuntimeError` in `HiLAMParallel` forward pass on hierarchical graphs by offsetting edge indices into the global mesh node index space ([#679](https://github.com/mllam/neural-lam/issues/679))
 
 - Fix `IndexError` in HiLAM forward pass by offsetting grid nodes in `zero_index_g2m`/`zero_index_m2g` by the total mesh-node count across all levels ([#642](https://github.com/mllam/neural-lam/issues/642)) @Sir-Sloth-The-Lazy
@@ -81,6 +85,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Maintenance
 
+- Add 100% type-hint coverage across `neural_lam/`, align all type annotations with PEP 585 and PEP 604, and adopt `ty` (astral-sh) for type checking [\#673](https://github.com/mllam/neural-lam/pull/673) @GiGiKoneti
+
 - Split the monolithic `neural_lam/utils.py` into a `neural_lam/utils/` package with one module per concern (`buffer_list`, `graph`, `networks`, `plot`, `logging`, `tensor`, `time`); `utils/__init__.py` re-exports the full public API so existing imports are unaffected. Pure code movement, no behavioural change. [\#682](https://github.com/mllam/neural-lam/pull/682) @Sir-Sloth-The-Lazy
 
 - Add comprehensive type hints to GraphLAM in `neural_lam/models/step_predictors/graph/graph_lam.py` [\#669](https://github.com/mllam/neural-lam/pull/669) @GiGiKoneti

neural_lam/config.py

@@ -3,7 +3,7 @@
 # Standard library
 import dataclasses
 from pathlib import Path
-from typing import Dict, Union
+from typing import cast
 
 # Third-party
 import dataclass_wizard
@@ -35,7 +35,7 @@ class DatastoreSelection:
     kind: str
     config_path: str
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         """
         Validate that the selected datastore kind is implemented.
 
@@ -56,11 +56,11 @@ class ManualStateFeatureWeighting:
 
     Attributes
     ----------
-    weights : Dict[str, float]
+    weights : dict[str, float]
         Manual weights for the state features.
     """
 
-    weights: Dict[str, float]
+    weights: dict[str, float]
 
 
 @dataclasses.dataclass
@@ -80,14 +80,14 @@ class OutputClamping:
 
     Attributes
     ----------
-    lower : Dict[str, float]
+    lower : dict[str, float]
         The minimum value to clamp each output feature to.
-    upper : Dict[str, float]
+    upper : dict[str, float]
         The maximum value to clamp each output feature to.
     """
 
-    lower: Dict[str, float] = dataclasses.field(default_factory=dict)
-    upper: Dict[str, float] = dataclasses.field(default_factory=dict)
+    lower: dict[str, float] = dataclasses.field(default_factory=dict)
+    upper: dict[str, float] = dataclasses.field(default_factory=dict)
 
 
 @dataclasses.dataclass
@@ -97,8 +97,8 @@ class TrainingConfig:
 
     Attributes
     ----------
-    state_feature_weighting : Union[ManualStateFeatureWeighting,
-                                    UniformFeatureWeighting]
+    state_feature_weighting :
+        ManualStateFeatureWeighting | UniformFeatureWeighting
         The method to use for weighting the state features in the loss
         function. Defaults to uniform weighting (`UniformFeatureWeighting`, i.e.
         all features are weighted equally).
@@ -107,9 +107,9 @@ class TrainingConfig:
         Defaults to an empty ``OutputClamping`` (no clamping).
     """
 
-    state_feature_weighting: Union[
-        ManualStateFeatureWeighting, UniformFeatureWeighting
-    ] = dataclasses.field(default_factory=UniformFeatureWeighting)
+    state_feature_weighting: (
+        ManualStateFeatureWeighting | UniformFeatureWeighting
+    ) = dataclasses.field(default_factory=UniformFeatureWeighting)
 
     output_clamping: OutputClamping = dataclasses.field(
         default_factory=OutputClamping
@@ -174,7 +174,7 @@ class InvalidConfigError(Exception):
 
 def load_config_and_datastore(
     config_path: str,
-) -> tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]:
+) -> tuple[NeuralLAMConfig, MDPDatastore | NpyFilesDatastoreMEPS]:
     """
     Load the neural-lam configuration and the datastore specified in the
     configuration.
@@ -186,7 +186,7 @@ def load_config_and_datastore(
 
     Returns
     -------
-    tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]
+    tuple[NeuralLAMConfig, MDPDatastore | NpyFilesDatastoreMEPS]
         The Neural-LAM configuration and the loaded datastore.
     """
     try:
@@ -204,4 +204,4 @@ def load_config_and_datastore(
         datastore_kind=config.datastore.kind, config_path=datastore_config_path
     )
 
-    return config, datastore
+    return config, cast(MDPDatastore | NpyFilesDatastoreMEPS, datastore)

neural_lam/create_graph.py

@@ -3,7 +3,6 @@
 # Standard library
 import os
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
-from typing import Optional
 
 # Third-party
 import matplotlib
@@ -27,7 +26,7 @@
 
 
 def plot_graph(
-    graph: pyg.data.Data, title: Optional[str] = None
+    graph: pyg.data.Data, title: str | None = None
 ) -> tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]:
     """
     Render a PyTorch Geometric graph using stored node coordinates.
@@ -357,10 +356,10 @@ def prepend_node_index(graph: networkx.Graph, new_index: int) -> networkx.Graph:
 def create_graph(
     graph_dir_path: str,
     xy: np.ndarray,
-    n_max_levels: Optional[int] = None,
-    hierarchical: Optional[bool] = False,
-    create_plot: Optional[bool] = False,
-):
+    n_max_levels: int | None = None,
+    hierarchical: bool = False,
+    create_plot: bool = False,
+) -> None:
     """
     Create graph components from `xy` grid coordinates and store in
     `graph_dir_path`.
@@ -653,8 +652,10 @@ def create_graph(
                 .reshape((n, n, 2))[1::nx, 1::nx, :]
                 .reshape(int(n / nx) ** 2, 2)
             )
-            ij = [tuple(x) for x in ij]
-            G[lev] = networkx.relabel_nodes(G[lev], dict(zip(G[lev].nodes, ij)))
+            ij_tuples = [tuple(x) for x in ij]
+            G[lev] = networkx.relabel_nodes(
+                G[lev], dict(zip(G[lev].nodes, ij_tuples))
+            )
             G_tot = networkx.compose(G_tot, G[lev])
 
         # Relabel mesh nodes to start with 0
@@ -864,10 +865,10 @@ def create_graph(
 def create_graph_from_datastore(
     datastore: BaseRegularGridDatastore,
     output_root_path: str,
-    n_max_levels: Optional[int] = None,
+    n_max_levels: int | None = None,
     hierarchical: bool = False,
     create_plot: bool = False,
-):
+) -> None:
     """
     Generate graph components for ``datastore`` and persist them on disk.
 
@@ -900,7 +901,7 @@ def create_graph_from_datastore(
     )
 
 
-def cli(input_args: Optional[list[str]] = None) -> None:
+def cli(input_args: list[str] | None = None) -> None:
     """
     Parse CLI arguments and call :func:`create_graph_from_datastore`.
 

neural_lam/custom_loggers.py

@@ -2,7 +2,6 @@
 
 # Standard library
 import os
-from typing import Optional
 
 # Third-party
 import matplotlib.pyplot as plt
@@ -74,7 +73,7 @@ def log_image(
         self,
         key: str,
         images: list[plt.Figure],
-        step: Optional[int] = None,
+        step: int | None = None,
     ) -> None:
         """
         Log one or more matplotlib figures as images to MLFlow.

neural_lam/datastore/__init__.py

@@ -1,5 +1,8 @@
 """Datastore backends for loading and serving weather model data."""
 
+# Standard library
+from pathlib import Path
+
 # Local
 from .base import BaseDatastore  # noqa
 from .mdp import MDPDatastore  # noqa
@@ -11,12 +14,13 @@
 ]
 
 DATASTORES = {
-    datastore.SHORT_NAME: datastore  # type: ignore
-    for datastore in DATASTORE_CLASSES
+    datastore.SHORT_NAME: datastore for datastore in DATASTORE_CLASSES
 }
 
 
-def init_datastore(datastore_kind, config_path):
+def init_datastore(
+    datastore_kind: str, config_path: str | Path
+) -> BaseDatastore:
     """
     Instantiate a datastore based on its short-name identifier.
 

neural_lam/datastore/base.py

@@ -111,7 +111,7 @@ def get_vars_units(self, category: str) -> list[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The units of the variables.
 
         """
@@ -127,7 +127,7 @@ def get_vars_names(self, category: str) -> list[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The names of the variables.
 
         """
@@ -143,7 +143,7 @@ def get_vars_long_names(self, category: str) -> list[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The long names of the variables.
 
         """
@@ -333,7 +333,7 @@ def get_xy_extent(self, category: str) -> list[float]:
 
         Returns
         -------
-        List[float]
+        list[float]
             The extent of the x, y coordinates.
 
         """
@@ -423,7 +423,7 @@ def expected_dim_order(
             The category of the dataset (state/forcing/static).
         Returns
         -------
-        List[str]
+        list[str]
             The expected dimension order for the dataarray or dataset.
 
         """
@@ -496,7 +496,7 @@ class BaseRegularGridDatastore(BaseDatastore):
     `stack_grid_coords` and `unstack_grid_coords` respectively).
     """
 
-    spatial_coordinates = ("x", "y")
+    spatial_coordinates: tuple[str, str] = ("x", "y")
 
     @cached_property
     @abc.abstractmethod
@@ -562,7 +562,7 @@ def unstack_grid_coords(
         da_or_ds_unstacked = da_or_ds.unstack("grid_index")
 
         # Ensure that the x, y dimensions are in the correct order
-        dims = da_or_ds_unstacked.dims
+        dims = list(da_or_ds_unstacked.dims)
         xy_dim_order = [d for d in dims if d in self.spatial_coordinates]
 
         if xy_dim_order != self.spatial_coordinates:
@@ -615,7 +615,7 @@ def stack_grid_coords(
         # dimension named in the format `{category}_feature`
         category = None
         for dim in da_or_ds_stacked.dims:
-            if dim.endswith("_feature"):
+            if isinstance(dim, str) and dim.endswith("_feature"):
                 if category is not None:
                     raise ValueError(
                         "Multiple dimensions ending with '_feature' found in "