Full awkward (and super fast) cherenkov function implementation
The cherenkov function computes the expected arrival time for direct cherenkov photons from a track on a PMT (define as a hit position and direction). km3pipe provides an implementation in km3pipe.physics
but it only works on an "event-per-event" basis, with the need of comparing a regular sized list of hits due to the numpy implicit conversion to a single track.
To get much better performances from the IO, and limiting the python overhead of looping over the events, I came up with this implementation that takes awkard arrays as input:
import awkward as ak
import numpy as np
from km3pipe.constants import (
C_WATER,
SIN_CHERENKOV,
TAN_CHERENKOV,
V_LIGHT_WATER,
C_LIGHT,
)
def cherenkov(trks, hits, do_direction = False):
""" Super fast implemention of km3pipe.physics.cherenkov"""
V_x = hits.pos_x - trks.pos_x
V_y = hits.pos_y - trks.pos_y
V_z = hits.pos_z - trks.pos_z
L = V_x * trks.dir_x + V_y * trks.dir_y + V_z * trks.dir_z
d_photon_closest = np.sqrt(V_x*V_x + V_y*V_y + V_z*V_z - L*L)
d_photon = d_photon_closest / SIN_CHERENKOV
d_track = L - d_photon_closest / TAN_CHERENKOV
t_photon = trks.t + d_track / C_LIGHT + d_photon / V_LIGHT_WATER
results = {
"d_photon_closest": d_photon_closest,
"d_photon":d_photon,
"d_track":d_track,
"t_photon":t_photon
}
if do_direction:
V_photon_x = V_x - d_track * trks.dir_x
V_photon_y = V_y - d_track * trks.dir_y
V_photon_z = V_z - d_track * trks.dir_z
norm = np.sqrt(V_photon_x**2 + V_photon_y**2 + V_photon_z**2)
dir_photon_x = V_photon_x / norm
dir_photon_y = V_photon_y / norm
dir_photon_z = V_photon_z / norm
cos_angle_photon = dir_photon_x * hits.dir_x + dir_photon_y * hits.dir_y + dir_photon_z * hits.dir_z
results['dir_photon_x'] = dir_photon_x
results['dir_photon_y'] = dir_photon_y
results['dir_photon_z'] = dir_photon_z
results['cos_angle'] = cos_angle_photon
return ak.Array(results)
This way, you can use it on km3io's OfflineReader
objects directly:
import km3io as kio
f = kio.OfflineReader(fname)
f = f[f.n_tracks > 0]
best_tracks = kio.tools.best_jmuon(f.tracks)
cherencov_hits = cherenkov(best_tracks, f.hits)
This is extremely fast, with about 17 seconds spent in the cherenkov
function for 185k tracks and 26,4M hits.
I have checked the results for few events with respect to the previous code, that looks ok.
Looking at the hits.t - cherencov_hits.t_photon
, I get e.g. the following plot:
@tgal I'm super happy with the result, because that's crazy fast I think, but now I don't really know what to do with it. The current implementation in km3pipe is really "single event blob" oriented, and that's the main way km3pipe is used currently. Should this function still be added but with a different name ?