Modernize linting stack: migrate to Ruff and streamline pre-commit hooks

.pre-commit-config.yaml

@@ -10,32 +10,37 @@ repos:
       - id: check-yaml
       - id: debug-statements
       - id: end-of-file-fixer
+      - id: no-commit-to-branch
+        args: [--branch, main]
       - id: trailing-whitespace
 
+  - repo: https://github.com/pre-commit/pygrep-hooks
+    rev: v1.10.0
+    hooks:
+      - id: python-check-blanket-noqa
+        description: Require specific codes on noqa comments
+
   - repo: https://github.com/codespell-project/codespell
     rev: v2.4.1
     hooks:
       - id: codespell
         description: Check for spelling errors
 
-  - repo: https://github.com/psf/black
-    rev: 25.11.0
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.4
     hooks:
-      - id: black
-        description: Format Python code
+      - id: ruff-check
+        args: [--fix]
+        description: Lint Python code with Ruff
+      - id: ruff-format
+        description: Format Python code with Ruff
 
-  - repo: https://github.com/PyCQA/isort
-    rev: 7.0.0
+  - repo: https://github.com/tox-dev/pyproject-fmt
+    rev: v2.21.0
     hooks:
-      - id: isort
-        description: Group and sort Python imports
+      - id: pyproject-fmt
+        description: Normalize pyproject.toml formatting
 
-  - repo: https://github.com/PyCQA/flake8
-    rev: 7.3.0
-    hooks:
-      - id: flake8
-        description: Check Python code for correctness, consistency and adherence to best practices
-        additional_dependencies: [Flake8-pyproject]
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.19.0
     hooks:

CHANGELOG.md

@@ -68,6 +68,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add comprehensive type hints to `neural_lam/create_graph.py` [\#618](https://github.com/mllam/neural-lam/pull/618) @GiGiKoneti
 
+- Modernize linting and formatting toolchain by migrating `black`/`isort`/`flake8` to Ruff and adding `pyproject-fmt`, `python-check-blanket-noqa`, and `no-commit-to-branch` pre-commit hooks [\#614](https://github.com/mllam/neural-lam/pull/614) @Shivampal157
+
 
 ## [v0.6.0](https://github.com/mllam/neural-lam/releases/tag/v0.6.0)
 

neural_lam/datastore/__init__.py

@@ -1,7 +1,7 @@
 # Local
-from .base import BaseDatastore  # noqa
-from .mdp import MDPDatastore  # noqa
-from .npyfilesmeps import NpyFilesDatastoreMEPS  # noqa
+from .base import BaseDatastore
+from .mdp import MDPDatastore
+from .npyfilesmeps import NpyFilesDatastoreMEPS
 
 DATASTORE_CLASSES = [
     MDPDatastore,

neural_lam/datastore/mdp.py

@@ -114,9 +114,9 @@ def __init__(self, config_path, n_boundary_points=30, reuse_existing=True):
             if dim_order is None:
                 dim_order = dim_order_
             else:
-                assert (
-                    dim_order == dim_order_
-                ), "all inputs must have the same dimension order"
+                assert dim_order == dim_order_, (
+                    "all inputs must have the same dimension order"
+                )
 
         self.spatial_coordinates = dim_order
 
@@ -489,9 +489,9 @@ def get_xy(self, category: str, stacked: bool) -> ndarray:
         da_xs = ds_category[xdim]
         da_ys = ds_category[ydim]
 
-        assert (
-            da_xs.ndim == da_ys.ndim == 1
-        ), f"{xdim} and {ydim} coordinates must be 1D"
+        assert da_xs.ndim == da_ys.ndim == 1, (
+            f"{xdim} and {ydim} coordinates must be 1D"
+        )
 
         da_x, da_y = xr.broadcast(da_xs, da_ys)
         da_xy = xr.concat([da_x, da_y], dim="grid_coord")

neural_lam/datastore/npyfilesmeps/__init__.py

@@ -1,2 +1,2 @@
 # Local
-from .store import NpyFilesDatastoreMEPS  # noqa
+from .store import NpyFilesDatastoreMEPS

neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py

@@ -223,9 +223,10 @@ def main(
         flux_squares.append(torch.mean(flux_batch**2).cpu())  # (,)
 
     if distributed and world_size > 1:
-        means_gathered, squares_gathered = [None] * world_size, [
-            None
-        ] * world_size
+        means_gathered, squares_gathered = (
+            [None] * world_size,
+            [None] * world_size,
+        )
         flux_means_gathered, flux_squares_gathered = (
             [None] * world_size,
             [None] * world_size,
@@ -322,9 +323,9 @@ def main(
         state_std = state_std.to(device)
 
     time_step_int, time_step_unit = get_integer_time(step_length)
-    assert (
-        time_step_unit == "hours"
-    ), "Only 'hours' time unit is supported by meps datastore."
+    assert time_step_unit == "hours", (
+        "Only 'hours' time unit is supported by meps datastore."
+    )
     used_subsample_len = (65 // time_step_int) * time_step_int
 
     diff_means, diff_squares = [], []
@@ -333,8 +334,9 @@ def main(
         loader_standard, disable=rank != 0
     ):
         if distributed:
-            init_batch, target_batch = init_batch.to(device), target_batch.to(
-                device
+            init_batch, target_batch = (
+                init_batch.to(device),
+                target_batch.to(device),
             )
         init_batch = (init_batch - state_mean) / state_std
         target_batch = (target_batch - state_mean) / state_std

neural_lam/datastore/plot_example.py

@@ -157,9 +157,9 @@ def _parse_dict(arg_str):
     )
     args = parser.parse_args()
 
-    assert (
-        args.datastore_config_path is not None
-    ), "Specify your datastore config with --datastore_config_path"
+    assert args.datastore_config_path is not None, (
+        "Specify your datastore config with --datastore_config_path"
+    )
 
     selection = dict(args.selection)
     index_selection = dict(args.index_selection)

neural_lam/gnn_layers.py

@@ -247,9 +247,9 @@ class SplitMLPs(nn.Module):
 
     def __init__(self, mlps, chunk_sizes):
         super().__init__()
-        assert len(mlps) == len(
-            chunk_sizes
-        ), "Number of MLPs must match the number of chunks"
+        assert len(mlps) == len(chunk_sizes), (
+            "Number of MLPs must match the number of chunks"
+        )
 
         self.mlps = nn.ModuleList(mlps)
         self.chunk_sizes = chunk_sizes

neural_lam/metrics.py

@@ -17,9 +17,9 @@ def get_metric(metric_name):
         Function implementing the requested metric.
     """
     metric_name_lower = metric_name.lower()
-    assert (
-        metric_name_lower in DEFINED_METRICS
-    ), f"Unknown metric: {metric_name}"
+    assert metric_name_lower in DEFINED_METRICS, (
+        f"Unknown metric: {metric_name}"
+    )
     return DEFINED_METRICS[metric_name_lower]
 
 

neural_lam/models/step_predictors/base.py

@@ -177,10 +177,10 @@ 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):
+            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 = []
@@ -195,11 +195,11 @@ def prepare_clamping_params(self, datastore: BaseDatastore):
 
         for feature_idx, feature in enumerate(state_feature_names):
             if feature in lower_lims and feature in upper_lims:
-                assert (
-                    lower_lims[feature] < upper_lims[feature]
-                ), f'Invalid clamping limits for feature "{feature}",\
+                assert lower_lims[feature] < upper_lims[feature], (
+                    f'Invalid clamping limits for feature "{feature}",\
                      lower: {lower_lims[feature]}, larger than\
                      upper: {upper_lims[feature]}'
+                )
                 sigmoid_lower_upper_idx.append(feature_idx)
                 sigmoid_lower_lims.append(
                     normalize_clamping_lim(lower_lims[feature], feature_idx)

neural_lam/models/step_predictors/graph/graph_lam.py

@@ -51,9 +51,9 @@ def __init__(
             m2g_gnn_type=m2g_gnn_type,
         )
 
-        assert (
-            not self.hierarchical
-        ), "GraphLAM does not use a hierarchical mesh graph"
+        assert not self.hierarchical, (
+            "GraphLAM does not use a hierarchical mesh graph"
+        )
 
         # grid_dim from data + static + batch_static
         mesh_dim = self.mesh_static_features.shape[1]

neural_lam/utils.py

@@ -297,12 +297,12 @@ def loads_file(fn: str) -> Any:
     m2g_features = m2g_features / longest_edge
 
     # Some checks for consistency
-    assert (
-        len(m2m_features) == n_levels
-    ), "Inconsistent number of levels in mesh"
-    assert (
-        len(mesh_static_features) == n_levels
-    ), "Inconsistent number of levels in mesh"
+    assert len(m2m_features) == n_levels, (
+        "Inconsistent number of levels in mesh"
+    )
+    assert len(mesh_static_features) == n_levels, (
+        "Inconsistent number of levels in mesh"
+    )
 
     if hierarchical:
         # Load up and down edges and features