generalize utilities to both parquet formats

RasmusOrsoe · Oct 30, 2024 · dbccfbf · dbccfbf
1 parent 822b7a5
commit dbccfbf
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/...eometry_tables/prometheus/icecube.parquet → ...metry_tables/prometheus/icecube86.parquet b/...eometry_tables/prometheus/icecube.parquet → ...metry_tables/prometheus/icecube86.parquet
diff --git a/src/graphnet/data/extractors/prometheus/utilities.py b/src/graphnet/data/extractors/prometheus/utilities.py
@@ -36,7 +36,9 @@ def compute_visible_inelasticity(mc_truth: pd.DataFrame) -> float:
                 f"event. Got {mc_truth.keys()}"
             )
 
-    final_type_1, final_type_2 = abs(mc_truth["final_state_type"])
+    # final_type_1, final_type_2 = abs(mc_truth["final_state_type"])
+    final_state = abs(mc_truth["final_state_type"])
+    final_type_1, final_type_2 = final_state[0], final_state[1]
     if mc_truth["interaction"] != 1:
         visible_inelasticity = 1.0
     elif not (final_type_1 == 13 or final_type_2 == 13):
@@ -64,7 +66,8 @@ def get_muon_direction(
     mc_truth: pd.DataFrame, transform_az: bool = False
 ) -> Tuple[float, float]:
     """Get angles of muon in nu_mu CC events."""
-    final_type_1, final_type_2 = abs(mc_truth["final_state_type"])
+    final_state = abs(mc_truth["final_state_type"])
+    final_type_1, final_type_2 = final_state[0], final_state[1]
     if mc_truth["interaction"] != 1:
         muon_zenith = -1
         muon_azimuth = -1

diff --git a/src/graphnet/data/readers/prometheus_reader.py b/src/graphnet/data/readers/prometheus_reader.py
@@ -73,12 +73,13 @@ def __call__(self, file_path: str) -> List[OrderedDict]:
                         try:
                             output = extractor(file[extractor._table][k])
                             extracted_event[extractor._extractor_name] = output
-                        except:  # noqa
+                        except Exception as e:  # noqa
                             self.warning(
                                 "Unable to apply "
                                 f"{extractor.__class__.__name__} to"
                                 f" {file_path}."
                             )
+
                 # Apply filter. If one filter returns False event is skipped.
                 if self._keep_event(extracted_event=extracted_event):
                     outputs.append(extracted_event)