[Feature Request] Separate processor_layers for fine- and coarse-grained layers in hi_lam

[jamesYu84]

Hi, thank you for the excellent work on neural-lam!

I am currently experimenting with the hi_lam model and trying to reproduce the WaveGraph paper settings on a Mediterranean region dataset. I noticed that the current implementation uses a single parameter --processor_layers to control the number of message-passing steps for all layers (both fine- and coarse-grained graphs).

According to the WaveGraph paper:

Fine-grained graph: uses N_v = 2 message-passing steps.

Coarse-grained graph: uses N_u = 8 message-passing steps.

However, --processor_layers (default=4 in hi_lam) applies the same number of steps to both the fine and coarse branches, which makes it impossible to replicate the paper's exact architecture.

Questions:

Is this a known design choice or limitation?

Would it be possible to add separate parameters, e.g., --fine_processor_layers and --coarse_processor_layers, to better align with the original WaveGraph architecture?

In the meantime, is there a workaround to achieve this separation (e.g., by modifying the config file or model definition)?

I'd be happy to test and provide feedback if this feature is added.

Environment:

neural-lam version: 59c3e77

Python: 3.10

CUDA: 12.4

GPU: RTX 3090 24GB

Thank you for your time!

[joeloskarsson]

Hi, I am not familiar with WaveGraph, but happy to try answering the best I can. Not sure if relevant, but the SeaCast model for the Mediterranean is based on this codebase, so maybe you could also get some inspiration there? https://www.nature.com/articles/s41598-025-31177-w

Note that the number of processor layers in hi_lam means the number of full up-and-down sweeps we do through the hierarchy (see e.g. figure 1 and appendix B.2 of https://arxiv.org/pdf/2309.17370.pdf). So this is somewhat different than just the number of message-passing steps at different coarseness levels. But perhaps this would achieve the same thing you are after? The steps up-and-down the hierarchy are intertwined with intra-level message-passing, e.g.

neural-lam/neural_lam/models/step_predictors/graph/hi_lam.py

Lines 231 to 234 in 59c3e77

	# Run same level processing on level l
	mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn(
	new_node_rep, new_node_rep, same_edge_rep
	)

, so perhaps the generalization you would like to do is to use multiple GNN layers and message passing steps for these parts?

You could likely achieve this by changing the GNN instantiation here

neural-lam/neural_lam/models/step_predictors/graph/hi_lam.py

Lines 89 to 103 in 59c3e77

	# Make down GNNs, both for down edges and same level
	self.mesh_down_gnns = nn.ModuleList(
	[self.make_down_gnns() for _ in range(processor_layers)]
	) # Nested lists (proc_steps, num_levels-1)
	self.mesh_down_same_gnns = nn.ModuleList(
	[self.make_same_gnns() for _ in range(processor_layers)]
	) # Nested lists (proc_steps, num_levels)
	# Make up GNNs, both for up edges and same level
	self.mesh_up_gnns = nn.ModuleList(
	[self.make_up_gnns() for _ in range(processor_layers)]
	) # Nested lists (proc_steps, num_levels-1)
	self.mesh_up_same_gnns = nn.ModuleList(
	[self.make_same_gnns() for _ in range(processor_layers)]
	) # Nested lists (proc_steps, num_levels)

, but I think you would also have to make some small changes to the actual forward-pass of the model.

Hmm, it might actually even be possible to get to the multi-message-passing on each level by just replacing the one InteractionNetwork instantiated per level in

neural-lam/neural_lam/models/step_predictors/graph/hi_lam.py

Lines 105 to 123 in 59c3e77

	def make_same_gnns(self) -> nn.ModuleList:
	"""
	Make intra-level GNNs.
	Returns
	-------
	nn.ModuleList
	List of GNNs for each level.
	"""
	return nn.ModuleList(
	[
	InteractionNet(
	edge_index,
	self.hidden_dim,
	hidden_layers=self.hidden_layers,
	)
	for edge_index in self.m2m_edge_index
	]
	)

with a stack of such networks, e.g. using

neural-lam/neural_lam/utils/networks.py

Line 43 in 59c3e77

def make_gnn_seq(

, and then not have to change the forward pass at all. But this requires testing.