Fixing graph processing logic

Author: sevmag

src/graphnet/models/data_representation/graphs/graph_definition.py

@@ -112,6 +112,9 @@ def __init__(
         self._sort_by = sort_by
         self._add_static_features = add_static_features
 
+        # make sure output feature names are set also in node definition
+        self._set_output_feature_names(self._input_feature_names)
+
     def _set_output_feature_names(
         self, input_feature_names: List[str]
     ) -> List[str]:
@@ -160,9 +163,8 @@ def forward(  # type: ignore
             loss_weight_default_value=loss_weight_default_value,
             data_path=data_path,
         )
-
         # Create graph & get new node feature names
-        data = self._node_definition(data.x)
+        data.x = self._node_definition(data.x)
         if self._sort_by is not None:
             data.x = data.x[data.x[:, self._sort_by].sort()[1]]
 
@@ -178,7 +180,6 @@ def forward(  # type: ignore
                 data,
                 self.output_feature_names,
             )
-
         return data
 
     def _add_features_individually(

src/graphnet/models/data_representation/graphs/nodes/nodes.py

@@ -34,7 +34,7 @@ def __init__(
             )
 
     @final
-    def forward(self, x: torch.tensor) -> Tuple[Data, List[str]]:
+    def forward(self, x: torch.tensor) -> torch.tensor:
         """Construct nodes from raw node features.
 
         Args:
@@ -45,9 +45,9 @@ def forward(self, x: torch.tensor) -> Tuple[Data, List[str]]:
         Returns:
             graph: a graph without edges
         """
-        graph = self._construct_nodes(x=x)
+        data = self._construct_nodes(x=x)
 
-        return graph
+        return data
 
     @property
     def _output_feature_names(self) -> List[str]:
@@ -110,7 +110,7 @@ def _define_output_feature_names(
         """
 
     @abstractmethod
-    def _construct_nodes(self, x: torch.tensor) -> Data:
+    def _construct_nodes(self, x: torch.tensor) -> torch.tensor:
         """Construct nodes from raw node features ´x´.
 
         Args:
@@ -133,8 +133,8 @@ def _define_output_feature_names(
     ) -> List[str]:
         return input_feature_names
 
-    def _construct_nodes(self, x: torch.Tensor) -> Data:
-        return Data(x=x)
+    def _construct_nodes(self, x: torch.Tensor) -> torch.Tensor:
+        return x
 
 
 class PercentileClusters(NodeDefinition):
@@ -195,7 +195,7 @@ def _get_indices_and_feature_names(
             new_feature_names.append("counts")
         return cluster_idx, summ_idx, new_feature_names
 
-    def _construct_nodes(self, x: torch.Tensor) -> Data:
+    def _construct_nodes(self, x: torch.Tensor) -> torch.Tensor:
         # Cast to Numpy
         x = x.numpy()
         # Construct clusters with percentile-summarized features
@@ -219,7 +219,7 @@ def _construct_nodes(self, x: torch.Tensor) -> Data:
             )  # noqa
             raise AttributeError
 
-        return Data(x=torch.tensor(array))
+        return torch.tensor(array)
 
 
 class NodeAsDOMTimeSeries(NodeDefinition):
@@ -270,7 +270,7 @@ def _define_output_feature_names(
     ) -> List[str]:
         return input_feature_names + ["new_node_col"]
 
-    def _construct_nodes(self, x: torch.Tensor) -> Data:
+    def _construct_nodes(self, x: torch.Tensor) -> torch.Tensor:
         """Construct nodes from raw node features ´x´."""
         # Cast to Numpy
         x = x.numpy()
@@ -309,7 +309,7 @@ def _construct_nodes(self, x: torch.Tensor) -> Data:
         new_node_col[0] = 1
         x = np.column_stack([x, new_node_col])
 
-        return Data(x=torch.tensor(x))
+        return torch.tensor(x)
 
 
 class IceMixNodes(NodeDefinition):
@@ -440,7 +440,7 @@ def _pulse_sampler(
 
         return ids
 
-    def _construct_nodes(self, x: torch.Tensor) -> Data:
+    def _construct_nodes(self, x: torch.Tensor) -> torch.Tensor:
 
         event_length = x.shape[0]
         if self.hlc_name is not None:
@@ -468,4 +468,4 @@ def _construct_nodes(self, x: torch.Tensor) -> Data:
         for idx, feature in enumerate(non_ice_features):
             graph[:event_length, idx] = x[ids, self.feature_indexes[feature]]
 
-        return Data(x=graph)
+        return graph

tests/models/test_node_definition.py

@@ -40,7 +40,7 @@ def test_percentile_cluster() -> None:
     # Apply node definition to torch tensor with raw pulses
     graph = node_definition(tensor)
     new_features = node_definition._output_feature_names
-    x_tilde = graph.x.numpy()
+    x_tilde = graph.numpy()
 
     # Calculate percentiles "the normal way" and compare that output of
     # node definition match.