From 8cb660bf5ca0e6e647238b233677f8b1c7a77aba Mon Sep 17 00:00:00 2001
From: RasmusOrsoe <rahn@outlook.dk>
Date: Tue, 17 Sep 2024 16:02:04 +0200
Subject: [PATCH] move az transform into func

---
 .../data/extractors/prometheus/prometheus_extractor.py      | 5 +++--
 src/graphnet/data/extractors/prometheus/utilities.py        | 6 +++++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/graphnet/data/extractors/prometheus/prometheus_extractor.py b/src/graphnet/data/extractors/prometheus/prometheus_extractor.py
index d32d8eea0..2bb4f3ebf 100644
--- a/src/graphnet/data/extractors/prometheus/prometheus_extractor.py
+++ b/src/graphnet/data/extractors/prometheus/prometheus_extractor.py
@@ -85,7 +85,9 @@ def __call__(self, event: pd.DataFrame) -> pd.DataFrame:
         """Extract event-level truth information."""
         # Extract data
         visible_inelasticity = compute_visible_inelasticity(event)
-        muon_zenith, muon_azimuth = get_muon_direction(event)
+        muon_zenith, muon_azimuth = get_muon_direction(
+            event, self._transform_az
+        )
         res = super().__call__(event=event)
         # transform azimuth from [-pi, pi] to [0, 2pi] if wanted
         if self._transform_az:
@@ -93,7 +95,6 @@ def __call__(self, event: pd.DataFrame) -> pd.DataFrame:
                 azimuth = np.asarray(res["initial_state_azimuth"]) + np.pi
                 azimuth = azimuth.tolist()  # back to list
                 res["initial_state_azimuth"] = azimuth
-                muon_azimuth += np.pi
         res["visible_inelasticity"] = [visible_inelasticity]
         res["muon_azimuth"] = [muon_azimuth]
         res["muon_zenith"] = [muon_zenith]
diff --git a/src/graphnet/data/extractors/prometheus/utilities.py b/src/graphnet/data/extractors/prometheus/utilities.py
index 74cda9d42..3c6b953b7 100644
--- a/src/graphnet/data/extractors/prometheus/utilities.py
+++ b/src/graphnet/data/extractors/prometheus/utilities.py
@@ -60,7 +60,9 @@ def compute_visible_inelasticity(mc_truth: pd.DataFrame) -> float:
     return visible_inelasticity
 
 
-def get_muon_direction(mc_truth: pd.DataFrame) -> Tuple[float, float]:
+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"])
     if mc_truth["interaction"] != 1:
@@ -77,6 +79,8 @@ def get_muon_direction(mc_truth: pd.DataFrame) -> Tuple[float, float]:
         muon_azimuth = mc_truth["final_state_azimuth"][
             abs(mc_truth["final_state_type"]) == 13
         ][0]
+        if transform_az:
+            muon_azimuth += np.pi
     return muon_zenith, muon_azimuth