Fix theta uncertainty bug

straxen/itp_map.py

@@ -55,7 +55,7 @@ def __call__(self, points):
 
         # fill non valid values with nan and compute the others
         # fill method slightly faster than multiplication of np.ones with nan
-        result[~valid] = float("nan")
+        result[~valid] = np.nan
 
         # Get distances to neighbours_to_use nearest neighbours
         distances, indices = self.kdtree.query(points[valid], self.neighbours_to_use)

straxen/plugins/events/_event_s1_positions_base.py

@@ -80,9 +80,9 @@ def compute(self, events):
         result = np.ones(len(events), dtype=self.dtype)
         result["time"], result["endtime"] = events["time"], strax.endtime(events)
 
-        result["event_x_" + self.algorithm] *= float("nan")
-        result["event_y_" + self.algorithm] *= float("nan")
-        result["event_z_" + self.algorithm] *= float("nan")
+        result["event_x_" + self.algorithm] *= np.nan
+        result["event_y_" + self.algorithm] *= np.nan
+        result["event_z_" + self.algorithm] *= np.nan
 
         model = self.get_tf_model()
 

straxen/plugins/events/_event_s2_positions_base.py

@@ -81,10 +81,10 @@ def get_tf_model(self):
     def compute(self, events):
         result = np.ones(len(events), dtype=self.dtype)
         result["time"], result["endtime"] = events["time"], strax.endtime(events)
-        result["event_s2_x_" + self.algorithm] *= float("nan")
-        result["event_s2_y_" + self.algorithm] *= float("nan")
-        result["event_alt_s2_x_" + self.algorithm] *= float("nan")
-        result["event_alt_s2_y_" + self.algorithm] *= float("nan")
+        result["event_s2_x_" + self.algorithm] *= np.nan
+        result["event_s2_y_" + self.algorithm] *= np.nan
+        result["event_alt_s2_x_" + self.algorithm] *= np.nan
+        result["event_alt_s2_y_" + self.algorithm] *= np.nan
 
         model = self.get_tf_model()
 

straxen/plugins/events/event_position_uncertainty.py

@@ -2,6 +2,7 @@
 import strax
 import straxen
 from straxen.plugins.defaults import DEFAULT_POSREC_ALGO
+from straxen.plugins.peaks.peak_positions_cnf import PeakPositionsCNF
 
 
 export, __all__ = strax.exporter()
@@ -96,8 +97,8 @@ def infer_dtype(self):
             dtype += [
                 (
                     (
-                        f"Naive flow position contour for {infoline[type_]}",
-                        f"{type_}_position_contour_cnf_naive",
+                        f"Position uncertainty contour for {infoline[type_]}",
+                        f"{type_}_position_contour_cnf",
                     ),
                     np.float32,
                     (self.n_poly + 1, 2),
@@ -107,16 +108,19 @@ def infer_dtype(self):
         # Add fields for flow position contour and corrected positions
         dtype += [
             (
-                (f"Flow position contour", "position_contour_cnf"),
+                (
+                    "Position uncertainty contour, field-distortion corrected (cm)",
+                    "position_contour_cnf_fdc",
+                ),
                 np.float32,
                 (self.n_poly + 1, 2),
             ),
             (
-                (f"Flow x position", "x_cnf_fdc"),
+                ("Reconstructed cnf S2 X position, field-distortion corrected (cm)", "x_cnf_fdc"),
                 np.float32,
             ),
             (
-                (f"Flow y position", "y_cnf_fdc"),
+                ("Reconstructed cnf S2 Y position, field-distortion corrected (cm)", "y_cnf_fdc"),
                 np.float32,
             ),
         ]
@@ -146,8 +150,8 @@ def compute(self, events, peaks):
 
         # Initialize contour fields with NaN values
         for type_ in ["s2", "alt_s2"]:
-            result[f"{type_}_position_contour_cnf_naive"] *= np.nan
-        result[f"position_contour_cnf"] *= np.nan
+            result[f"{type_}_position_contour_cnf"] *= np.nan
+        result["position_contour_cnf_fdc"] *= np.nan
 
         # Split peaks by containment in events
         split_peaks = strax.split_by_containment(peaks, events)
@@ -158,16 +162,16 @@ def compute(self, events, peaks):
                 type_index = event[f"{type_}_index"]
                 if type_index != -1:
                     # Store naive flow position contour
-                    result[f"{type_}_position_contour_cnf_naive"][event_i] = sp[
-                        "position_contours_cnf"
-                    ][type_index]
+                    result[f"{type_}_position_contour_cnf"][event_i] = sp["position_contour_cnf"][
+                        type_index
+                    ]
 
                 if type_ == "s2":
                     if self.use_fdc_for_contour:
                         # Apply full field distortion correction to contour
                         contour_plus_xy = np.concatenate(
                             [
-                                sp["position_contours_cnf"][type_index],
+                                sp["position_contour_cnf"][type_index],
                                 np.array([[sp["x_cnf"][type_index], sp["y_cnf"][type_index]]]),
                             ],
                             axis=0,
@@ -184,14 +188,14 @@ def compute(self, events, peaks):
                         scaled_2d_contour = contour_with_z[:, :2] * scale[:, np.newaxis]
 
                         # Store corrected contour and positions
-                        result["position_contour_cnf"][event_i] = scaled_2d_contour[:-1]
+                        result["position_contour_cnf_fdc"][event_i] = scaled_2d_contour[:-1]
                         result["x_cnf_fdc"][event_i] = scaled_2d_contour[-1, 0]
                         result["y_cnf_fdc"][event_i] = scaled_2d_contour[-1, 1]
                     else:
                         # Apply simple scaling based on field distortion correction
                         scale = event["r_field_distortion_correction"] / event["r_naive"] + 1
-                        result["position_contour_cnf"][event_i] = (
-                            sp["position_contours_cnf"][type_index] * scale
+                        result["position_contour_cnf_fdc"][event_i] = (
+                            sp["position_contour_cnf"][type_index] * scale
                         )
                         result["x_cnf_fdc"][event_i] = sp["x_cnf"][type_index] * scale
                         result["y_cnf_fdc"][event_i] = sp["y_cnf"][type_index] * scale
@@ -224,9 +228,9 @@ class EventPositionUncertainty(strax.Plugin):
 
     """
 
-    __version__ = "0.0.1"
+    __version__ = "0.0.2"
 
-    depends_on = ("event_info", "event_position_contour")
+    depends_on = ("event_info", "event_positions", "event_position_contour")
     provides = "event_position_uncertainty"
 
     def infer_dtype(self):
@@ -261,11 +265,11 @@ def infer_dtype(self):
         # Add fields for corrected position uncertainties
         dtype += [
             (
-                ("Flow position uncertainty in r (cm)", "r_position_uncertainty"),
+                ("Position uncertainty in r (cm)", "r_position_uncertainty"),
                 np.float32,
             ),
             (
-                ("Flow position uncertainty in theta (rad)", "theta_position_uncertainty"),
+                ("Position uncertainty in theta (rad)", "theta_position_uncertainty"),
                 np.float32,
             ),
         ]
@@ -289,21 +293,19 @@ def compute(self, events):
         # Calculate uncertainties for main and alternative S2 signals
         for type_ in ["s2", "alt_s2"]:
             # Calculate radial uncertainties
-            r_array = np.linalg.norm(events[f"{type_}_position_contour_cnf_naive"], axis=2)
+            r_array = np.linalg.norm(events[f"{type_}_position_contour_cnf"], axis=2)
             r_min = np.min(r_array, axis=1)
             r_max = np.max(r_array, axis=1)
 
             # Calculate angular uncertainties
             theta_array = np.arctan2(
-                events[f"{type_}_position_contour_cnf_naive"][..., 1],
-                events[f"{type_}_position_contour_cnf_naive"][..., 0],
+                events[f"{type_}_position_contour_cnf"][..., 1],
+                events[f"{type_}_position_contour_cnf"][..., 0],
             )
-            theta_min = np.min(theta_array, axis=1)
-            theta_max = np.max(theta_array, axis=1)
-            theta_diff = theta_max - theta_min
 
-            # Correct for circular nature of angles
-            theta_diff[theta_diff > np.pi] -= 2 * np.pi
+            avg_theta = np.arctan2(events[f"{type_}_y_cnf"], events[f"{type_}_x_cnf"])
+
+            theta_diff = PeakPositionsCNF.calculate_theta_diff(theta_array, avg_theta)
 
             # Store uncertainties
             result[f"{type_}_r_position_uncertainty"] = (r_max - r_min) / 2

straxen/plugins/peaks/_peak_positions_base.py

@@ -81,8 +81,8 @@ def compute(self, peaks):
         result = np.ones(len(peaks), dtype=self.dtype)
         result["time"], result["endtime"] = peaks["time"], strax.endtime(peaks)
 
-        result["x_" + self.algorithm] *= float("nan")
-        result["y_" + self.algorithm] *= float("nan")
+        result["x_" + self.algorithm] *= np.nan
+        result["y_" + self.algorithm] *= np.nan
         model = self.get_tf_model()
 
         if model is None:

straxen/plugins/peaks/_peak_s1_positions_base.py

@@ -80,9 +80,9 @@ def compute(self, peaks):
         result = np.ones(len(peaks), dtype=self.dtype)
         result["time"], result["endtime"] = peaks["time"], strax.endtime(peaks)
 
-        result["x_" + self.algorithm] *= float("nan")
-        result["y_" + self.algorithm] *= float("nan")
-        result["z_" + self.algorithm] *= float("nan")
+        result["x_" + self.algorithm] *= np.nan
+        result["y_" + self.algorithm] *= np.nan
+        result["z_" + self.algorithm] *= np.nan
 
         model = self.get_tf_model()
 

straxen/plugins/peaks/peak_basics.py

@@ -98,7 +98,7 @@ def compute(self, peaks):
         # Negative-area peaks get NaN AFT
         m = p["area"] > 0
         r["area_fraction_top"][m] = area_top[m] / area_total[m]
-        r["area_fraction_top"][~m] = float("nan")
+        r["area_fraction_top"][~m] = np.nan
         r["rise_time"] = -p["area_decile_from_midpoint"][:, 1]
 
         if self.check_peak_sum_area_rtol is not None:

straxen/plugins/peaks/peak_positions_cnf.py

@@ -29,7 +29,7 @@ class PeakPositionsCNF(PeakPositionsBaseNT):
 
     """
 
-    __version__ = "0.0.3"
+    __version__ = "0.0.4"
     depends_on = "peaks"
     provides = "peak_positions_cnf"
     algorithm = "cnf"
@@ -68,6 +68,32 @@ class PeakPositionsCNF(PeakPositionsBaseNT):
         help="Compiled JAX function",
     )
 
+    @staticmethod
+    def calculate_theta_diff(theta_array, avg_theta):
+        """Calculate the difference between maximum and minimum angles from an array of angles by
+        normalizing the angular difference into the range [0, 2π).
+
+        Parameters:
+            theta_array : np.ndarray
+            A 2D numpy array where each row represents a set of angles in radians.
+            avg_theta : np.ndarray
+            A 1D numpy array representing the average angle in radians for each
+            row in `theta_array`.
+
+        Returns:
+            theta_diff : np.ndarray
+            A 1D numpy array with the difference between the maximum and minimum angles in radians
+            for each row in `theta_array`
+
+        """
+        # Correction to handle circular nature of angles
+        theta_array_shift = (theta_array - avg_theta[..., np.newaxis] + np.pi) % (2 * np.pi)
+        theta_min = np.min(theta_array_shift, axis=1)
+        theta_max = np.max(theta_array_shift, axis=1)
+        theta_diff = theta_max - theta_min
+
+        return theta_diff
+
     def infer_dtype(self):
         """Define the data type for the output.
 
@@ -91,7 +117,7 @@ def infer_dtype(self):
                 np.float32,
             ),
             (
-                ("Position uncertainty contour", f"position_contours_{self.algorithm}"),
+                ("Position uncertainty contour (cm)", f"position_contour_{self.algorithm}"),
                 np.float32,
                 (self.n_poly + 1, 2),
             ),
@@ -172,9 +198,9 @@ def compute(self, peaks):
         result["time"], result["endtime"] = peaks["time"], strax.endtime(peaks)
 
         # Set default values to NaN
-        result[f"x_{self.algorithm}"] *= float("nan")
-        result[f"y_{self.algorithm}"] *= float("nan")
-        result[f"position_contours_{self.algorithm}"] *= float("nan")
+        result[f"x_{self.algorithm}"] *= np.nan
+        result[f"y_{self.algorithm}"] *= np.nan
+        result[f"position_contour_{self.algorithm}"] *= np.nan
         result[f"r_uncertainty_{self.algorithm}"] *= np.nan
         result[f"theta_uncertainty_{self.algorithm}"] *= np.nan
 
@@ -202,18 +228,18 @@ def compute(self, peaks):
         # Write output to the result array
         result[f"x_{self.algorithm}"][peak_mask] = xy[:, 0]
         result[f"y_{self.algorithm}"][peak_mask] = xy[:, 1]
-        result[f"position_contours_{self.algorithm}"][peak_mask] = contours
+        result[f"position_contour_{self.algorithm}"][peak_mask] = contours
 
         # Calculate uncertainties in r and theta
         r_array = np.linalg.norm(contours, axis=2)
         r_min = np.min(r_array, axis=1)
         r_max = np.max(r_array, axis=1)
 
         theta_array = np.arctan2(contours[..., 1], contours[..., 0])
-        theta_min = np.min(theta_array, axis=1)
-        theta_max = np.max(theta_array, axis=1)
-        theta_diff = theta_max - theta_min
-        theta_diff[theta_diff > np.pi] -= 2 * np.pi
+
+        avg_theta = np.arctan2(result[f"y_{self.algorithm}"], result[f"x_{self.algorithm}"])
+
+        theta_diff = self.calculate_theta_diff(theta_array, avg_theta)
 
         result[f"r_uncertainty_{self.algorithm}"][peak_mask] = (r_max - r_min) / 2
         result[f"theta_uncertainty_{self.algorithm}"][peak_mask] = np.abs(theta_diff) / 2