Refactor nv plugins

straxen/plugins/events_mv/events_mv.py

@@ -10,7 +10,7 @@
 class muVETOEvents(nVETOEvents):
     """Plugin which computes the boundaries of veto events."""
 
-    depends_on = "hitlets_mv"
+    depends_on = "hitlets_mv"  # type: ignore
     provides = "events_mv"
     data_kind = "events_mv"
 

straxen/plugins/events_nv/__init__.py

@@ -3,3 +3,6 @@
 
 from . import events_nv
 from .events_nv import *
+
+from . import event_waveform_nv
+from .event_waveform_nv import *

straxen/plugins/events_nv/event_waveform_nv.py

@@ -0,0 +1,112 @@
+import strax
+import straxen
+import numpy as np
+from immutabledict import immutabledict
+
+export, __all__ = strax.exporter()
+
+
+@export
+class nVETOEventWaveform(strax.Plugin):
+    """Plugin which computes the summed waveform as well as some shape properties of the NV
+    events."""
+
+    __version__ = "0.0.1"
+
+    depends_on = "events_nv", "records_nv"
+    provides = "event_waveform_nv"
+    data_kind = "events_nv"
+    compressor = "zstd"
+
+    gain_model_nv = straxen.URLConfig(
+        default="cmt://to_pe_model_nv?version=ONLINE&run_id=plugin.run_id",
+        infer_type=False,
+        help="PMT gain model. Specify as (model_type, model_config, nT = True)",
+    )
+
+    channel_map = straxen.URLConfig(
+        track=False,
+        type=immutabledict,
+        help="immutabledict mapping subdetector to (min, max) channel number",
+    )
+
+    def infer_dtype(self):
+        return veto_event_waveform_dtype()
+
+    def setup(self):
+        self.channel_range = self.channel_map["nveto"]
+        self.n_channel = (self.channel_range[1] - self.channel_range[0]) + 1
+
+        to_pe = self.gain_model_nv
+
+        # Create to_pe array of size max channel:
+        self.to_pe = np.zeros(self.channel_range[1] + 1, dtype=np.float32)
+        self.to_pe[self.channel_range[0] :] = to_pe[:]
+
+    def compute(self, events_nv, records_nv, start, end):
+        events = events_nv
+        events_waveform = np.zeros(len(events), self.dtype)
+
+        # Compute shape like properties:
+        _tmp_events = np.zeros(len(events_nv), dtype=_temp_event_data_type())
+        strax.copy_to_buffer(events_nv, _tmp_events, "_temp_nv_evts_wf_cpy")
+        _tmp_events["length"] = (events_nv["endtime"] - events_nv["time"]) // 2
+        _tmp_events["dt"] = 2
+        print(records_nv)
+        strax.simple_summed_waveform(records_nv, _tmp_events, self.to_pe)
+        strax.compute_widths(_tmp_events)
+
+        strax.copy_to_buffer(_tmp_events, events_waveform, "_temp_nv_evts_cpy")
+        events_waveform["range_50p_area"] = _tmp_events["width"][:, 5]
+        events_waveform["range_90p_area"] = _tmp_events["width"][:, 9]
+        events_waveform["rise_time"] = -_tmp_events["area_decile_from_midpoint"][:, 1]
+        del _tmp_events
+
+        return events_waveform
+
+
+def veto_event_waveform_dtype(
+    n_samples_wf: int = 200,
+) -> list:
+    dtype = []
+    dtype += strax.time_dt_fields  # because mutable
+    dtype += [
+        (("Waveform data in PE/sample (not PE/ns!)", "data"), np.float32, n_samples_wf),
+        (("Width (in ns) of the central 50% area of the peak", "range_50p_area"), np.float32),
+        (("Width (in ns) of the central 90% area of the peak", "range_90p_area"), np.float32),
+        (("Time between 10% and 50% area quantiles [ns]", "rise_time"), np.float32),
+    ]
+    return dtype
+
+
+def _temp_event_data_type(
+    n_samples_wf: int = 150,
+    n_pmts: int = 120,
+    n_widths: int = 11,
+) -> list:
+    """Temp.
+
+    data type which adds field required to use some of the functions used to compute the shape of
+    the summed waveform for the TPC.
+
+    """
+    dtype = veto_event_waveform_dtype()
+    dtype += [
+        (
+            ("Dummy, total area of all hitlets in event [pe]", "area"),
+            np.float32,
+        ),
+        (
+            ("Dummy top waveform data in PE/sample (not PE/ns!)", "data_top"),
+            np.float32,
+            n_samples_wf,
+        ),
+        (("Peak widths in range of central area fraction [ns]", "width"), np.float32, n_widths),
+        (
+            ("Peak widths: time between nth and 5th area decile [ns]", "area_decile_from_midpoint"),
+            np.float32,
+            n_widths,
+        ),
+    ]
+
+    return dtype

straxen/plugins/events_nv/events_nv.py

@@ -12,7 +12,7 @@
 class nVETOEvents(strax.OverlapWindowPlugin):
     """Plugin which computes the boundaries of veto events."""
 
-    __version__ = "0.0.3"
+    __version__ = "0.1.0"
 
     depends_on = "hitlets_nv"
     provides = "events_nv"
@@ -60,6 +60,7 @@ def compute(self, hitlets_nv, start, end):
             n_channel=self.n_channel,
         )
 
+        # Compute basic properties:
         if len(hitlets_ids_in_event):
             compute_nveto_event_properties(
                 events, hitlets_nv, hitlets_ids_in_event, start_channel=self.channel_range[0]
@@ -69,16 +70,13 @@ def compute(self, hitlets_nv, start, end):
         n_events = len(events)
         events[self.name_event_number] = np.arange(n_events) + self.events_seen
         self.events_seen += n_events
-
-        # Don't extend beyond the chunk boundaries
-        # This will often happen for events near the invalid boundary of the
-        # overlap processing (which should be thrown away)
-        events["time"] = np.clip(events["time"], start, end)
-        events["endtime"] = np.clip(events["endtime"], start, end)
         return events
 
 
-def veto_event_dtype(name_event_number: str = "event_number_nv", n_pmts: int = 120) -> list:
+def veto_event_dtype(
+    name_event_number: str = "event_number_nv",
+    n_pmts: int = 120,
+) -> list:
     dtype = []
     dtype += strax.time_fields  # because mutable
     dtype += [
@@ -95,6 +93,14 @@ def veto_event_dtype(name_event_number: str = "event_number_nv", n_pmts: int = 1
             np.float32,
         ),
         (("Weighted variance of time [ns]", "center_time_spread"), np.float32),
+        (
+            ("Minimal amplitude-to-amplitude gap between neighboring hitlets [ns]", "min_gap"),
+            np.int8,
+        ),
+        (
+            ("Maximal amplitude-to-amplitude gap between neighboring hitlets [ns]", "max_gap"),
+            np.int8,
+        ),
     ]
     return dtype
 
@@ -116,17 +122,22 @@ def compute_nveto_event_properties(
 
     """
     for e, (s_i, e_i) in zip(events, contained_hitlets_ids):
-        hitlet = hitlets[s_i:e_i]
-        event_area = np.sum(hitlet["area"])
+        hitlets_in_event = hitlets[s_i:e_i]
+        event_area = np.sum(hitlets_in_event["area"])
         e["area"] = event_area
-        e["n_hits"] = len(hitlet)
-        e["n_contributing_pmt"] = len(np.unique(hitlet["channel"]))
+        e["n_hits"] = len(hitlets_in_event)
+        e["n_contributing_pmt"] = len(np.unique(hitlets_in_event["channel"]))
+
+        t = hitlets_in_event["time"] - hitlets_in_event[0]["time"]
+
+        dt = np.diff(hitlets_in_event["time"] + hitlets_in_event["time_amplitude"])
+        e["min_gap"] = np.min(dt)
+        e["max_gap"] = np.max(dt)
 
-        t = hitlet["time"] - hitlet[0]["time"]
         if event_area:
-            e["center_time"] = np.sum(t * hitlet["area"]) / event_area
+            e["center_time"] = np.sum(t * hitlets_in_event["area"]) / event_area
             if e["n_hits"] > 1 and e["center_time"]:
-                w = hitlet["area"] / e["area"]  # normalized weights
+                w = hitlets_in_event["area"] / e["area"]  # normalized weights
                 # Definition of variance
                 e["center_time_spread"] = np.sqrt(
                     np.sum(w * np.power(t - e["center_time"], 2)) / np.sum(w)
@@ -135,7 +146,7 @@ def compute_nveto_event_properties(
                 e["center_time_spread"] = np.inf
 
         # Compute per channel properties:
-        for hit in hitlet:
+        for hit in hitlets_in_event:
             ch = hit["channel"] - start_channel
             e["area_per_channel"][ch] += hit["area"]
 
@@ -226,4 +237,5 @@ def _make_event(hitlets: np.ndarray, hitlet_ids: np.ndarray, res: np.ndarray):
     for ei, ids in enumerate(hitlet_ids):
         hit = hitlets[ids[0] : ids[1]]
         res[ei]["time"] = hit[0]["time"]
-        res[ei]["endtime"] = np.max(strax.endtime(hit))
+        endtime = np.max(strax.endtime(hit))
+        res[ei]["endtime"] = endtime

straxen/plugins/hitlets_mv/hitlets_mv.py

@@ -65,30 +65,6 @@ class muVETOHitlets(nVETOHitlets):
         ),
     )
 
-    entropy_template_mv = straxen.URLConfig(
-        default="flat",
-        track=True,
-        infer_type=False,
-        child_option=True,
-        parent_option_name="entropy_template_nv",
-        help=(
-            'Template data is compared with in conditional entropy. Can be either "flat" or a '
-            "template array."
-        ),
-    )
-
-    entropy_square_data_mv = straxen.URLConfig(
-        default=False,
-        track=True,
-        infer_type=False,
-        child_option=True,
-        parent_option_name="entropy_square_data_nv",
-        help=(
-            "Parameter which decides if data is first squared before normalized and compared to "
-            "the template."
-        ),
-    )
-
     gain_model_mv = straxen.URLConfig(
         default="cmt://to_pe_model_mv?version=ONLINE&run_id=plugin.run_id",
         infer_type=False,