diff --git a/straxen/itp_map.py b/straxen/itp_map.py index bf86cf741..33607285e 100644 --- a/straxen/itp_map.py +++ b/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) diff --git a/straxen/plugins/events/_event_s1_positions_base.py b/straxen/plugins/events/_event_s1_positions_base.py index 94ad306a5..be45b1958 100644 --- a/straxen/plugins/events/_event_s1_positions_base.py +++ b/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() diff --git a/straxen/plugins/events/_event_s2_positions_base.py b/straxen/plugins/events/_event_s2_positions_base.py index 187b62f85..923993acd 100644 --- a/straxen/plugins/events/_event_s2_positions_base.py +++ b/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() diff --git a/straxen/plugins/events/event_position_uncertainty.py b/straxen/plugins/events/event_position_uncertainty.py index 1272e7708..7f06aa173 100644 --- a/straxen/plugins/events/event_position_uncertainty.py +++ b/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 diff --git a/straxen/plugins/peaks/_peak_positions_base.py b/straxen/plugins/peaks/_peak_positions_base.py index dbca11920..e71f741ec 100644 --- a/straxen/plugins/peaks/_peak_positions_base.py +++ b/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: diff --git a/straxen/plugins/peaks/_peak_s1_positions_base.py b/straxen/plugins/peaks/_peak_s1_positions_base.py index 619ebb19f..668ed4ea0 100644 --- a/straxen/plugins/peaks/_peak_s1_positions_base.py +++ b/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() diff --git a/straxen/plugins/peaks/peak_basics.py b/straxen/plugins/peaks/peak_basics.py index 6ccd48c8d..8f604c311 100644 --- a/straxen/plugins/peaks/peak_basics.py +++ b/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: diff --git a/straxen/plugins/peaks/peak_positions_cnf.py b/straxen/plugins/peaks/peak_positions_cnf.py index 6c681ad98..cf1c65251 100644 --- a/straxen/plugins/peaks/peak_positions_cnf.py +++ b/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,7 +228,7 @@ 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) @@ -210,10 +236,10 @@ def compute(self, peaks): 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