[Bug] TypeError: GraphLAM.__init__() got an unexpected keyword argument 'mesh_up_gnn_type'

Description

When running model training for the default flat graph model (--model graph_lam), the execution crashes immediately at startup during the initialization of the step predictor.

This occurs because train_model.py unconditionally passes hierarchical-only parameters (mesh_up_gnn_type and mesh_down_gnn_type) to the predictor constructor, which the GraphLAM class does not accept in its signature.

In addition, there is a secondary bug in load_forecaster_module_from_checkpoint: when reloading a model from a checkpoint for evaluation, custom GNN parameters are not reconstructed, causing the model to default back to InteractionNet even if custom GNN types were configured during training.

Traceback

Seed set to 42
2026-06-27 07:02:10.308 | WARNING  | neural_lam.datastore.mdp:__init__:81 - Config file has been modified since zarr was created. The old zarr archive (in tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.zarr) will be used.
2026-06-27 07:02:10.345 | INFO     | neural_lam.utils:log_on_rank_zero:587 - The loaded datastore contains the following features:
2026-06-27 07:02:10.345 | INFO     | neural_lam.utils:log_on_rank_zero:587 -  state   : u100m v100m r2m t2m
2026-06-27 07:02:10.345 | INFO     | neural_lam.utils:log_on_rank_zero:587 -  forcing : swavr0m
2026-06-27 07:02:10.345 | INFO     | neural_lam.utils:log_on_rank_zero:587 -  static  : lsm orography
2026-06-27 07:02:10.346 | INFO     | neural_lam.utils:log_on_rank_zero:587 - With the following splits (over time):
2026-06-27 07:02:10.354 | INFO     | neural_lam.utils:log_on_rank_zero:587 -  train   : 2022-04-01T00:00 to 2022-04-04T00:00
2026-06-27 07:02:10.359 | INFO     | neural_lam.utils:log_on_rank_zero:587 -  val     : 2022-04-04T00:00 to 2022-04-07T00:00
2026-06-27 07:02:10.364 | INFO     | neural_lam.utils:log_on_rank_zero:587 -  test    : 2022-04-07T00:00 to 2022-04-10T00:00
2026-06-27 07:02:10.364 | ERROR    | __main__:<module>:555 - An error has been caught in function '<module>', process 'MainProcess' (4507), thread 'MainThread' (8412765824):
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/mynimbus/neural-lam/neural_lam/train_model.py", line 555, in <module>
    main()
  File "/Users/mynimbus/neural-lam/neural_lam/train_model.py", line 443, in main
    predictor = predictor_class(
TypeError: GraphLAM.__init__() got an unexpected keyword argument 'mesh_up_gnn_type'

Minimal Reproducible Example

Run the following training command targeting the graph_lam architecture:

uv run python -m neural_lam.train_model \
  --config_path tests/datastore_examples/mdp/danra_100m_winds/config.yaml \
  --model graph_lam \
  --graph dummy_graph \
  --epochs 1 \
  --batch_size 2

Root Cause Analysis

1. In neural_lam/train_model.py (lines 443-459), flat and hierarchical parameters are merged into a single hardcoded instantiation call:

   predictor = predictor_class(
       datastore=datastore,
       # ...
       g2m_gnn_type=args.g2m_gnn_type,
       m2g_gnn_type=args.m2g_gnn_type,
       mesh_up_gnn_type=args.mesh_up_gnn_type,
       mesh_down_gnn_type=args.mesh_down_gnn_type,
   )
   

   If predictor_class is GraphLAM, it lacks mesh_up_gnn_type and mesh_down_gnn_type in its signature, causing a TypeError.
2. In load_forecaster_module_from_checkpoint (lines 41-72), GNN type parameters are not passed to the constructor at all, meaning the reloaded predictor is reconstructed using default values instead of the GNN settings stored in the checkpoint's hyperparameters.

Proposed Fix

Use python's inspect module to dynamically filter constructor parameters before instantiating predictor_class in both main() and load_forecaster_module_from_checkpoint().

1. In main():

   import inspect
   
   predictor_kwargs = {
       "datastore": datastore,
       "graph_name": args.graph,
       "hidden_dim": args.hidden_dim,
       "hidden_layers": args.hidden_layers,
       "processor_layers": args.processor_layers,
       "mesh_aggr": args.mesh_aggr,
       "num_past_forcing_steps": args.num_past_forcing_steps,
       "num_future_forcing_steps": args.num_future_forcing_steps,
       "output_std": args.output_std,
       "output_clamping_lower": config.training.output_clamping.lower,
       "output_clamping_upper": config.training.output_clamping.upper,
       "g2m_gnn_type": args.g2m_gnn_type,
       "m2g_gnn_type": args.m2g_gnn_type,
       "mesh_up_gnn_type": args.mesh_up_gnn_type,
       "mesh_down_gnn_type": args.mesh_down_gnn_type,
   }
   
   # Filter kwargs to match constructor signature
   sig = inspect.signature(predictor_class)
   filtered_kwargs = {
       k: v for k, v in predictor_kwargs.items()
       if k in sig.parameters
   }
   predictor = predictor_class(**filtered_kwargs)
   

2. In load_forecaster_module_from_checkpoint(): Apply the same filtering pattern and extract GNN settings from args using getattr to maintain backward-compatibility with older checkpoint files.

cc: @joeloskarsson @sadamov