add shape property

Author: sevmag

src/graphnet/models/data_representation/images/mappings/pixel_mappings.py

@@ -47,6 +47,18 @@ def _set_image_feature_names(self, input_feature_names: List[str]) -> None:
         """Set the final image feature names."""
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def shape(
+        self,
+    ) -> List[List[int]]:
+        """Return the shape of the output images as a list of tuples.
+
+        In the dimensions (F,D,H,W) where F is the number of features
+        per pixel. And D,H,W are the dimension of the image
+        """
+        pass
+
 
 class IC86PixelMapping(PixelMapping):
     """Mapping for the IceCube86.
@@ -230,6 +242,20 @@ def _set_image_feature_names(self, input_feature_names: List[str]) -> None:
             if infeature not in [self._string_label, self._dom_number_label]
         ]
 
+    @property
+    def shape(
+        self,
+    ) -> List[List[int]]:
+        """Return the shape of the output images as a list of tuples."""
+        ret = []
+        if self._include_main_array:
+            ret.append([self._nb_cnn_features, 10, 10, 60])
+        if self._include_upper_dc:
+            ret.append([self._nb_cnn_features, 1, 8, 10])
+        if self._include_lower_dc:
+            ret.append([self._nb_cnn_features, 1, 8, 50])
+        return ret
+
 
 class ExamplePrometheusMapping(PixelMapping):
     """Mapping for the Prometheus detector.
@@ -363,3 +389,10 @@ def _set_image_feature_names(self, input_feature_names: List[str]) -> None:
             for infeature in input_feature_names
             if infeature not in [self._string_label, self._sensor_number_label]
         ]
+
+    @property
+    def shape(
+        self,
+    ) -> List[List[int]]:
+        """Return the shape of the output images as a list of tuples."""
+        return [[self._nb_cnn_features, 8, 9, 22]]

tests/models/test_pixel_mapping.py

@@ -68,11 +68,24 @@ def test_pixel_mappings() -> None:
     # Apply node definition to torch tensor with raw pulses
     picture = pixel_mapping(dummy_data, pixel_feature_names)
     new_features = pixel_mapping.image_feature_names
+    n_features = len(new_features)
 
     # Check the output
     basic_checks_picture(picture, dtype)
 
     # More checks
+    assert (
+        len(pixel_mapping.shape) == 3
+    ), f"Expected shape to be 3 got {len(pixel_mapping.shape)}"
+    assert pixel_mapping.shape == [
+        (n_features, 10, 10, 60),
+        (n_features, 1, 8, 10),
+        (n_features, 1, 8, 50),
+    ], (
+        f"Expected shape to be [({n_features},10,10,60), "
+        f"({n_features},1,8,10), ({n_features},1,8,50)] got "
+        f"{pixel_mapping.shape}"
+    )
     assert isinstance(
         new_features, list
     ), f"Output should be a list of feature names got {type(new_features)}"