ztplot.py

#!/usr/bin/env python
# coding=utf-8
# vim: ts=4 sw=4 et
"""
Creates z-t-plots for every DU.

Usage:
    ztplot.py [options]
    ztplot.py (-h | --help)

Options:
    -l LIGIER_IP    The IP of the ligier [default: 127.0.0.1].
    -p LIGIER_PORT  The port of the ligier [default: 5553].
    -d DET_ID       Detector ID [default: 29].
    -o PLOT_DIR     The directory to save the plot [default: plots].
    -h --help       Show this screen.

"""
from __future__ import division

from datetime import datetime
import os
import queue
import shutil
import threading

import matplotlib
# Force matplotlib to not use any Xwindows backend.
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np

import km3pipe as kp
from km3pipe.io.daq import is_3dmuon, is_3dshower, is_mxshower
from km3modules.hits import count_multiplicities
import km3pipe.style
km3pipe.style.use('km3pipe')

from km3pipe.logger import logging

lock = threading.Lock()


class ZTPlot(kp.Module):
    def configure(self):
        self.plots_path = self.require('plots_path')
        self.ytick_distance = self.get('ytick_distance', default=200)
        self.min_dus = self.get('min_dus', default=1)
        self.min_doms = self.get('min_doms', default=4)
        self.det_id = self.require('det_id')
        self.t0set = None
        self.calib = None
        self.max_z = None

        self.sds = kp.db.StreamDS()

        self.index = 0

        self._update_calibration()

        self.run = True
        self.max_queue = 3
        self.queue = queue.Queue()
        self.thread = threading.Thread(target=self.plot, daemon=True)
        self.thread.start()

    def _update_calibration(self):
        self.print("Updating calibration")
        self.t0set = self.sds.t0sets(detid=self.det_id).iloc[-1]['CALIBSETID']
        self.calib = kp.calib.Calibration(det_id=self.det_id, t0set=self.t0set)
        self.max_z = round(np.max(self.calib.detector.pmts.pos_z) + 10, -1)

    def process(self, blob):
        if 'Hits' not in blob:
            return blob

        self.index += 1
        if self.index % 1000 == 0:
            self._update_calibration()

        hits = blob['Hits']
        hits = self.calib.apply(hits)
        event_info = blob['EventInfo']

        n_triggered_dus = len(np.unique(hits[hits.triggered == True].du))
        n_triggered_doms = len(np.unique(hits[hits.triggered == True].dom_id))
        if n_triggered_dus < self.min_dus or n_triggered_doms < self.min_doms:
            print(f"Skipping event with {n_triggered_dus} DUs "
                  f"and {n_triggered_doms} DOMs.")
            return blob

        print("OK")
        # print("Event queue size: {0}".format(self.queue.qsize()))
        if self.queue.qsize() < self.max_queue:
            self.queue.put((event_info, hits))

        return blob

    def plot(self):
        while self.run:
            try:
                event_info, hits = self.queue.get(timeout=50)
            except queue.Empty:
                continue
            with lock:
                self.create_plot(event_info, hits)

    def create_plot(self, event_info, hits):
        print(self.__class__.__name__ + ": updating plot.")
        dus = set(hits.du)
        doms = set(hits.dom_id)
        fontsize = 16

        hits = hits.append_columns('multiplicity',
                                   np.ones(len(hits))).sorted(by='time')

        for dom in doms:
            dom_hits = hits[hits.dom_id == dom]
            mltps, m_ids = count_multiplicities(dom_hits.time)
            hits['multiplicity'][hits.dom_id == dom] = mltps

        time_offset = np.min(hits[hits.triggered == True].time)
        hits.time -= time_offset

        n_plots = len(dus)
        n_cols = int(np.ceil(np.sqrt(n_plots)))
        n_rows = int(n_plots / n_cols) + (n_plots % n_cols > 0)
        marker_fig, marker_axes = plt.subplots()  # for the marker size hack...
        fig, axes = plt.subplots(ncols=n_cols,
                                 nrows=n_rows,
                                 sharex=True,
                                 sharey=True,
                                 figsize=(16, 8),
                                 constrained_layout=True)

        axes = [axes] if n_plots == 1 else axes.flatten()

        dom_zs = self.calib.detector.pmts.pos_z[
            (self.calib.detector.pmts.du == min(dus))
            & (self.calib.detector.pmts.channel_id == 0)]

        for ax, du in zip(axes, dus):
            for z in dom_zs:
                ax.axhline(z, lw=1, color='b', ls='--', alpha=0.15)
            du_hits = hits[hits.du == du]
            trig_hits = du_hits[du_hits.triggered == True]

            ax.scatter(du_hits.time,
                       du_hits.pos_z,
                       s=du_hits.multiplicity * 30,
                       c='#09A9DE',
                       label='hit',
                       alpha=0.5)
            ax.scatter(trig_hits.time,
                       trig_hits.pos_z,
                       s=trig_hits.multiplicity * 30,
                       alpha=0.8,
                       marker="+",
                       c='#FF6363',
                       label='triggered hit')
            ax.set_title('DU{0}'.format(int(du)),
                         fontsize=fontsize,
                         fontweight='bold')
            # The only way I could create a legend with matching marker sizes
            max_multiplicity = int(np.max(du_hits.multiplicity))
            custom_markers = [
                marker_axes.scatter(
                    [], [], s=mult * 30, color='#09A9DE', lw=0, alpha=0.5)
                for mult in range(0, max_multiplicity)
            ] + [marker_axes.scatter([], [], s=30, marker="+", c='#FF6363')]
            ax.legend(custom_markers, ['multiplicity'] +
                      ["       %d" % m
                       for m in range(1, max_multiplicity)] + ['triggered'],
                      scatterpoints=1,
                      markerscale=1,
                      loc='upper left',
                      frameon=True,
                      framealpha=0.7)

        for idx, ax in enumerate(axes):
            ax.set_ylim(0, self.max_z)
            ax.tick_params(labelsize=fontsize)
            ax.yaxis.set_major_locator(
                ticker.MultipleLocator(self.ytick_distance))
            xlabels = ax.get_xticklabels()
            for label in xlabels:
                label.set_rotation(45)

            if idx % n_cols == 0:
                ax.set_ylabel('z [m]', fontsize=fontsize)
            if idx >= len(axes) - n_cols:
                ax.set_xlabel('time [ns]', fontsize=fontsize)

        trigger_params = ' '.join([
            trig
            for trig, trig_check in (("MX", is_mxshower), ("3DM", is_3dmuon),
                                     ("3DS", is_3dshower))
            if trig_check(int(event_info.trigger_mask[0]))
        ])

        plt.suptitle(
            "z-t-Plot for DetID-{0} (t0set: {1}), Run {2}, FrameIndex {3}, "
            "TriggerCounter {4}, Overlays {5}, Trigger: {8}"
            "\n{7} UTC (time offset: {6} ns)".format(
                event_info.det_id[0], self.t0set, event_info.run_id[0],
                event_info.frame_index[0], event_info.trigger_counter[0],
                event_info.overlays[0], time_offset,
                datetime.utcfromtimestamp(event_info.utc_seconds),
                trigger_params),
            fontsize=fontsize,
            y=1.05)

        filename = 'ztplot'
        f = os.path.join(self.plots_path, filename + '.png')
        f_tmp = os.path.join(self.plots_path, filename + '_tmp.png')
        plt.savefig(f_tmp, dpi=120, bbox_inches="tight")
        if len(doms) > 4:
            plt.savefig(os.path.join(self.plots_path, filename + '_5doms.png'))
        plt.close('all')
        shutil.move(f_tmp, f)

    def finish(self):
        self.run = False


def main():
    from docopt import docopt
    args = docopt(__doc__)

    det_id = int(args['-d'])
    plots_path = args['-o']
    ligier_ip = args['-l']
    ligier_port = int(args['-p'])

    pipe = kp.Pipeline()
    pipe.attach(kp.io.ch.CHPump,
                host=ligier_ip,
                port=ligier_port,
                tags='IO_EVT, IO_SUM',
                timeout=60 * 60 * 24 * 7,
                max_queue=2000)
    pipe.attach(kp.io.daq.DAQProcessor)
    pipe.attach(ZTPlot, det_id=det_id, plots_path=plots_path)
    pipe.drain()


if __name__ == '__main__':
    main()