import strax
import numpy as np
[docs]class EventwBayesClass(strax.Plugin):
"""Append at event level the posterior probability for an S1, S2, alt_S1 and alt_S2."""
provides = "event_w_bayes_class"
depends_on = ("peak_classification_bayes", "event_basics")
data_kind = "events"
__version__ = "0.0.1"
[docs] def infer_dtype(self):
dtype = []
dtype += strax.time_fields
for name in ["s1", "s2", "alt_s1", "alt_s2"]:
dtype += [(f"{name}_ln_prob_s1", np.float32, f"Given an {name}, s1 ln probability")]
dtype += [(f"{name}_ln_prob_s2", np.float32, f"Given an {name}, s2 ln probability")]
return dtype
[docs] def compute(self, peaks, events):
result = np.empty(len(events), dtype=self.dtype)
for name in ["s1", "s2", "alt_s1", "alt_s2"]:
result[f"{name}_ln_prob_s1"] = np.nan
result[f"{name}_ln_prob_s2"] = np.nan
# Select peaks based on their start time
mask = np.in1d(peaks["time"], events[f"{name}_time"])
mask_ev = np.in1d(events[f"{name}_time"], peaks["time"])
result[f"{name}_ln_prob_s1"][mask_ev] = peaks["ln_prob_s1"][mask]
result[f"{name}_ln_prob_s2"][mask_ev] = peaks["ln_prob_s2"][mask]
result["time"] = events["time"]
result["endtime"] = events["endtime"]
return result