From 203c54b19ff57a319bd118869c21fbc08de47c60 Mon Sep 17 00:00:00 2001 From: Junaid Shah Date: Fri, 27 Feb 2026 22:15:04 +0530 Subject: [PATCH 1/2] docs: improve __init__ docstrings and parameter descriptions --- neural_lam/models/base_graph_model.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/neural_lam/models/base_graph_model.py b/neural_lam/models/base_graph_model.py index 35b1ab126..eb7b7a208 100644 --- a/neural_lam/models/base_graph_model.py +++ b/neural_lam/models/base_graph_model.py @@ -16,6 +16,21 @@ class BaseGraphModel(ARModel): """ def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore): + """ + Initializes the BaseGraphModel for graph-based weather forecasting. + + This constructor loads the graph structure, registers static features as buffers, + initializes MLP embedders for grid and mesh nodes, and sets up the + Encoder-Processor-Decoder GNN architecture. + + Args: + args (Namespace): Command-line arguments containing hyperparameters like + hidden_dim and hidden_layers. + config (NeuralLAMConfig): Configuration object containing training and + model settings. + datastore (BaseDatastore): Datastore instance to handle data loading + and graph paths. + """ super().__init__(args, config=config, datastore=datastore) # Load graph with static features From aee63b96162e9377efabd8fd81b2a6241a49127c Mon Sep 17 00:00:00 2001 From: Junaid Shah Date: Fri, 27 Feb 2026 22:59:46 +0530 Subject: [PATCH 2/2] docs: add hierarchical model docstrings for project #2 --- neural_lam/models/hi_lam.py | 98 +++++++++++++++++++++++++++---------- 1 file changed, 73 insertions(+), 25 deletions(-) diff --git a/neural_lam/models/hi_lam.py b/neural_lam/models/hi_lam.py index c340c95da..f5727a27b 100644 --- a/neural_lam/models/hi_lam.py +++ b/neural_lam/models/hi_lam.py @@ -12,10 +12,22 @@ class HiLAM(BaseHiGraphModel): """ Hierarchical graph model with message passing that goes sequentially down and up the hierarchy during processing. - The Hi-LAM model from Oskarsson et al. (2023) + + The Hi-LAM model from Oskarsson et al. (2023) implements a multi-scale + approach to weather forecasting using hierarchical Graph Neural Networks. """ def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore): + """ + Initializes the Hi-LAM model with hierarchical processing layers. + + Args: + args (Namespace): Command-line arguments containing model hyperparameters + like hidden_dim and processor_layers. + config (NeuralLAMConfig): Configuration object for the Neural-LAM model. + datastore (BaseDatastore): Datastore object providing access to hierarchical + graph structures and edge indices. + """ super().__init__(args, config=config, datastore=datastore) # Make down GNNs, both for down edges and same level @@ -36,7 +48,13 @@ def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore): def make_same_gnns(self, args): """ - Make intra-level GNNs. + Creates GNN layers for intra-level (same level) message passing. + + Args: + args (Namespace): Model arguments specifying dimensions and layers. + + Returns: + nn.ModuleList: A list of InteractionNet layers for each hierarchy level. """ return nn.ModuleList( [ @@ -51,7 +69,13 @@ def make_same_gnns(self, args): def make_up_gnns(self, args): """ - Make GNNs for processing steps up through the hierarchy. + Creates GNN layers for processing steps upward through the hierarchy. + + Args: + args (Namespace): Model arguments specifying dimensions and layers. + + Returns: + nn.ModuleList: A list of InteractionNet layers for upward edges. """ return nn.ModuleList( [ @@ -66,7 +90,13 @@ def make_up_gnns(self, args): def make_down_gnns(self, args): """ - Make GNNs for processing steps down through the hierarchy. + Creates GNN layers for processing steps downward through the hierarchy. + + Args: + args (Namespace): Model arguments specifying dimensions and layers. + + Returns: + nn.ModuleList: A list of InteractionNet layers for downward edges. """ return nn.ModuleList( [ @@ -88,8 +118,20 @@ def mesh_down_step( same_gnns, ): """ - Run down-part of vertical processing, sequentially alternating between - processing using down edges and same-level edges. + Executes the downward part of hierarchical processing. + + Alternates between message passing across downward edges (inter-level) + and same-level edges (intra-level). + + Args: + mesh_rep_levels (list): List of node representations at each level. + mesh_same_rep (list): List of same-level edge representations. + mesh_down_rep (list): List of downward edge representations. + down_gnns (nn.ModuleList): GNN layers for downward processing. + same_gnns (nn.ModuleList): GNN layers for same-level processing. + + Returns: + tuple: Updated (mesh_rep_levels, mesh_same_rep, mesh_down_rep). """ # Run same level processing on level L mesh_rep_levels[-1], mesh_same_rep[-1] = same_gnns[-1]( @@ -119,7 +161,6 @@ def mesh_down_step( mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn( new_node_rep, new_node_rep, same_edge_rep ) - # (B, N_mesh[l], d_h) and (B, M_same[l], d_h) return mesh_rep_levels, mesh_same_rep, mesh_down_rep @@ -127,8 +168,17 @@ def mesh_up_step( self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, up_gnns, same_gnns ): """ - Run up-part of vertical processing, sequentially alternating between - processing using up edges and same-level edges. + Executes the upward part of hierarchical processing. + + Args: + mesh_rep_levels (list): List of node representations. + mesh_same_rep (list): List of same-level edge representations. + mesh_up_rep (list): List of upward edge representations. + up_gnns (nn.ModuleList): GNN layers for upward processing. + same_gnns (nn.ModuleList): GNN layers for same-level processing. + + Returns: + tuple: Updated (mesh_rep_levels, mesh_same_rep, mesh_up_rep). """ # Run same level processing on level 0 @@ -152,13 +202,11 @@ def mesh_up_step( new_node_rep, mesh_up_rep[level_l - 1] = up_gnn( send_node_rep, rec_node_rep, up_edge_rep ) - # (B, N_mesh[l], d_h) and (B, M_up[l-1], d_h) # 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 ) - # (B, N_mesh[l], d_h) and (B, M_same[l], d_h) return mesh_rep_levels, mesh_same_rep, mesh_up_rep @@ -166,17 +214,19 @@ def hi_processor_step( self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep ): """ - Internal processor step of hierarchical graph models. - Between mesh init and read out. - - Each input is list with representations, each with shape - - mesh_rep_levels: (B, N_mesh[l], d_h) - mesh_same_rep: (B, M_same[l], d_h) - mesh_up_rep: (B, M_up[l -> l+1], d_h) - mesh_down_rep: (B, M_down[l <- l+1], d_h) - - Returns same lists + Main internal processor step for the hierarchical model. + + This step coordinates the sequence of downward and upward message + passing across all hierarchical levels. + + Args: + mesh_rep_levels (list): Tensors of shape (B, N_mesh[l], d_h) + mesh_same_rep (list): Tensors of shape (B, M_same[l], d_h) + mesh_up_rep (list): Tensors of shape (B, M_up[l -> l+1], d_h) + mesh_down_rep (list): Tensors of shape (B, M_down[l <- l+1], d_h) + + Returns: + tuple: Updated representation lists for levels, same, up, and down edges. """ for down_gnns, down_same_gnns, up_gnns, up_same_gnns in zip( self.mesh_down_gnns, @@ -202,6 +252,4 @@ def hi_processor_step( up_same_gnns, ) - # NOTE: We return all, even though only down edges really are used - # later - return mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep + return mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep \ No newline at end of file