Adds PropagationNet GNN layer and makes it optionally usable in existing deterministic models

[Sir-Sloth-The-Lazy]

Describe your changes

Adds PropagationNet GNN layer and makes it optionally usable in existing deterministic models, as outlined in #62.

It is integrated into the existing model hierarchy from #208 and can be enabled via the vertical_propnets flag.

Depends on #208.
For changes on top of #208 only, see:
Sir-Sloth-The-Lazy/neural-lam@refactor/model-class-hierarchy-issue-49...refactor/batch-fold-ensemble-prep

Issue Link

Contributes to #62

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug
  • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
• author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
• Once the PR is ready to be merged, squash commits and merge the PR.

[Sir-Sloth-The-Lazy]

@joeloskarsson @sadamov @observingClouds please have a look , if this qualifies as the next step in ensemble prep 😄 . Would be grateful for your feedback !

[Debadri-das]

@joeloskarsson @sadamov would request your further review on this PR!

[joeloskarsson]

Thanks for looking at this, and sorry for being slow with reviews 😅 Shared some first thoughts here on how I think we can best integrate this. Happy to hear input also from others, as there are some non-trivial design choices around this (e.g. how to choose the GNN type for each sub-graph).

[Sir-Sloth-The-Lazy]

Refactor: replace vertical_propnets bool with fine-grained GNN type strings

Overview

Replaced the single --vertical_propnets boolean flag with four per-direction string parameters, enabling independent GNN selection per message-passing direction.

---

Changes

• Per-direction GNN type parameters
  Replaced --vertical_propnets with four string args: --g2m_gnn_type, --m2g_gnn_type, --mesh_up_gnn_type, --mesh_down_gnn_type. Each defaults to "InteractionNet" and can be set to "PropagationNet".

• GNN_TYPES registry + get_gnn_class() lookup (interaction_net.py)
  Added a registry to resolve GNN type strings to their corresponding classes, making it straightforward to register new GNN types in the future.

• PropagationNet forward() deduplication
  Eliminated the duplicated forward() method by extracting a node_residual_target() hook:

  • InteractionNet.node_residual_target() → returns the receiver representation (rec_rep)
  • PropagationNet.node_residual_target() → returns the aggregated edge messages (edge_rep_aggr)
  • Shared forward() now delegates to this hook.

• CLI wiring + propagation through model hierarchy
  Threaded the per-direction types through train_model.py, BaseGraphModel, BaseHiGraphModel, HiLAM, HiLAMParallel, and GraphLAM.

• Integration tests updated
  Tests now exercise the string-based API, including per-direction overrides and invalid-type error handling.

---

Files Changed (8)

File	Nature of change
interaction_net.py	Added GNN_TYPES registry, get_gnn_class(), node_residual_target() hook; unified forward()
base_graph_model.py	Threaded per-direction GNN type args
base_hi_graph_model.py	Threaded per-direction GNN type args
graph_lam.py	Wired new GNN type params
hi_lam.py	Wired new GNN type params
hi_lam_parallel.py	Wired new GNN type params
train_model.py	Replaced --vertical_propnets with four --*_gnn_type CLI args
test_propagation_net.py	Updated tests for string-based API and error handling

[Sir-Sloth-The-Lazy]

Proposal: GNN type selection per sub-graph

Context

With the current commit (3472387), four per-direction CLI args cover the encoder/decoder
and inter-level GNNs. However, three sub-graphs remain hardcoded to InteractionNet:

Sub-graph	Location	Currently configurable?
m2m same-level	HiLAM.make_same_gnns	No
m2m processor	GraphLAM.__init__	No
processor (all edges)	HiLAMParallel.__init__	No

Note on m2m same-level edges: The original --vertical_propnets flag deliberately
excluded same-level m2m processing, nodes at the same mesh level are peers with no
clear sender→receiver directionality, so PropagationNet (designed for directional
vertical message passing) may not be semantically meaningful here. Any global default
that switches everything would affect m2m same-level too, which is worth keeping in mind.

---

Option A: Single --gnn_type default + per-direction overrides

Add one global --gnn_type arg (default "InteractionNet"). Per-direction args
(--g2m_gnn_type, --m2m_gnn_type, etc.) default to None and override the global
when set. Currently-hardcoded sub-graphs in HiLAM and GraphLAM fall back to the
global default.

Resolution logic:

g2m_gnn_type = g2m_gnn_type or gnn_type  # per-direction wins if set

Pros:

• Clean UX for the common case — --gnn_type PropagationNet switches everything in one flag
• Fine-grained control still available via per-direction overrides
• Small diff on top of the current implementation
• Closes the hardcoded gaps (except HiLAMParallel) without adding more args

Cons:

• None-vs-default fallback logic adds a layer of indirection to understand
• A blanket --gnn_type PropagationNet would also switch m2m same-level edges, which may
  not be desirable (mitigated by --m2m_gnn_type InteractionNet override, but the user
  has to know to do this)
• Still CLI-arg-driven, which will eventually need migration if the codebase moves fully
  to config files

---

Option B: YAML config-driven (via NeuralLAMConfig)

Move GNN type selection into the existing NeuralLAMConfig YAML config. The codebase
already uses dataclass_wizard.YAMLWizard with Union-based type selection for training
config (loss weighting, output clamping). Model architecture config would follow the same
pattern:

model:
  hidden_dim: 64
  processor_layers: 4
  gnn_types:
    g2m: PropagationNet
    m2g: InteractionNet
    mesh_up: PropagationNet
    mesh_down: InteractionNet
    m2m: InteractionNet

Pros:

• Consistent with the direction the codebase is already moving (NeuralLAMConfig handles
  datastore, loss weighting, output clamping via YAML)
• Scales cleanly to any number of sub-graphs or future model config without CLI bloat
• Config files are versioned, shareable, and reproducible — easier to track experiment setups
• No fallback/override logic needed; each sub-graph has an explicit entry
• Makes the m2m same-level choice explicit rather than inheriting from a global default

Cons:

• Larger refactor — needs new dataclasses, migration of existing CLI model args
• Scope creep risk if done within this PR (--hidden_dim, --processor_layers would
  ideally move at the same time)
• Diverges from the current CLI-driven workflow that other model args still use

---

Recommendation

Option A now, Option B later.

Option A is a small, self-contained change that fits the existing CLI pattern and closes
the hardcoded gaps in HiLAM and GraphLAM immediately. Option B is the right
long-term direction but is better done as a dedicated refactor that moves all model
hyperparameters into NeuralLAMConfig at once, rather than partially migrating just
GNN types.

HiLAMParallel remains out of scope for both options , its processor fuses all
edge types (m2m + up + down) into a single InteractionNet with
edge_chunk_sizes/aggr_chunk_sizes, so per-direction GNN type selection is not
applicable without rethinking its architecture.

[joeloskarsson]

Looking really nice now! Previous comments addressed, I just added some small suggestions here that is really just polishing.

The outstanding question is indeed what to do with the remaining GNN layers, if we should make them all configurable. In my opinion there is not that much point in making the GNNs for m2m edges propagation-networks, so for most practical use cases the present options are sufficient. But it is a bit strange that some of the GNNs can be changed but not all. I think I am leaning towards trying to merge this as is atm, and then leaving a further generalization (e.g. through a model-config file) as a future change. But this requires some discussion I think. I will propose this for the next release version, and also see what others at the dev meeting think about the configurability question.

[Sir-Sloth-The-Lazy]

Hi @joeloskarsson made the changes , agreed with your verdict and hope you found the proposal a little helpful ! ;)

[joeloskarsson]

From the dev-meeting there seems to be agreement to stick to the current design of cli-arguments for choosing some gnn layers, and leaving the broader config to later.

[Sir-Sloth-The-Lazy]

Understood , I would love to work on that when the time comes 😄

[joeloskarsson]

The implementation here is all good now, so approving this. This should however go in after #208, so need to wait after possible conflict resolution with that (although might not be much as it is built on top of that).

The one thing missing is a changelog entry :)