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

neural_lam/interaction_net.py

@@ -131,6 +131,83 @@ def aggregate(self, inputs, index, ptr, dim_size):
         return aggr, inputs
 
 
+class PropagationNet(InteractionNet):
+    """
+    Alternative version of InteractionNet that incentivizes the propagation
+    of information from sender nodes to receivers.
+    """
+
+    # pylint: disable=arguments-differ
+    # Disable to override args/kwargs from superclass
+
+    def __init__(
+        self,
+        edge_index,
+        input_dim,
+        update_edges=True,
+        hidden_layers=1,
+        hidden_dim=None,
+        edge_chunk_sizes=None,
+        aggr_chunk_sizes=None,
+        aggr="sum",
+    ):
+        # Use mean aggregation in propagation version to avoid instability
+        super().__init__(
+            edge_index,
+            input_dim,
+            update_edges=update_edges,
+            hidden_layers=hidden_layers,
+            hidden_dim=hidden_dim,
+            edge_chunk_sizes=edge_chunk_sizes,
+            aggr_chunk_sizes=aggr_chunk_sizes,
+            aggr="mean",
+        )
+
+    def forward(self, send_rep, rec_rep, edge_rep):
+        """
+        Apply propagation network to update the representations of receiver
+        nodes, and optionally the edge representations.
+
+        send_rep: (N_send, d_h), vector representations of sender nodes
+        rec_rep: (N_rec, d_h), vector representations of receiver nodes
+        edge_rep: (M, d_h), vector representations of edges used
+
+        Returns:
+        rec_rep: (N_rec, d_h), updated vector representations of receiver
+            nodes
+        (optionally) edge_rep: (M, d_h), updated vector representations
+            of edges
+        """
+        # Always concatenate to [rec_nodes, send_nodes] for propagation,
+        # but only aggregate to rec_nodes
+        node_reps = torch.cat((rec_rep, send_rep), dim=-2)
+        edge_rep_aggr, edge_diff = self.propagate(
+            self.edge_index, x=node_reps, edge_attr=edge_rep
+        )
+        rec_diff = self.aggr_mlp(
+            torch.cat((rec_rep, edge_rep_aggr), dim=-1)
+        )
+
+        # Residual connections
+        rec_rep = edge_rep_aggr + rec_diff  # residual is to aggregation
+
+        if self.update_edges:
+            edge_rep = edge_rep + edge_diff
+            return rec_rep, edge_rep
+
+        return rec_rep
+
+    def message(self, x_j, x_i, edge_attr):
+        """
+        Compute messages from node j to node i.
+        """
+        # Residual connection is to sender node, propagating information
+        # to edge
+        return x_j + self.edge_mlp(
+            torch.cat((edge_attr, x_j, x_i), dim=-1)
+        )
+
+
 class SplitMLPs(nn.Module):
     """
     Module that feeds chunks of input through different MLPs.

neural_lam/models/__init__.py

@@ -1,6 +1,11 @@
 # Local
-from .base_graph_model import BaseGraphModel
-from .base_hi_graph_model import BaseHiGraphModel
+from .forecaster_module import ForecasterModule
 from .graph_lam import GraphLAM
 from .hi_lam import HiLAM
 from .hi_lam_parallel import HiLAMParallel
+
+MODELS = {
+    "graph_lam": GraphLAM,
+    "hi_lam": HiLAM,
+    "hi_lam_parallel": HiLAMParallel,
+}

neural_lam/models/ar_forecaster.py

@@ -0,0 +1,84 @@
+# Third-party
+import torch
+
+# Local
+from ..datastore import BaseDatastore
+from .forecaster import Forecaster
+from .step_predictor import StepPredictor
+
+
+class ARForecaster(Forecaster):
+    """
+    Subclass of Forecaster that uses an auto-regressive strategy to
+    unroll a forecast. Makes use of a StepPredictor at each AR step.
+    """
+
+    def __init__(self, predictor: StepPredictor, datastore: BaseDatastore):
+        super().__init__()
+        self.predictor = predictor
+
+        # Register boundary/interior masks on the forecaster, not the predictor
+        boundary_mask = (
+            torch.tensor(datastore.boundary_mask.values, dtype=torch.float32)
+            .unsqueeze(0)
+            .unsqueeze(-1)
+        )
+        self.register_buffer("boundary_mask", boundary_mask, persistent=False)
+        self.register_buffer(
+            "interior_mask", 1.0 - self.boundary_mask, persistent=False
+        )
+
+    @property
+    def predicts_std(self) -> bool:
+        return self.predictor.predicts_std
+
+    def forward(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        boundary_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Unroll the autoregressive model.
+        boundary_states is used ONLY to overwrite boundary nodes at each step.
+        The interior prediction at step i must not depend on
+        boundary_states[:, i] in any other way.
+        """
+
+        prev_prev_state = init_states[:, 0]
+        prev_state = init_states[:, 1]
+        prediction_list = []
+        pred_std_list = []
+        pred_steps = forcing_features.shape[1]
+
+        for i in range(pred_steps):
+            forcing = forcing_features[:, i]
+            boundary_state = boundary_states[:, i]
+
+            pred_state, pred_std = self.predictor(
+                prev_state, prev_prev_state, forcing
+            )
+
+            # Overwrite boundary with true state using ARForecaster's mask
+            new_state = (
+                self.boundary_mask * boundary_state
+                + self.interior_mask * pred_state
+            )
+
+            prediction_list.append(new_state)
+            if pred_std is not None:
+                pred_std_list.append(pred_std)
+
+            # Update conditioning states
+            prev_prev_state = prev_state
+            prev_state = new_state
+
+        prediction = torch.stack(prediction_list, dim=1)
+        # If predictor outputs std, stack it; otherwise return None so
+        # ForecasterModule can substitute the constant per_var_std
+        if pred_std_list:
+            pred_std = torch.stack(pred_std_list, dim=1)
+        else:
+            pred_std = None
+
+        return prediction, pred_std