summaryrefslogtreecommitdiff
path: root/fileio/root.py
diff options
context:
space:
mode:
Diffstat (limited to 'fileio/root.py')
-rw-r--r--fileio/root.py186
1 files changed, 98 insertions, 88 deletions
diff --git a/fileio/root.py b/fileio/root.py
index 4c9d9bb..507c249 100644
--- a/fileio/root.py
+++ b/fileio/root.py
@@ -14,61 +14,66 @@ def tvector3_to_ndarray(vec):
def make_photon_with_arrays(size):
'''Returns a new chroma.event.Photons object for `size` number of
photons with empty arrays set for all the photon attributes.'''
- return event.Photons(positions=np.empty((size,3), dtype=np.float32),
- directions=np.empty((size,3), dtype=np.float32),
- polarizations=np.empty((size,3), dtype=np.float32),
+ return event.Photons(pos=np.empty((size,3), dtype=np.float32),
+ dir=np.empty((size,3), dtype=np.float32),
+ pol=np.empty((size,3), dtype=np.float32),
wavelengths=np.empty(size, dtype=np.float32),
- times=np.empty(size, dtype=np.float32),
- histories=np.empty(size, dtype=np.uint32),
+ t=np.empty(size, dtype=np.float32),
+ flags=np.empty(size, dtype=np.uint32),
last_hit_triangles=np.empty(size, dtype=np.int32))
+def root_vertex_to_python_vertex(vertex):
+ "Returns a chroma.event.Vertex object from a root Vertex object."
+ return event.Vertex(str(vertex.particle_name),
+ tvector3_to_ndarray(vertex.pos),
+ tvector3_to_ndarray(vertex.dir),
+ tvector3_to_ndarray(vertex.pol),
+ vertex.ke,
+ vertex.t0)
def root_event_to_python_event(ev):
'''Returns a new chroma.event.Event object created from the
contents of the ROOT event `ev`.'''
- pyev = event.Event(ev.event_id)
-
- # MC
- pyev.particle_name = str(ev.mc.particle)
- pyev.gen_position = tvector3_to_ndarray(ev.mc.gen_position)
- pyev.gen_direction = tvector3_to_ndarray(ev.mc.gen_direction)
- pyev.gen_total_energy = ev.mc.gen_total_energy
-
- pyev.nphoton = ev.mc.nphoton
-
- for subtrack in ev.mc.subtrack:
- pysubtrack = event.Subtrack(str(subtrack.particle),
- tvector3_to_ndarray(subtrack.position),
- tvector3_to_ndarray(subtrack.direction),
- subtrack.start_time,
- subtrack.total_energy)
- pyev.subtracks.append(pysubtrack)
-
- # photon start
- if ev.mc.photon_start.size() > 0:
- photons = make_photon_with_arrays(ev.mc.photon_start.size())
- ROOT.get_photons(ev.mc.photon_start, photons.positions.ravel(), photons.directions.ravel(),
- photons.polarizations.ravel(), photons.wavelengths, photons.times,
- photons.histories, photons.last_hit_triangles)
- pyev.photon_start = photons
-
- # photon stop
- if ev.mc.photon_stop.size() > 0:
- photons = make_photon_with_arrays(ev.mc.photon_stop.size())
- ROOT.get_photons(ev.mc.photon_stop, photons.positions.ravel(), photons.directions.ravel(),
- photons.polarizations.ravel(), photons.wavelengths, photons.times,
- photons.histories, photons.last_hit_triangles)
- pyev.photon_stop = photons
-
- # hits
- max_channel_id = ev.max_channel_id
- hit = np.empty(shape=max_channel_id+1, dtype=np.int32)
- t = np.empty(shape=max_channel_id+1, dtype=np.float32)
- q = np.empty(shape=max_channel_id+1, dtype=np.float32)
- histories = np.empty(shape=max_channel_id+1, dtype=np.uint32)
-
- ev.get_channels(max_channel_id+1, hit, t, q, histories)
- pyev.channels = event.Channels(hit.astype(bool), t, q, histories)
+ pyev = event.Event(ev.id)
+ pyev.primary_vertex = root_vertex_to_python_vertex(ev.primary_vertex)
+
+ for vertex in ev.vertices:
+ pyev.vertices.append(root_vertex_to_python_vertex(vertex))
+
+ # photon begin
+ if ev.photons_beg.size() > 0:
+ photons = make_photon_with_arrays(ev.photons_beg.size())
+ ROOT.get_photons(ev.photons_beg,
+ photons.pos.ravel(),
+ photons.dir.ravel(),
+ photons.pol.ravel(),
+ photons.wavelengths,
+ photons.t,
+ photons.last_hit_triangles,
+ photons.flags)
+ pyev.photons_beg = photons
+
+ # photon end
+ if ev.photons_end.size() > 0:
+ photons = make_photon_with_arrays(ev.photons_end.size())
+ ROOT.get_photons(ev.photons_end,
+ photons.pos.ravel(),
+ photons.dir.ravel(),
+ photons.pol.ravel(),
+ photons.wavelengths,
+ photons.t,
+ photons.last_hit_triangles,
+ photons.flags)
+ pyev.photons_end = photons
+
+ # channels
+ hit = np.empty(ev.nchannels, dtype=np.int32)
+ t = np.empty(ev.nchannels, dtype=np.float32)
+ q = np.empty(ev.nchannels, dtype=np.float32)
+ flags = np.empty(ev.nchannels, dtype=np.uint32)
+
+ ROOT.get_channels(ev, hit, t, q, flags)
+ pyev.channels = event.Channels(hit.astype(bool), t, q, flags)
return pyev
class RootReader(object):
@@ -140,48 +145,53 @@ class RootWriter(object):
self.T.Branch('ev', self.ev)
def write_event(self, pyev):
- '''Write an event.Event object to the ROOT tree as a ROOT.Event object.'''
- self.ev.event_id = pyev.event_id
-
- self.ev.mc.particle = pyev.particle_name
- self.ev.mc.gen_position.SetXYZ(*pyev.gen_position)
- self.ev.mc.gen_direction.SetXYZ(*pyev.gen_direction)
- self.ev.mc.gen_total_energy = pyev.gen_total_energy
- self.ev.mc.nphoton = pyev.nphoton
-
- if pyev.photon_start is not None:
- photons = pyev.photon_start
- ROOT.fill_photons(self.ev.mc.photon_start,
- len(photons.positions),
- np.ravel(photons.positions),
- np.ravel(photons.directions),
- np.ravel(photons.polarizations),
- photons.wavelengths, photons.times,
- photons.histories, photons.last_hit_triangles)
- if pyev.photon_stop is not None:
- photons = pyev.photon_stop
- ROOT.fill_photons(self.ev.mc.photon_stop,
- len(photons.positions),
- np.ravel(photons.positions),
- np.ravel(photons.directions),
- np.ravel(photons.polarizations),
- photons.wavelengths, photons.times,
- photons.histories, photons.last_hit_triangles)
-
- self.ev.mc.subtrack.resize(0)
- if pyev.subtracks is not None:
- self.ev.mc.subtrack.resize(len(pyev.subtracks))
- for i, subtrack in enumerate(pyev.subtracks):
- self.ev.mc.subtrack[i].name = subtrack.particle_name
- self.ev.mc.subtrack[i].position.SetXYZ(*subtrack.position)
- self.ev.mc.subtrack[i].direction.SetXYZ(*subtrack.direction)
- self.ev.mc.subtrack[i].start_time = subtrack.start_time
- self.ev.mc.subtrack[i].total_energy = subtrack.total_energy
-
- ROOT.fill_hits(self.ev, len(pyev.channels.t), pyev.channels.t, pyev.channels.q, pyev.channels.histories)
+ "Write an event.Event object to the ROOT tree as a ROOT.Event object."
+ self.ev.id = pyev.id
+
+ self.ev.primary_vertex.particle_name = pyev.primary_vertex.particle_name
+ self.ev.primary_vertex.pos.SetXYZ(*pyev.primary_vertex.pos)
+ self.ev.primary_vertex.dir.SetXYZ(*pyev.primary_vertex.dir)
+ if pyev.primary_vertex.pol is not None:
+ self.ev.primary_vertex.pol.SetXYZ(*pyev.primary_vertex.pol)
+ self.ev.primary_vertex.ke = pyev.primary_vertex.ke
+
+ if pyev.photons_beg is not None:
+ photons = pyev.photons_beg
+ ROOT.fill_photons(self.ev.photons_beg,
+ len(photons.pos),
+ photons.pos.ravel(),
+ photons.dir.ravel(),
+ photons.pol.ravel(),
+ photons.wavelengths, photons.t,
+ photons.last_hit_triangles, photons.flags)
+
+ if pyev.photons_end is not None:
+ photons = pyev.photons_end
+ ROOT.fill_photons(self.ev.photons_end,
+ len(photons.pos),
+ photons.pos.ravel(),
+ photons.dir.ravel(),
+ photons.pol.ravel(),
+ photons.wavelengths, photons.t,
+ photons.last_hit_triangles, photons.flags)
+
+ self.ev.vertices.resize(0)
+ if pyev.vertices is not None:
+ self.ev.vertices.resize(len(pyev.vertices))
+ for i, vertex in enumerate(pyev.vertices):
+ self.ev.vertices[i].particle_name = vertex.particle_name
+ self.ev.vertices[i].pos.SetXYZ(*vertex.pos)
+ self.ev.vertices[i].dir.SetXYZ(*vertex.dir)
+ if vertex.pol is not None:
+ self.ev.vertices[i].pol.SetXYZ(*vertex.pol)
+ self.ev.vertices[i].ke = vertex.ke
+ self.ev.vertices[i].t0 = vertex.t0
+
+ if pyev.channels is not None:
+ ROOT.fill_channels(self.ev, np.count_nonzero(pyev.channels.hit), np.arange(len(pyev.channels.t))[pyev.channels.hit].astype(np.int32), pyev.channels.t, pyev.channels.q, pyev.channels.flags, len(pyev.channels.hit))
+
self.T.Fill()
def close(self):
self.T.Write()
self.file.Close()
-
generated by cgit (git 2.34.1) at 2025-05-09 01:06:22 +0000