forward() in GraphDefinition

Author: sevmag

src/graphnet/models/data_representation/data_representation.py

@@ -7,7 +7,6 @@
 from numpy.random import default_rng, Generator
 
 from graphnet.models.detector import Detector
-from graphnet.utilities.decorators import final
 from graphnet.models import Model
 from abc import abstractmethod
 
@@ -109,7 +108,6 @@ def __init__(
         else:
             self.rng = default_rng()
 
-    @final
     def forward(  # type: ignore
         self,
         input_features: np.ndarray,
@@ -121,10 +119,10 @@ def forward(  # type: ignore
         loss_weight_default_value: Optional[float] = None,
         data_path: Optional[str] = None,
     ) -> Data:
-        """Construct graph as ´Data´ object.
+        """Construct data as ´Data´ object.
 
         Args:
-            input_features: Input features for graph construction.
+            input_features: Input features for data construction.
                 Shape ´[num_rows, d]´
             input_feature_names: name of each column. Shape ´[,d]´.
             truth_dicts: Dictionary containing truth labels.
@@ -138,7 +136,7 @@ def forward(  # type: ignore
             data_path: Path to dataset data files. Defaults to None.
 
         Returns:
-            graph
+            data
         """
         # Checks
         self._validate_input(
@@ -197,9 +195,6 @@ def forward(  # type: ignore
                 data=data, custom_label_functions=custom_label_functions
             )
 
-        # Do final processing steps
-        data = self._forward_end(data, data_feature_names)
-
         # DEPRECATION STAMP GRAPH_DEFINITION: REMOVE AT 2.0 LAUNCH
         # See https://github.com/graphnet-team/graphnet/issues/647
         data["graph_definition"] = self.__class__.__name__
@@ -439,18 +434,6 @@ def _create_data(self, input_features: torch.Tensor) -> Data:
         """
         raise NotImplementedError
 
-    def _forward_end(
-        self,
-        data: Data,
-        data_feature_names: List[str],
-    ) -> Data:
-        """Place to add any final data processing steps.
-
-        Override this method to add any final processing steps in the end of
-        the forward pass.
-        """
-        return data
-
     def _label_repeater(self, label: torch.Tensor, data: Data) -> torch.Tensor:
         """Handle the label repetition.
 

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

@@ -5,9 +5,10 @@
 and can be passed to dataloaders during training and deployment.
 """
 
-from typing import List, Optional, Dict, Union, Tuple
+from typing import List, Optional, Dict, Union, Tuple, Any, Callable
 import torch
 from numpy.random import Generator
+import numpy as np
 
 from graphnet.models.detector import Detector
 from .edges import EdgeDefinition
@@ -144,6 +145,55 @@ def _create_data(
 
         return data, data_feature_names
 
+    def forward(  # type: ignore
+        self,
+        input_features: np.ndarray,
+        input_feature_names: List[str],
+        truth_dicts: Optional[List[Dict[str, Any]]] = None,
+        custom_label_functions: Optional[Dict[str, Callable[..., Any]]] = None,
+        loss_weight_column: Optional[str] = None,
+        loss_weight: Optional[float] = None,
+        loss_weight_default_value: Optional[float] = None,
+        data_path: Optional[str] = None,
+    ) -> Data:
+        """Construct graph as ´Data´ object.
+
+        Args:
+            input_features: Input features for graph construction.
+                Shape ´[num_rows, d]´
+            input_feature_names: name of each column. Shape ´[,d]´.
+            truth_dicts: Dictionary containing truth labels.
+            custom_label_functions: Custom label functions.
+            loss_weight_column: Name of column that holds loss weight.
+                                Defaults to None.
+            loss_weight: Loss weight associated with event. Defaults to None.
+            loss_weight_default_value: default value for loss weight.
+                    Used in instances where some events have
+                    no pre-defined loss weight. Defaults to None.
+            data_path: Path to dataset data files. Defaults to None.
+
+        Returns:
+            graph
+        """
+        data = super().forward(
+            input_features=input_features,
+            input_feature_names=input_feature_names,
+            truth_dicts=truth_dicts,
+            custom_label_functions=custom_label_functions,
+            loss_weight_column=loss_weight_column,
+            loss_weight=loss_weight,
+            loss_weight_default_value=loss_weight_default_value,
+            data_path=data_path,
+        )
+
+        if self._add_static_features:
+            data = self._add_features_individually(
+                data,
+                self._node_definition._output_feature_names,
+            )
+
+        return data
+
     def _forward_end(
         self,
         data: Data,