Method forward

climart/models/GraphNet/graph_network_block.py:136–205

(self, graph: GraphComponentToTensor)

    #       batch_size = edge_feats_old.shape[0]

    def forward(self, graph: GraphComponentToTensor) -> GraphComponentToTensor:
        if self.use_edge_features:
            edge_feats_old = graph[EDGES]  # edge_feats_old have shape (b, #edges, #edge-feats)
            if edge_feats_old.shape[gn_constants.SPATIAL_DIM] != self._n_edges:
                    raise ValueError(f"Edge features imply {edge_feats_old.shape[gn_constants.SPATIAL_DIM]} edges, "
                                     f"while sender and receiver lists imply {self._n_edges} edges.")
        node_feats_old = graph[NODES]  # node_feats_old have shape (b, #nodes, #node-feats)
        n_nodes = node_feats_old.shape[gn_constants.SPATIAL_DIM]
        if self.use_global_features:
            global_feats_old = graph[GLOBALS]  # global_feats_old have shape (b, #global-feats)
            batch_size, n_glob_feats = global_feats_old.shape

        out = {c: None for c in self.components}

        # ----------------------- Update edges
        # self.senders and self.receivers are a sequence of indices with #edges elements
        sender_feats = node_feats_old.index_select(index=self._senders, dim=gn_constants.SPATIAL_DIM)
        receiver_feats = node_feats_old.index_select(index=self._receivers, dim=gn_constants.SPATIAL_DIM)

        # print('E:', edge_feats_old.shape, 'V:', node_feats_old.shape, 'U:', global_feats_old.shape)
        # print('Senders', sender_feats.shape, 'Recvs', receiver_feats.shape, global_feats_unsqueezed.shape)
        mlp_input_e = torch.cat([
            edge_feats_old,  # (b, #edges, #edge-feats)
            sender_feats,  # (b, #edges, #node-feats)
            receiver_feats,  # (b, #edges, #node-feats)
            repeat(global_feats_old, 'b g -> b e g', e=self._n_edges)  # (1, self.n_edges, 1) (b, 1, #global-feats)
        ], dim=-1)

        mlp_input_e = rearrange(mlp_input_e, 'b e d1 -> (b e) d1')

        out[EDGES] = self.update_funcs[EDGES](mlp_input_e)
        out[EDGES] = rearrange(out[EDGES], '(b e) d2 -> b e d2', b=batch_size, e=self._n_edges)
        if self.residual[EDGES]:
            out[EDGES] += edge_feats_old
        # ----------------------- Update nodes
        aggregated_edge_feats_for_node = self.aggregator_funcs[AggregationTypes.AGG_E_TO_N](
            out[EDGES], index=self._receivers
        )

        mlp_input_n = torch.cat([
            aggregated_edge_feats_for_node,  # (b, #nodes, #edge-feats)
            node_feats_old,  # (b, #nodes, #node-feats)
            repeat(global_feats_old, 'b g -> b n g', n=n_nodes)  # (b, #nodes, #global-feats)
        ], dim=-1)
        mlp_input_n = rearrange(mlp_input_n, 'b n d1 -> (b n) d1')

        # print('Agg E:', aggregated_edge_feats_for_node.shape, 'V:', node_feats_old.shape, mlp_input_n.shape)
        out[NODES] = self.update_funcs[NODES](mlp_input_n)
        out[NODES] = rearrange(out[NODES], '(b n) d2 -> b n d2', b=batch_size, n=n_nodes)
        if self.residual[NODES]:
            out[NODES] += node_feats_old
        # ----------------------- Update global features
        aggregated_edge_feats_for_global = self.aggregator_funcs[AggregationTypes.AGG_E_TO_U](out[EDGES])
        aggregated_node_feats_for_global = self.aggregator_funcs[AggregationTypes.AGG_N_TO_U](out[NODES])
        # print('Agg EU:', aggregated_edge_feats_for_global.shape, 'VU:', aggregated_node_feats_for_global.shape)

        mlp_input_u = torch.cat([
            aggregated_edge_feats_for_global,  # (b, #edge-feats)