events.py

import strax
import numpy as np
import straxen

export, __all__ = strax.exporter()


@export
class Events(strax.OverlapWindowPlugin):
    """Plugin which defines an "event" in our TPC.

    An event is defined by peak(s) in fixed range of time around a peak
    which satisfies certain conditions:
        1. The triggering peak must have a certain area.
        2. The triggering peak must have less than
           "trigger_max_competing" peaks. (A competing peak must have a
           certain area fraction of the triggering peak and must be in a
           window close to the main peak)

    Note:
        The time range which defines an event gets chopped at the chunk
        boundaries. This happens at invalid boundaries of the

    """

    __version__ = "0.1.1"

    depends_on = ("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"),
    ]

    events_seen = 0

    electron_drift_velocity = straxen.URLConfig(
        default="cmt://electron_drift_velocity?version=ONLINE&run_id=plugin.run_id",
        cache=True,
        help="Vertical electron drift velocity in cm/ns (1e4 m/ms)",
    )

    trigger_min_area = straxen.URLConfig(
        default=100,
        type=(int, float),
        help="Peaks must have more area (PE) than this to cause events",
    )

    trigger_max_competing = straxen.URLConfig(
        default=7,
        type=int,
        help="Peaks must have FEWER nearby larger or slightly smaller peaks to cause events",
    )

    left_event_extension = straxen.URLConfig(
        default=int(0.25e6),
        type=(int, float),
        help=(
            "Extend events this many ns to the left from each "
            "triggering peak. This extension is added to the maximum "
            "drift time."
        ),
    )

    right_event_extension = straxen.URLConfig(
        default=int(0.25e6),
        type=(int, float),
        help="Extend events this many ns to the right from each triggering peak.",
    )

    max_drift_length = straxen.URLConfig(
        default=straxen.tpc_z,
        type=(int, float),
        help="Total length of the TPC from the bottom of gate to the top of cathode wires [cm]",
    )

    exclude_s1_as_triggering_peaks = straxen.URLConfig(
        default=True,
        type=bool,
        help="If true exclude S1s as triggering peaks.",
    )

    event_s1_min_coincidence = straxen.URLConfig(
        default=2,
        infer_type=False,
        help=(
            "Event level S1 min coincidence. Should be >= "
            "s1_min_coincidence in the peaklet classification"
        ),
    )

    s1_min_coincidence = straxen.URLConfig(
        default=2, type=int, help="Minimum tight coincidence necessary to make an S1"
    )

    diagnose_overlapping = straxen.URLConfig(
        track=False, default=True, infer_type=False, help="Enable runtime checks for disjointness"
    )

    def setup(self):
        if self.s1_min_coincidence > self.event_s1_min_coincidence:
            raise ValueError(
                "Peak s1 coincidence requirement should be smaller "
                "or equal to event_s1_min_coincidence"
            )
        self.drift_time_max = int(self.max_drift_length / self.electron_drift_velocity)
        # Left_extension and right_extension should be computed in setup to be
        # reflected in cutax too.
        self.left_extension = self.left_event_extension + self.drift_time_max
        self.right_extension = self.right_event_extension

    def get_window_size(self):
        # Take a large window for safety, events can have long tails
        return 10 * (self.left_event_extension + self.drift_time_max + self.right_event_extension)

    def _is_triggering(self, peaks):
        _is_triggering = peaks["area"] > self.trigger_min_area
        _is_triggering &= peaks["n_competing"] <= self.trigger_max_competing
        if self.exclude_s1_as_triggering_peaks:
            _is_triggering &= peaks["type"] == 2
        else:
            is_not_s1 = peaks["type"] != 1
            has_tc_large_enough = peaks["tight_coincidence"] >= self.event_s1_min_coincidence
            _is_triggering &= is_not_s1 | has_tc_large_enough
        return _is_triggering

    def compute(self, peaks, start, end):
        _is_triggering = self._is_triggering(peaks)

        triggers = peaks[_is_triggering]

        # Join nearby triggers
        t0, t1 = strax.find_peak_groups(
            triggers,
            gap_threshold=self.left_extension + self.right_extension + 1,
            left_extension=self.left_extension,
            right_extension=self.right_extension,
        )

        # 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)
        t0 = np.clip(t0, start, end)
        t1 = np.clip(t1, start, end)

        result = np.zeros(len(t0), self.dtype)
        result["time"] = t0
        result["endtime"] = t1
        result["event_number"] = np.arange(len(result)) + self.events_seen

        if not result.size > 0:
            print("Found chunk without events?!")

        if self.diagnose_overlapping and len(result):
            # Check if the event windows overlap
            _event_window_do_not_overlap = (strax.endtime(result)[:-1] - result["time"][1:]) <= 0
            assert np.all(_event_window_do_not_overlap), "Events not disjoint"

        self.events_seen += len(result)

        return result