From 9dbd0fdc991860b21101208e9f7ffdb33043730d Mon Sep 17 00:00:00 2001 From: dachengx Date: Wed, 24 Apr 2024 19:17:17 +0800 Subject: [PATCH 1/3] Add events building --- axidence/plugins/salting/event_building.py | 122 +++++++++++++++++++-- axidence/plugins/salting/salting_events.py | 41 ++++--- axidence/plugins/salting/salting_peaks.py | 5 +- 3 files changed, 142 insertions(+), 26 deletions(-) diff --git a/axidence/plugins/salting/event_building.py b/axidence/plugins/salting/event_building.py index 91b6152..911445c 100644 --- a/axidence/plugins/salting/event_building.py +++ b/axidence/plugins/salting/event_building.py @@ -1,24 +1,130 @@ import numpy as np +from tqdm import tqdm import strax from strax import Plugin -from straxen import Events +import straxen +from straxen import Events, EventBasics +from axidence.plugin import ExhaustPlugin -class SaltedEvents(Events): + +class SaltedEvents(Events, ExhaustPlugin): __version__ = "0.0.0" depends_on = ("salting_peaks", "salting_peak_proximity", "peak_basics", "peak_proximity") provides = "events" data_kind = "events" save_when = strax.SaveWhen.EXPLICIT - dtype = [ - ("event_number", np.int64, "Event number in this dataset"), - ("time", np.int64, "Event start time in ns since the unix epoch"), - ("endtime", np.int64, "Event end time in ns since the unix epoch"), - ] + n_drift_time_window = straxen.URLConfig( + default=5, + type=int, + track=True, + help="How many max drift time will the event builder extend", + ) + + disable_tqdm = straxen.URLConfig( + default=True, + type=bool, + track=False, + help="Whether to disable tqdm", + ) + + def __init__(self): + super().__init__() + self.dtype = super().dtype + [ + (("Salting number of events", "salt_number"), np.int64), + (("Whether the salting event can trigger", "is_triggering"), bool), + ] + + def setup(self): + super().setup() + + # intersection depends_on's dtype.names will be needed in event building + self.needed_fields = set.intersection( + *tuple( + set.union(*tuple(set(self.deps[d].dtype_for(d).names) for d in dk)) + for dk in self.dependencies_by_kind().values() + ) + ) + dk = list(self.dependencies_by_kind().values())[0] + dtype_reference = strax.merged_dtype([self.deps[d].dtype_for(d) for d in dk]) + self._peaks_dtype = [] + for x in dtype_reference: + if x[0][1] in self.needed_fields: + self._peaks_dtype.append(x) + if len(self._peaks_dtype) != len(self.needed_fields): + raise ValueError( + "Weird! Could not find all needed fields " + f"{self.needed_fields} in {dtype_reference}!" + ) + + self.window = self.n_drift_time_window * self.drift_time_max + + if self.window < self.left_extension + self.right_extension: + raise ValueError( + f"The window {self.window}ns is too small to extend the event " + f"while the gap_threshold is about {self.left_extension + self.right_extension}!" + ) + + def get_window_size(self): + return max(super().get_window_size(), self.window * 10) def compute(self, salting_peaks, peaks, start, end): - pass + if salting_peaks["salt_number"][0] != 0: + raise ValueError( + "Expected salt_number to start from 0 because " + f"{self.__class__.__name__} is a ExhaustPlugin plugin!" + ) + + n_events = len(salting_peaks) // 2 + if np.unique(salting_peaks["salt_number"]).size != n_events: + raise ValueError("Expected salt_number to be half of the input peaks number!") + + # combine salting_peaks and peaks + _peaks = np.empty(len(salting_peaks) + len(peaks), dtype=self._peaks_dtype) + for n in set(self.needed_fields): + _peaks[n] = np.hstack([salting_peaks[n], peaks[n]]) + _peaks = np.sort(_peaks, order="time") + # use S2s as anchors + anchor_peaks = salting_peaks[1::2] + windows = strax.touching_windows(_peaks, anchor_peaks, window=self.window) + + # check if the salting event can trigger + # for now, only S2 can trigger + if not self.exclude_s1_as_triggering_peaks: + raise NotImplementedError("Only S2 can trigger for now!") + result = np.empty(n_events, self.dtype) + if np.unique(anchor_peaks["type"]).size != 1: + raise ValueError("Expected only one type of anchor peaks!") + + # prepare for an empty event + empty_event = np.empty(1, dtype=self.dtype) + EventBasics.set_nan_defaults(empty_event) + + # iterate through all anchor peaks + _result = [] + for i in tqdm(range(n_events), disable=self.disable_tqdm): + left_i, right_i = windows[i] + _events = super().compute(_peaks[left_i:right_i], start, end) + _events = strax.split_touching_windows(_events, anchor_peaks[i : i + 1]) + if result["is_triggering"][i]: + if len(_events) != 1 or len(_events[0]) != 1: + raise ValueError(f"Expected 1 event, got {_events}!") + _result.append(_events[0]) + else: + empty_event["time"] = anchor_peaks["time"][i] + empty_event["endtime"] = anchor_peaks["endtime"][i] + _result.append(empty_event) + _result = np.hstack(_result) + + for n in _result.dtype.names: + result[n] = _result[n] + + # assign the most important parameters + result["is_triggering"] = self._is_triggering(anchor_peaks) + result["salt_number"] = salting_peaks["salt_number"][::2] + result["event_number"] = salting_peaks["salt_number"][::2] + return result class SaltedEventBasics(Plugin): diff --git a/axidence/plugins/salting/salting_events.py b/axidence/plugins/salting/salting_events.py index 00cad42..9584d92 100644 --- a/axidence/plugins/salting/salting_events.py +++ b/axidence/plugins/salting/salting_events.py @@ -100,6 +100,13 @@ class SaltingEvents(EventPositions, EventBasics): help="Will assign the events's ``s1_n_hits`` and ``s1_tight_coincidence`` by this", ) + n_drift_time_window = straxen.URLConfig( + default=5, + type=int, + track=True, + help="How many max drift time will the event builder extend", + ) + def _set_posrec_save(self): self.posrec_save = [] self.pos_rec_labels = [] @@ -154,6 +161,9 @@ def sample_time(self): """Sample the time according to the start and end of the run.""" self.event_time_interval = units.s // self.salting_rate + if units.s / self.salting_rate < self.drift_time_max * self.n_drift_time_window * 2: + raise ValueError("Salting rate is too high according the drift time window!") + time = np.arange( self.run_start + self.veto_length_run_start, self.run_end - self.veto_length_run_end, @@ -184,7 +194,8 @@ def set_chunk_splitting(self): def setup(self): """Sample the features of events.""" - super().setup() + super(EventPositions, self).setup() + super(SaltingEvents, self).setup() self.init_rng() self.init_run_meta() @@ -195,8 +206,6 @@ def setup(self): self.salting_events["salt_number"] = np.arange(self.n_events) self.salting_events["time"] = time self.salting_events["endtime"] = time - self.salting_events["s1_center_time"] = time - self.salting_events["s2_center_time"] = time self.salting_events["s1_n_hits"] = self.s1_n_hits_tight_coincidence self.salting_events["s1_tight_coincidence"] = self.s1_n_hits_tight_coincidence @@ -206,21 +215,21 @@ def setup(self): self.salting_events["x"] = np.cos(theta) * r self.salting_events["y"] = np.sin(theta) * r self.salting_events["z"] = -self.rng.random(size=self.n_events) * straxen.tpc_z - self.salting_events["s2_x"], self.salting_events["s2_y"], self.salting_events["z_naive"] = ( - self.inverse_field_distortion( - self.salting_events["x"], - self.salting_events["y"], - self.salting_events["z"], - ) + s2_x, s2_y, z_naive = self.inverse_field_distortion( + self.salting_events["x"], + self.salting_events["y"], + self.salting_events["z"], ) + self.salting_events["s2_x"] = s2_x + self.salting_events["s2_y"] = s2_y + self.salting_events["z_naive"] = z_naive self.salting_events["drift_time"] = ( - -( - self.salting_events["z_naive"] - - self.electron_drift_velocity * self.electron_drift_time_gate - ) - / self.electron_drift_velocity - + self.electron_drift_time_gate - ) + self.electron_drift_velocity * self.electron_drift_time_gate + - self.salting_events["z_naive"] + ) / self.electron_drift_velocity + + self.salting_events["s1_center_time"] = time - self.salting_events["drift_time"] + self.salting_events["s2_center_time"] = time s1_area_range = (float(self.s1_area_range[0]), float(self.s1_area_range[1])) s2_area_range = (float(self.s2_area_range[0]), float(self.s2_area_range[1])) diff --git a/axidence/plugins/salting/salting_peaks.py b/axidence/plugins/salting/salting_peaks.py index e7a2b67..af0263e 100644 --- a/axidence/plugins/salting/salting_peaks.py +++ b/axidence/plugins/salting/salting_peaks.py @@ -49,8 +49,6 @@ def infer_dtype(self): def compute(self, salting_events): """Copy features of salting_events into salting_peaks.""" salting_peaks = np.empty(len(salting_events) * 2, dtype=self.dtype) - for n in "time endtime".split(): - salting_peaks[n] = np.repeat(salting_events[n], 2) for n in "center_time area".split(): salting_peaks[n] = np.vstack( [ @@ -58,6 +56,9 @@ def compute(self, salting_events): salting_events[f"s2_{n}"], ] ).T.flatten() + salting_peaks["time"] = salting_peaks["center_time"] + # add one to prevent error about non-positive length + salting_peaks["endtime"] = salting_peaks["time"] + 1 for n in "n_hits tight_coincidence".split(): salting_peaks[n] = np.vstack( [ From 0d0ad3681e4241aedcaeae080d935f3950317a1c Mon Sep 17 00:00:00 2001 From: dachengx Date: Wed, 24 Apr 2024 20:40:39 +0800 Subject: [PATCH 2/3] Add `SaltedEventBasics` --- axidence/plugins/salting/event_building.py | 104 ++++++++++++++++----- axidence/utils.py | 24 +++++ 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/axidence/plugins/salting/event_building.py b/axidence/plugins/salting/event_building.py index 911445c..8f22185 100644 --- a/axidence/plugins/salting/event_building.py +++ b/axidence/plugins/salting/event_building.py @@ -1,10 +1,10 @@ import numpy as np from tqdm import tqdm import strax -from strax import Plugin import straxen from straxen import Events, EventBasics +from axidence.utils import needed_dtype from axidence.plugin import ExhaustPlugin @@ -38,25 +38,9 @@ def __init__(self): def setup(self): super().setup() - - # intersection depends_on's dtype.names will be needed in event building - self.needed_fields = set.intersection( - *tuple( - set.union(*tuple(set(self.deps[d].dtype_for(d).names) for d in dk)) - for dk in self.dependencies_by_kind().values() - ) + self.needed_fields, self._peaks_dtype = needed_dtype( + self.deps, self.dependencies_by_kind, set.intersection ) - dk = list(self.dependencies_by_kind().values())[0] - dtype_reference = strax.merged_dtype([self.deps[d].dtype_for(d) for d in dk]) - self._peaks_dtype = [] - for x in dtype_reference: - if x[0][1] in self.needed_fields: - self._peaks_dtype.append(x) - if len(self._peaks_dtype) != len(self.needed_fields): - raise ValueError( - "Weird! Could not find all needed fields " - f"{self.needed_fields} in {dtype_reference}!" - ) self.window = self.n_drift_time_window * self.drift_time_max @@ -89,6 +73,8 @@ def compute(self, salting_peaks, peaks, start, end): anchor_peaks = salting_peaks[1::2] windows = strax.touching_windows(_peaks, anchor_peaks, window=self.window) + _is_triggering = self._is_triggering(anchor_peaks) + # check if the salting event can trigger # for now, only S2 can trigger if not self.exclude_s1_as_triggering_peaks: @@ -107,31 +93,99 @@ def compute(self, salting_peaks, peaks, start, end): left_i, right_i = windows[i] _events = super().compute(_peaks[left_i:right_i], start, end) _events = strax.split_touching_windows(_events, anchor_peaks[i : i + 1]) - if result["is_triggering"][i]: + if _is_triggering[i]: if len(_events) != 1 or len(_events[0]) != 1: raise ValueError(f"Expected 1 event, got {_events}!") _result.append(_events[0]) else: empty_event["time"] = anchor_peaks["time"][i] empty_event["endtime"] = anchor_peaks["endtime"][i] - _result.append(empty_event) + _result.append(empty_event.copy()) _result = np.hstack(_result) for n in _result.dtype.names: result[n] = _result[n] # assign the most important parameters - result["is_triggering"] = self._is_triggering(anchor_peaks) + result["is_triggering"] = _is_triggering result["salt_number"] = salting_peaks["salt_number"][::2] result["event_number"] = salting_peaks["salt_number"][::2] + + if np.any(np.diff(result["time"]) < 0): + raise ValueError("Expected time to be sorted!") return result -class SaltedEventBasics(Plugin): +class SaltedEventBasics(EventBasics, ExhaustPlugin): __version__ = "0.0.0" - depends_on = ("events", "salting_peaks", "peak_basics", "peak_positions") + depends_on = ( + "events", + "salting_peaks", + "salting_peak_proximity", + "peak_basics", + "peak_proximity", + "peak_positions", + ) provides = "event_basics" data_kind = "events" save_when = strax.SaveWhen.EXPLICIT - dtype = strax.time_fields + def infer_dtype(self): + dtype = super().infer_dtype() + dtype += [ + (("Salting number of main S1", "s1_salt_number"), np.int64), + (("Salting number of main S2", "s2_salt_number"), np.int64), + (("Salting number of alternative S1", "alt_s1_salt_number"), np.int64), + (("Salting number of alternative S2", "alt_s2_salt_number"), np.int64), + (("Salting number of events", "salt_number"), np.int64), + (("Whether the salting event can trigger", "is_triggering"), bool), + ] + return dtype + + def setup(self): + super().setup() + + self.needed_fields, self._peaks_dtype = needed_dtype( + self.deps, self.dependencies_by_kind, set.union + ) + + self.peak_properties = tuple( + list(self.peak_properties) + [("salt_number", np.int64, "Salting number of peaks")] + ) + + def pick_fields(self, field, peaks): + if field in peaks.dtype.names: + _field = peaks[field] + else: + if np.issubdtype(np.dtype(self._peaks_dtype)[field], np.integer): + _field = np.full(len(peaks), -1) + else: + _field = np.full(len(peaks), np.nan) + return _field + + def compute(self, events, salting_peaks, peaks): + if salting_peaks["salt_number"][0] != 0: + raise ValueError( + "Expected salt_number to start from 0 because " + f"{self.__class__.__name__} is a ExhaustPlugin plugin!" + ) + + # combine salting_peaks and peaks + _peaks = np.empty(len(salting_peaks) + len(peaks), dtype=self._peaks_dtype) + for n in set(self.needed_fields): + _peaks[n] = np.hstack([self.pick_fields(n, salting_peaks), self.pick_fields(n, peaks)]) + _peaks = np.sort(_peaks, order="time") + + result = np.zeros(len(events), dtype=self.dtype) + self.set_nan_defaults(result) + + split_peaks = strax.split_by_containment(_peaks, events) + + result["time"] = events["time"] + result["endtime"] = events["endtime"] + result["salt_number"] = events["salt_number"] + result["event_number"] = events["event_number"] + + self.fill_events(result, events, split_peaks) + result["is_triggering"] = events["is_triggering"] + return result diff --git a/axidence/utils.py b/axidence/utils.py index 0d0af60..7d11fe5 100644 --- a/axidence/utils.py +++ b/axidence/utils.py @@ -19,3 +19,27 @@ def copy_dtype(dtype_reference, required_names): if not found: raise ValueError(f"Could not find {n} in dtype_reference!") return dtype + + +def needed_dtype(deps, dependencies_by_kind, func): + # intersection depends_on's dtype.names will be needed in event building + needed_fields = func( + *tuple( + set.union(*tuple(set(deps[d].dtype_for(d).names) for d in dk)) + for dk in dependencies_by_kind().values() + ) + ) + dtype_reference = list( + func( + *tuple( + set.union(*tuple(set(deps[d].dtype_for(d).descr) for d in dk)) + for dk in dependencies_by_kind().values() + ) + ) + ) + _peaks_dtype = copy_dtype(dtype_reference, needed_fields) + if len(_peaks_dtype) != len(needed_fields): + raise ValueError( + f"Weird! Could not find all needed fields {needed_fields} in {dtype_reference}!" + ) + return needed_fields, _peaks_dtype From 144de091c3ffb9108153b6d4639050301bfe5afd Mon Sep 17 00:00:00 2001 From: dachengx Date: Wed, 24 Apr 2024 20:42:01 +0800 Subject: [PATCH 3/3] Add sanity check of `SaltedEventBasics` --- axidence/plugins/salting/event_building.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/axidence/plugins/salting/event_building.py b/axidence/plugins/salting/event_building.py index 8f22185..fee0b29 100644 --- a/axidence/plugins/salting/event_building.py +++ b/axidence/plugins/salting/event_building.py @@ -188,4 +188,7 @@ def compute(self, events, salting_peaks, peaks): self.fill_events(result, events, split_peaks) result["is_triggering"] = events["is_triggering"] + + if np.all(result["s1_salt_number"] < 0) or np.all(result["s2_salt_number"] < 0): + raise ValueError("Found zero triggered salting peaks!") return result