#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# vim:set ts=4 sts=4 sw=4 et:
"""MC Helpers.
"""
from __future__ import absolute_import, print_function, division

from km3pipe import Module
from km3pipe.mc import pdg2name
from km3pipe.math import zenith, azimuth

NEUTRINOS = {
    'nu_e',
    'anu_e',
    'nu_mu',
    'anu_mu',
}    # noqa


class McTruth(Module):
    """Extract MC info of 1st MC track.

    Parameters
    ----------
    most_energetic_primary: bool, default=True
    """

    def configure(self):
        self.most_energetic_primary = bool(self.get('most_energetic_primary')
                                           ) or True

    @classmethod
    def t2f(cls, row):
        return pdg2name(row['type'])

    @classmethod
    def is_nu(cls, flavor):
        return flavor in NEUTRINOS

    def process(self, blob):
        mc = blob['McTracks'].conv_to('pandas')
        if self.most_energetic_primary:
            mc.sort_values('energy', ascending=False, inplace=True)
        mc = mc.head(1)
        mc['zenith'] = zenith(mc[['dir_x', 'dir_y', 'dir_z']])
        mc['azimuth'] = azimuth(mc[['dir_x', 'dir_y', 'dir_z']])
        flavor = mc.apply(self.t2f, axis=1)
        mc['is_neutrino'] = flavor.apply(self.is_nu)
        blob['McTruth'] = mc
        return blob

@@ -155,16 +155,36 @@
 vision%0A%0A
+import numpy as np%0A%0A
 from km3
@@ -276,16 +276,77 @@
 zimuth%0A%0A
+__author__ = %22Michael Moser and Tamas Gal and Moritz Lotze%22%0A%0A
 NEUTRINO
@@ -407,19 +407,38 @@
 u',%0A
-%7D    # noqa
+    'nu_tau',%0A    'anu_tau',%0A%7D
 %0A%0A%0Ac
@@ -1336,16 +1336,16 @@
 .is_nu)%0A
-
         
@@ -1385,8 +1385,1384 @@
 rn blob%0A
+%0A%0Adef time_converter(hits_times_jte, evt_timestamp_in_ns, evt_mc_time):%0A    %22%22%22%0A    Function that converts JTE times of hits to MC times.%0A%0A    Parameters%0A    ----------%0A    hits_times_jte : np.ndarray%0A        Hits array with JTE times.%0A    evt_timestamp_in_ns : int%0A        Total timestamp of the event in nanoseconds.%0A    evt_mc_time : int%0A        Mc time of the event in nanoseconds.%0A%0A    Returns%0A    -------%0A    ndarray%0A        Converted hits array with MC times.%0A    %22%22%22%0A    hits_times_mc = hits_times_jte + evt_timestamp_in_ns - evt_mc_time%0A    return hits_times_mc%0A%0A%0Aclass MCTimeCorrector(Module):%0A    %22%22%22%0A    Module that converts JTE hit times to MC times.%0A%0A    Parameters%0A    ----------%0A    hits_key : str, optional%0A        Name of the Hits to convert (default: 'Hits').%0A    event_info_key : str, optional%0A        Name of the EventInfo to store this in (default: 'EventInfo').%0A    %22%22%22%0A%0A    def configure(self):%0A        self.hits_key = self.get('hits_key', default='Hits')%0A        self.event_info_key = self.get('event_info_key', default='EventInfo')%0A        self.convert_time = np.frompyfunc(time_converter, 1, 1)%0A%0A    def process(self, blob):%0A        event_info = blob%5Bself.event_info_key%5D%0A        hits = blob%5Bself.hits_key%5D%0A        hits%5B'time'%5D = self.convert_time(%0A            hits.time, event_info.timestamp_in_ns, event_info.mc_time%0A        )%0A        return blob%0A

Add slightly modified converter implementation