Fixing output_feature_names issue

Author: sevmag

src/graphnet/models/data_representation/data_representation.py

@@ -173,9 +173,7 @@ def forward(  # type: ignore
         input_features = self._detector(input_features, input_feature_names)
 
         # Create data & get new final data feature names
-        data, data_feature_names = self._create_data(
-            input_features=input_features
-        )
+        data = self._create_data(input_features=input_features)
 
         # Attach number of pulses as static attribute.
         data.n_pulses = torch.tensor(len(input_features), dtype=torch.int32)

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

@@ -124,7 +124,7 @@ def _create_data(
         self, input_features: torch.Tensor
     ) -> Tuple[Data, List[str]]:
         # Create graph & get new node feature names
-        data, data_feature_names = self._node_definition(input_features)
+        data = self._node_definition(input_features)
         if self._sort_by is not None:
             data.x = data.x[data.x[:, self._sort_by].sort()[1]]
 
@@ -135,7 +135,7 @@ def _create_data(
         if self._edge_definition is not None:
             data = self._edge_definition(data)
 
-        return data, data_feature_names
+        return data
 
     def forward(  # type: ignore
         self,
@@ -181,7 +181,7 @@ def forward(  # type: ignore
         if self._add_static_features:
             data = self._add_features_individually(
                 data,
-                self._node_definition._output_feature_names,
+                self.output_feature_names,
             )
 
         return data

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

@@ -46,8 +46,14 @@ def forward(self, x: torch.tensor) -> Tuple[Data, List[str]]:
             graph: a graph without edges
         """
         graph = self._construct_nodes(x=x)
+
+        return graph
+
+    @property
+    def _output_feature_names(self) -> List[str]:
+        """Return output feature names."""
         try:
-            self._output_feature_names
+            self._hidden_output_feature_names
         except AttributeError as e:
             self.error(
                 f"""{self.__class__.__name__} was instantiated without
@@ -57,7 +63,7 @@ def forward(self, x: torch.tensor) -> Tuple[Data, List[str]]:
                        with `input_feature_names`."""
             )  # noqa
             raise e
-        return graph, self._output_feature_names
+        return self._hidden_output_feature_names
 
     @property
     def nb_outputs(self) -> int:
@@ -85,7 +91,7 @@ def set_output_feature_names(self, input_feature_names: List[str]) -> None:
             input_feature_names: List of column names of the input to the
             node definition.
         """
-        self._output_feature_names = self._define_output_feature_names(
+        self._hidden_output_feature_names = self._define_output_feature_names(
             input_feature_names
         )