diff options
Diffstat (limited to 'event.py')
| -rw-r--r-- | event.py | 102 |
1 files changed, 53 insertions, 49 deletions
@@ -1,69 +1,73 @@ import numpy as np -from chroma.sample import uniform_sphere +class Vertex(object): + def __init__(self, particle_name, pos, dir, pol, ke, t0=0.0): + self.particle_name = particle_name + self.pos = pos + self.dir = dir + self.pol = pol + self.ke = ke + self.t0 = t0 class Photons(object): - def __init__(self, positions, directions, wavelengths, polarizations=None, times=None, last_hit_triangles=None, histories=None): - self.positions = positions - - nphotons = len(positions) - - assert len(directions) == nphotons - self.directions = directions - - assert len(wavelengths) == nphotons + def __init__(self, pos, dir, pol, wavelengths, t=None, last_hit_triangles=None, flags=None): + self.pos = pos + self.dir = dir + self.pol = pol self.wavelengths = wavelengths - if polarizations is not None: - self.polarizations = polarizations + if t is None: + self.t = np.zeros(len(pos), dtype=np.float32) 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') + self.t = t - self.times = times - self.last_hit_triangles = last_hit_triangles - self.histories = histories + if last_hit_triangles is None: + self.last_hit_triangles = np.empty(len(pos), dtype=np.int32) + self.last_hit_triangles.fill(-1) + else: + self.last_hit_triangles = last_hit_triangles + + if flags is None: + self.flags = np.zeros(len(pos), dtype=np.uint32) + else: + self.flags = flags -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])) + def __add__(self, other): + pos = np.concatenate((self.pos, other.pos)) + dir = np.concatenate((self.dir, other.dir)) + pol = np.concatenate((self.pol, other.pol)) + wavelengths = np.concatenate((self.wavelengths, other.wavelengths)) + t = np.concatenate((self.t, other.t)) + last_hit_triangles = np.concatenate((self.last_hit_triangles, other.last_hit_triangles)) + flags = np.concatenate((self.flags, other.flags)) + return Photons(pos, dir, pol, wavelengths, t, last_hit_triangles, flags) class Channels(object): - def __init__(self, hit, t, q, histories=None): + def __init__(self, hit, t, q, flags=None): self.hit = hit self.t = t self.q = q - self.histories=histories + self.flags = flags 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 + def __init__(self, id=0, primary_vertex=None, vertices=None, photons_beg=None, photons_end=None, channels=None): + self.id = id + + self.nphotons = None + + self.primary_vertex = primary_vertex + + if vertices is not None: + if np.iterable(vertices): + self.vertices = vertices + else: + self.vertices = [vertices] + else: + self.vertices = [] - self.subtracks = [] - self.nphoton = 0 - self.photon_start = None - self.photon_stop = None - self.channels = None + self.photons_beg = photons_beg + self.photons_end = photons_end + self.channels = channels |