summaryrefslogtreecommitdiff
path: root/event.py
blob: 633fe26dda8334fb8edf6838cdd4ba43b0343781 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import numpy as np

from chroma.sample import uniform_sphere

class Photons(object):
    def __init__(self, positions, directions, wavelengths, polarizations=None, times=None, last_hit_triangles=None, histories=None):
        nphotons = len(positions)

        self.positions = positions
        self.directions = directions
        self.wavelengths = wavelengths

        if polarizations is not None:
            self.polarizations = polarizations
        else:
            # polarizations are isotropic in plane perpendicular
            # to the direction vectors
            polarizations = np.cross(directions, uniform_sphere(nphotons))
            # normalize polarization vectors
            polarizations /= np.tile(np.apply_along_axis(np.linalg.norm,1,polarizations),[3,1]).transpose()
            self.polarizations = np.asarray(polarizations, order='C')

        if times is None:
            self.times = np.zeros(nphotons, dtype=np.float32)
        else:
            self.times = times

        if last_hit_triangles is None:
            self.last_hit_triangles = np.empty(nphotons, dtype=np.int32)
            self.last_hit_triangles.fill(-1)
        else:
            self.last_hit_triangles = last_hit_triangles

        if histories is None:
            self.histories = np.zeros(nphotons, dtype=np.uint32)
        else:
            self.histories = histories

def concatenate_photons(photons):
    '''Merge a list of Photons objects into one long list of photons'''
    return Photons(positions=np.concatenate([p.positions for p in photons]),
                   directions=np.concatenate([p.directions for p in photons]),
                   polarizations=np.concatenate([p.polarizations for p in photons]),
                   times=np.concatenate([p.times for p in photons]),
                   wavelengths=np.concatenate([p.wavelengths for p in photons]))

class Channels(object):
    def __init__(self, hit, t, q, histories=None):
        self.hit = hit
        self.t = t
        self.q = q
        self.histories=histories

    def hit_channels(self):
        return self.hit.nonzero(), self.t[self.hit], self.q[self.hit]

class Subtrack(object):
    def __init__(self, particle_name, position, direction, start_time, total_energy):
        self.particle_name = particle_name
        self.position = position
        self.direction = direction
        self.start_time = start_time
        self.total_energy = total_energy

class Event(object):
    def __init__(self, event_id, particle_name=None, gen_position=None, gen_direction=None, gen_total_energy=None, ):
        self.event_id = event_id
        self.particle_name = particle_name
        self.gen_position = gen_position
        self.gen_direction = gen_direction
        self.gen_total_energy = gen_total_energy

        self.subtracks = []
        self.nphoton = 0
        self.photon_start = None
        self.photon_stop = None
        self.channels = None