modules.py

"""
Custom km3pipe modules for making nn input files.
"""

import numpy as np
import km3pipe as kp
import km3modules as km
import orcasong.plotting.plot_binstats as plot_binstats

__author__ = 'Stefan Reck'


class McInfoMaker(kp.Module):
    """
    Store mc info as float64 in the blob.

    Attributes
    ----------
    mc_info_extr : function
        Function to extract the info. Takes the blob as input, outputs
        a dict with the desired mc_infos.
    store_as : str
        Store the mcinfo with this name in the blob.

    """

    def configure(self):
        self.mc_info_extr = self.require('mc_info_extr')
        self.store_as = self.require('store_as')
        self.to_float64 = self.get("to_float64", default=True)

    def process(self, blob):
        track = self.mc_info_extr(blob)
        if self.to_float64:
            dtypes = []
            for key, v in track.items():
                if key in ("group_id", "event_id"):
                    dtypes.append((key, type(v)))
                else:
                    dtypes.append((key, np.float64))
        else:
            dtypes = None
        kp_hist = kp.dataclasses.Table(
            track, dtype=dtypes,  h5loc='y', name='event_info')
        if len(kp_hist) != 1:
            self.log.warning(
                "Warning: Extracted mc_info should have len 1, "
                "but it has len {}".format(len(kp_hist))
            )
        blob[self.store_as] = kp_hist
        return blob


class TimePreproc(kp.Module):
    """
    Preprocess the time in the blob in various ways.

    Attributes
    ----------
    add_t0 : bool
        If true, t0 will be added to times of hits and mchits.
    center_time : bool
        If true, center hit and mchit times with the time of the first
        triggered hit.

    """

    def configure(self):
        self.add_t0 = self.get('add_t0', default=False)
        self.center_time = self.get('center_time', default=True)

        self._has_mchits = None
        self._print_flags = set()

    def process(self, blob):
        if self._has_mchits is None:
            self._has_mchits = "McHits" in blob
        if self.add_t0:
            blob = self.add_t0_time(blob)
        if self.center_time:
            blob = self.center_hittime(blob)

        return blob

    def add_t0_time(self, blob):
        self._print_once("Adding t0 to hit times")
        blob["Hits"].time = np.add(blob["Hits"].time, blob["Hits"].t0)
        if self._has_mchits:
            self._print_once("Adding t0 to mchit times")
            blob["McHits"].time = np.add(
                blob["McHits"].time, blob["McHits"].t0)

        return blob

    def center_hittime(self, blob):
        hits_time = blob["Hits"].time
        hits_triggered = blob["Hits"].triggered
        t_first_trigger = np.min(hits_time[hits_triggered != 0])

        self._print_once("Centering time of Hits with first triggered hit")
        blob["Hits"].time = np.subtract(hits_time, t_first_trigger)

        if self._has_mchits:
            self._print_once("Centering time of McHits with first triggered hit")
            mchits_time = blob["McHits"].time
            blob["McHits"].time = np.subtract(mchits_time, t_first_trigger)

        return blob

    def _print_once(self, text):
        if text not in self._print_flags:
            self._print_flags.add(text)
            self.cprint(text)


class ImageMaker(kp.Module):
    """
    Make a n-d histogram from "Hits", and store it in the blob as 'samples'.

    Attributes
    ----------
    bin_edges_list : List
        List with the names of the fields to bin, and the respective bin edges,
        including the left- and right-most bin edge.
    hit_weights : str, optional
        Use blob["Hits"][hit_weights] as weights for samples in histogram.

    """

    def configure(self):
        self.bin_edges_list = self.require('bin_edges_list')
        self.hit_weights = self.get('hit_weights')
        self.store_as = "samples"

    def process(self, blob):
        data, bins, name = [], [], ""

        for bin_name, bin_edges in self.bin_edges_list:
            data.append(blob["Hits"][bin_name])
            bins.append(bin_edges)
            name += bin_name + "_"

        if self.hit_weights is not None:
            weights = blob["Hits"][self.hit_weights]
        else:
            weights = None

        histogram = np.histogramdd(data, bins=bins, weights=weights)[0]

        hist_one_event = histogram[np.newaxis, ...].astype(np.uint8)
        kp_hist = kp.dataclasses.NDArray(
            hist_one_event, h5loc='x', title=name + "event_images")

        blob[self.store_as] = kp_hist
        return blob


class BinningStatsMaker(kp.Module):
    """
    Generate a histogram of the number of hits for each binning field name.

    E.g. if the bin_edges_list contains "pos_z", this will make a histogram
    of #Hits vs. "pos_z", together with how many hits were outside
    of the bin edges in both directions.

    Per default, the resolution of the histogram (width of bins) will be
    higher then the given bin edges, and the edges will be stored seperatly.
    The time is the exception: The plotted bins have exactly the
    given bin edges.

    Attributes
    ----------
    bin_edges_list : List
        List with the names of the fields to bin, and the respective bin edges,
        including the left- and right-most bin edge.
    res_increase : int
        Increase the number of bins by this much in the hists (so that one
        can see if the edges have been placed correctly). Is never used
        for the time binning (field name "time").
    bin_plot_freq : int
        Extract data for the histograms only every given number of blobs
        (reduces time the pipeline takes to complete).

    """

    def configure(self):
        self.bin_edges_list = self.require('bin_edges_list')
        self.res_increase = self.get("res_increase", default=5)
        self.bin_plot_freq = 1

        self.hists = {}
        for bin_name, org_bin_edges in self.bin_edges_list:
            # dont space bin edges for time
            if bin_name == "time":
                bin_edges = org_bin_edges
            else:
                bin_edges = self._space_bin_edges(org_bin_edges)

            self.hists[bin_name] = {
                "hist": np.zeros(len(bin_edges) - 1),
                "hist_bin_edges": bin_edges,
                "bin_edges": org_bin_edges,
                # below smallest edge, above largest edge:
                "cut_off": np.zeros(2),
            }

        self.i = 0

    def _space_bin_edges(self, bin_edges):
        """
        Increase resolution of given binning.
        """
        increased_n_bins = (len(bin_edges) - 1) * self.res_increase + 1
        bin_edges = np.linspace(
            bin_edges[0], bin_edges[-1], increased_n_bins)

        return bin_edges

    def process(self, blob):
        """
        Extract data from blob for the hist plots.
        """
        if self.i % self.bin_plot_freq == 0:
            for bin_name, hists_data in self.hists.items():
                hist_bin_edges = hists_data["hist_bin_edges"]

                hits = blob["Hits"]
                data = hits[bin_name]
                # get how much is cut off due to these limits
                out_pos = data[data > np.max(hist_bin_edges)].size
                out_neg = data[data < np.min(hist_bin_edges)].size

                # get all hits which are not cut off by other bin edges
                data = hits[bin_name][self._is_in_limits(
                    hits, excluded=bin_name)]
                hist = np.histogram(data, bins=hist_bin_edges)[0]

                self.hists[bin_name]["hist"] += hist
                self.hists[bin_name]["cut_off"] += np.array([out_neg, out_pos])

        self.i += 1
        return blob

    def finish(self):
        """
        Append the hists, which are the stats of the binning.

        Its a dict with each binning field name containing the following
        ndarrays:

        bin_edges : The actual bin edges.
        cut_off : How many events were cut off in positive and negative
            direction due to this binning.
        hist_bin_edges : The bin edges for the plot in finer resolution then
            the actual bin edges.
        hist : The number of hist in each bin of the hist_bin_edges.

        """
        return self.hists

    def _is_in_limits(self, hits, excluded=None):
        """ Get which hits are in the limits defined by ALL bin edges
        (except for given one). """
        inside = None
        for dfield, edges in self.bin_edges_list:
            if dfield == excluded:
                continue
            is_in = np.logical_and(hits[dfield] >= min(edges),
                                   hits[dfield] <= max(edges))
            if inside is None:
                inside = is_in
            else:
                inside = np.logical_and(inside, is_in)
        return inside


class PointMaker(kp.Module):
    """
    Store individual hit info from "Hits" in the blob as 'samples'.

    Used for graph networks.

    Attributes
    ----------
    max_n_hits : int
        Maximum number of hits that gets saved per event. If an event has
        more, some will get cut!
    time_window : tuple, optional
        Two ints (start, end). Hits outside of this time window will be cut
        away (base on 'Hits/time').
        Default: Keep all hits.
    hit_infos : tuple, optional
        Which entries in the '/Hits' Table will be kept. E.g. pos_x, time, ...
        Default: Keep all entries.
    dset_n_hits : str, optional
        If given, store the number of hits that are in the time window
        as a new column called 'n_hits_intime' in the dataset with
        this name (usually this is EventInfo).

    """
    def configure(self):
        self.max_n_hits = self.require("max_n_hits")
        self.hit_infos = self.get("hit_infos", default=None)
        self.time_window = self.get("time_window", default=None)
        self.dset_n_hits = self.get("dset_n_hits", default=None)
        self.store_as = "samples"

    def process(self, blob):
        if self.hit_infos is None:
            self.hit_infos = blob["Hits"].dtype.names
        points, n_hits = self.get_points(blob)
        blob[self.store_as] = kp.NDArray(
            np.expand_dims(points, 0), h5loc="x", title="nodes")
        if self.dset_n_hits:
            blob[self.dset_n_hits] = blob[self.dset_n_hits].append_columns(
                "n_hits_intime", n_hits)
        return blob

    def get_points(self, blob):
        """
        Get the desired hit infos from the blob.

        Returns
        -------
        points : np.array
            The hit infos of this event as a 2d matrix. No of rows are
            fixed to the given max_n_hits. Each of the self.extract_keys,
            is in one column + an additional column which is 1 for
            actual hits, and 0 for if its a padded row.
        n_hits : int
            Number of hits in the given time window.

        """
        points = np.zeros(
            (self.max_n_hits, len(self.hit_infos) + 1), dtype="float32")

        hits = blob["Hits"]
        if self.time_window is not None:
            # remove hits outside of time window
            hits = hits[np.logical_and(
                hits["time"] >= self.time_window[0],
                hits["time"] <= self.time_window[1],
            )]

        n_hits = len(hits)
        if n_hits > self.max_n_hits:
            # if there are too many hits, take random ones, but keep order
            indices = np.arange(n_hits)
            np.random.shuffle(indices)
            which = indices[:self.max_n_hits]
            which.sort()
            hits = hits[which]

        for i, which in enumerate(self.hit_infos):
            data = hits[which]
            points[:n_hits, i] = data
        # last column is whether there was a hit or no
        points[:n_hits, -1] = 1.
        return points, n_hits

    def finish(self):
        return {"hit_infos": tuple(self.hit_infos) + ("is_valid", )}


class EventSkipper(kp.Module):
    """
    Skip events based on blob content.

    Attributes
    ----------
    event_skipper : callable
        Function that takes the blob as an input, and returns a bool.
        If the bool is true, the blob will be skipped.

    """

    def configure(self):
        self.event_skipper = self.require('event_skipper')
        self._not_skipped = 0
        self._skipped = 0

    def process(self, blob):
        blob = self._remove_groupid(blob)
        if self.event_skipper(blob):
            self._skipped += 1
            return
        else:
            self._not_skipped += 1
            return blob

    def _remove_groupid(self, blob):
        """
        Workaround until bug https://git.km3net.de/km3py/km3pipe/-/issues/203
        in km3pipe is fixed: Drop all group_ids
        """
        if "GroupInfo" in blob:
            del blob["GroupInfo"]
        for key in blob.keys():
            try:
                blob[key] = blob[key].drop_columns("group_id")
            except AttributeError:
                continue
        return blob

    def finish(self):
        tot_events = self._skipped + self._not_skipped
        self.cprint(
            f"Skipped {self._skipped}/{tot_events} events "
            f"({self._skipped/tot_events:.4%})."
        )


class DetApplier(kp.Module):
    """
    Apply calibration to the Hits and McHits with a detx file.

    Attributes
    ----------
    det_file : str
        Path to a .detx detector geometry file.

    """

    def configure(self):
        self.det_file = self.require("det_file")
        self.cprint(f"Calibrating with {self.det_file}")
        self.calib = kp.calib.Calibration(filename=self.det_file)
        self._calib_checked = False

    def process(self, blob):
        if self._calib_checked is False:
            if "pos_x" in blob["Hits"]:
                self.log.warn(
                    "Warning: Using a det file, but pos_x in Hits detected. "
                    "Is the file already calibrated? This might lead to "
                    "errors with t0."
                )
            self._calib_checked = True
        blob["Hits"] = self.calib.apply(blob["Hits"], correct_slewing=True)
        if "McHits" in blob:
            blob["McHits"] = self.calib.apply(blob["McHits"])
            # TODO remove once https://git.km3net.de/km3py/km3pipe/-/issues/239 is solved
            self.subtract_t0_mctime(blob)

        return blob

    def subtract_t0_mctime(self, blob):
        blob["McHits"].time = np.subtract(
            blob["McHits"].time, blob["McHits"].t0)
        return blob


class HitRotator(kp.Module):
    """
    Rotates hits by angle theta.

    Attributes
    ----------
    theta : float
        Angle by which hits are rotated (radian).

    """

    def configure(self):
        self.theta = self.require('theta')

    def process(self, blob):
        x = blob['Hits']['x']
        y = blob['Hits']['y']

        rot_matrix = np.array([[np.cos(self.theta), - np.sin(self.theta)],
                               [np.sin(self.theta), np.cos(self.theta)]])

        x_rot = []
        y_rot = []

        for i in range(0, len(x)):
            vec = np.array([[x[i]], [y[i]]])
            rot = np.dot(rot_matrix, vec)
            x_rot.append(rot[0][0])
            y_rot.append(rot[1][0])

        blob['Hits']['x'] = x_rot
        blob['Hits']['y'] = y_rot

        return blob