move az transform into func

src/graphnet/data/extractors/prometheus/prometheus_extractor.py

@@ -85,15 +85,16 @@ 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:
             if len(res["initial_state_azimuth"]) > 0:
                 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]

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