Add event building: SaltedEvents and SaltedEventBasics

axidence/plugins/salting/event_building.py

@@ -1,31 +1,194 @@
 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.utils import needed_dtype
+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()
+        self.needed_fields, self._peaks_dtype = needed_dtype(
+            self.deps, self.dependencies_by_kind, set.intersection
+        )
+
+        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)
+
+        _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:
+            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 _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.copy())
+        _result = np.hstack(_result)
+
+        for n in _result.dtype.names:
+            result[n] = _result[n]
 
+        # assign the most important parameters
+        result["is_triggering"] = _is_triggering
+        result["salt_number"] = salting_peaks["salt_number"][::2]
+        result["event_number"] = salting_peaks["salt_number"][::2]
 
-class SaltedEventBasics(Plugin):
+        if np.any(np.diff(result["time"]) < 0):
+            raise ValueError("Expected time to be sorted!")
+        return result
+
+
+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"]
+
+        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

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]))

axidence/plugins/salting/salting_peaks.py

@@ -49,15 +49,16 @@ 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(
                 [
                     salting_events[f"s1_{n}"],
                     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(
                 [

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