Fix graph LAM GNN kwarg handling

CHANGELOG.md

@@ -43,6 +43,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

neural_lam/models/step_predictors/graph/graph_lam.py

@@ -36,6 +36,7 @@ def __init__(
         output_clamping_upper: dict[str, float] | None = None,
         g2m_gnn_type: str = "InteractionNet",
         m2g_gnn_type: str = "InteractionNet",
+        **_kwargs: object,
     ) -> None:
         """
         Initialize the GraphLAM model.

neural_lam/train_model.py

@@ -62,6 +62,12 @@ def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore):
         output_std=args.output_std,
         output_clamping_lower=config.training.output_clamping.lower,
         output_clamping_upper=config.training.output_clamping.upper,
+        g2m_gnn_type=getattr(args, "g2m_gnn_type", "InteractionNet"),
+        m2g_gnn_type=getattr(args, "m2g_gnn_type", "InteractionNet"),
+        mesh_up_gnn_type=getattr(args, "mesh_up_gnn_type", "InteractionNet"),
+        mesh_down_gnn_type=getattr(
+            args, "mesh_down_gnn_type", "InteractionNet"
+        ),
     )
     forecaster = ARForecaster(predictor, datastore)
     return ForecasterModule.load_from_checkpoint(

tests/test_train_model_warnings.py

@@ -1,11 +1,12 @@
 # Standard library
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 # Third-party
 import pytest
 
 # First-party
-from neural_lam.train_model import main
+from neural_lam.train_model import load_forecaster_module_from_checkpoint, main
 
 
 @pytest.mark.parametrize(
@@ -87,3 +88,57 @@ def capture_init(_self, **kwargs):
         "create_gif" in captured_kwargs
     ), "create_gif was not forwarded to ForecasterModule"
     assert captured_kwargs["create_gif"] is True
+
+
+def test_checkpoint_loader_restores_gnn_type_kwargs():
+    """Checkpoint reload must preserve custom GNN choices from saved args."""
+    args = SimpleNamespace(
+        model="hi_lam",
+        graph="hierarchical",
+        hidden_dim=4,
+        hidden_layers=1,
+        processor_layers=1,
+        mesh_aggr="sum",
+        num_past_forcing_steps=1,
+        num_future_forcing_steps=1,
+        output_std=False,
+        g2m_gnn_type="PropagationNet",
+        m2g_gnn_type="PropagationNet",
+        mesh_up_gnn_type="PropagationNet",
+        mesh_down_gnn_type="InteractionNet",
+    )
+    config = SimpleNamespace(
+        training=SimpleNamespace(
+            output_clamping=SimpleNamespace(lower={}, upper={})
+        )
+    )
+    datastore = MagicMock()
+    captured_kwargs = {}
+
+    class DummyPredictor:
+        def __init__(self, **kwargs):
+            captured_kwargs.update(kwargs)
+
+    loaded_module = MagicMock()
+
+    with (
+        patch(
+            "neural_lam.train_model.torch.load",
+            return_value={"hyper_parameters": {"args": args}},
+        ),
+        patch("neural_lam.train_model.MODELS", {"hi_lam": DummyPredictor}),
+        patch("neural_lam.train_model.ARForecaster"),
+        patch(
+            "neural_lam.train_model.ForecasterModule.load_from_checkpoint",
+            return_value=loaded_module,
+        ),
+    ):
+        result = load_forecaster_module_from_checkpoint(
+            "model.ckpt", config, datastore
+        )
+
+    assert result is loaded_module
+    assert captured_kwargs["g2m_gnn_type"] == "PropagationNet"
+    assert captured_kwargs["m2g_gnn_type"] == "PropagationNet"
+    assert captured_kwargs["mesh_up_gnn_type"] == "PropagationNet"
+    assert captured_kwargs["mesh_down_gnn_type"] == "InteractionNet"