From 084dfd08b714faefaea77cb7dc04d2e93dc04b1d Mon Sep 17 00:00:00 2001 From: Stan Seibert Date: Fri, 16 Sep 2011 14:27:46 -0400 Subject: File reorganization to move toward standard python package layout --- fileio/__init__.py | 0 fileio/root.C | 147 --------------------------------------- fileio/root.py | 199 ----------------------------------------------------- 3 files changed, 346 deletions(-) delete mode 100644 fileio/__init__.py delete mode 100644 fileio/root.C delete mode 100644 fileio/root.py (limited to 'fileio') diff --git a/fileio/__init__.py b/fileio/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/fileio/root.C b/fileio/root.C deleted file mode 100644 index 2b39b1b..0000000 --- a/fileio/root.C +++ /dev/null @@ -1,147 +0,0 @@ -#include -#include -#include -#include - -struct Vertex { - std::string particle_name; - TVector3 pos; - TVector3 dir; - TVector3 pol; - double ke; - double t0; - - ClassDef(Vertex, 1); -}; - -struct Photon { - double t; - TVector3 pos; - TVector3 dir; - TVector3 pol; - double wavelength; // nm - unsigned int flag; - int last_hit_triangle; - - ClassDef(Photon, 1); -}; - -struct Channel { - Channel() : id(-1), t(-1e9), q(-1e9) { }; - int id; - double t; - double q; - unsigned int flag; - - ClassDef(Channel, 1); -}; - -struct Event { - int id; - unsigned int nhit; - unsigned int nchannels; - - Vertex primary_vertex; - - std::vector vertices; - std::vector photons_beg; - std::vector photons_end; - std::vector channels; - - ClassDef(Event, 1); -}; - -void fill_channels(Event *ev, unsigned int nhit, unsigned int *ids, float *t, - float *q, unsigned int *flags, unsigned int nchannels) -{ - ev->nhit = 0; - ev->nchannels = nchannels; - ev->channels.resize(0); - - Channel ch; - unsigned int id; - for (unsigned int i=0; i < nhit; i++) { - ev->nhit++; - id = ids[i]; - ch.id = id; - ch.t = t[id]; - ch.q = q[id]; - ch.flag = flags[id]; - ev->channels.push_back(ch); - } -} - -void get_channels(Event *ev, int *hit, float *t, float *q, unsigned int *flags) -{ - for (unsigned int i=0; i < ev->nchannels; i++) { - hit[i] = 0; - t[i] = -1e9f; - q[i] = -1e9f; - flags[i] = 0; - } - - unsigned int id; - for (unsigned int i=0; i < ev->channels.size(); i++) { - id = ev->channels[i].id; - - if (id < ev->nchannels) { - hit[id] = 1; - t[id] = ev->channels[i].t; - q[id] = ev->channels[i].q; - flags[id] = ev->channels[i].flag; - } - } -} - -void get_photons(const std::vector &photons, float *pos, float *dir, - float *pol, float *wavelengths, float *t, - int *last_hit_triangles, unsigned int *flags) -{ - for (unsigned int i=0; i < photons.size(); i++) { - const Photon &photon = photons[i]; - pos[3*i] = photon.pos.X(); - pos[3*i+1] = photon.pos.Y(); - pos[3*i+2] = photon.pos.Z(); - - dir[3*i] = photon.dir.X(); - dir[3*i+1] = photon.dir.Y(); - dir[3*i+2] = photon.dir.Z(); - - pol[3*i] = photon.pol.X(); - pol[3*i+1] = photon.pol.Y(); - pol[3*i+2] = photon.pol.Z(); - - wavelengths[i] = photon.wavelength; - t[i] = photon.t; - flags[i] = photon.flag; - last_hit_triangles[i] = photon.last_hit_triangle; - } -} - -void fill_photons(std::vector &photons, - unsigned int nphotons, float *pos, float *dir, - float *pol, float *wavelengths, float *t, - int *last_hit_triangles, unsigned int *flags) -{ - photons.resize(nphotons); - - for (unsigned int i=0; i < nphotons; i++) { - Photon &photon = photons[i]; - photon.t = t[i]; - photon.pos.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]); - photon.dir.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]); - photon.pol.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]); - photon.wavelength = wavelengths[i]; - photon.last_hit_triangle = last_hit_triangles[i]; - photon.flag = flags[i]; - - } -} - -#ifdef __MAKECINT__ -#pragma link C++ class vector; -#pragma link C++ class vector; -#pragma link C++ class vector; -#endif - - diff --git a/fileio/root.py b/fileio/root.py deleted file mode 100644 index af86deb..0000000 --- a/fileio/root.py +++ /dev/null @@ -1,199 +0,0 @@ -import ROOT -import os.path -import numpy as np - -ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g')) - -import ROOT -import chroma.event as event - -def tvector3_to_ndarray(vec): - '''Convert a ROOT.TVector3 into a numpy np.float32 array''' - return np.array((vec.X(), vec.Y(), vec.Z()), dtype=np.float32) - -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(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), - 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), - pos=tvector3_to_ndarray(vertex.pos), - dir=tvector3_to_ndarray(vertex.dir), - ke=vertex.ke, - t0=vertex.t0, - pol=tvector3_to_ndarray(vertex.pol)) - -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.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): - '''Reader of Chroma events from a ROOT file. This class can be used to - navigate up and down the file linearly or in a random access fashion. - All returned events are instances of the chroma.event.Event class. - - It implements the iterator protocol, so you can do - - for ev in RootReader('electron.root'): - # process event here - ''' - - def __init__(self, filename): - '''Open ROOT file named `filename` containing TTree `T`.''' - self.f = ROOT.TFile(filename) - self.T = self.f.T - self.i = -1 - - def __len__(self): - '''Returns number of events in this file.''' - return self.T.GetEntries() - - def next(self): - '''Return the next event in the file. Raises StopIteration - when you get to the end.''' - if self.i + 1 >= len(self): - raise StopIteration - - self.i += 1 - self.T.GetEntry(self.i) - return root_event_to_python_event(self.T.ev) - - def prev(self): - '''Return the next event in the file. Raises StopIteration if - that would go past the beginning.''' - if self.i <= 0: - self.i = -1 - raise StopIteration - - self.i -= 1 - self.T.GetEntry(self.i) - return root_event_to_python_event(self.T.ev) - - def current(self): - '''Return the current event in the file.''' - self.T.GetEntry(self.i) # just in case? - return root_event_to_python_event(self.T.ev) - - def jump_to(self, index): - '''Return the event at `index`. Updates current location.''' - if index < 0 or index >= len(self): - raise IndexError - - self.T.GetEntry(self.i) - return root_event_to_python_event(self.T.ev) - - def index(self): - '''Return the current event index''' - return self.i - -class RootWriter(object): - def __init__(self, filename): - self.filename = filename - self.file = ROOT.TFile(filename, 'RECREATE') - - self.T = ROOT.TTree('T', 'Chroma events') - self.ev = ROOT.Event() - 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.id = pyev.id - - if pyev.primary_vertex is not None: - 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() -- cgit