def __init__(self, filename): self.filename = filename self.file = ROOT.TFile(filename, 'RECREATE') self.T = ROOT.TTree('T', 'Chroma events') self.ev = ROOT.Event_short() self.T.Branch('ev', self.ev)
def __init__(self, filename): self.filename = filename self.file = ROOT.TFile(filename, 'RECREATE') self.T = ROOT.TTree('T', 'PDFs') self.pdf = ROOT.PMT_pdf() self.T.Branch('pdf', self.pdf)
def python_vertex_to_root_vertex(pvertex, rvertex): rvertex.particle_name = pvertex.particle_name rvertex.pos.SetXYZ(*pvertex.pos) rvertex.dir.SetXYZ(*pvertex.dir) if pvertex.pol is not None: rvertex.pol.SetXYZ(*vertex.pol) rvertex.ke = pvertex.ke rvertex.t0 = pvertex.t0 rvertex.trackid = pvertex.trackid rvertex.pdgcode = pvertex.pdgcode if pvertex.steps: ROOT.fill_steps(rvertex, len(pvertex.steps.x), pvertex.steps.x, pvertex.steps.y, pvertex.steps.z, pvertex.steps.t, pvertex.steps.dx, pvertex.steps.dy, pvertex.steps.dz, pvertex.steps.ke, pvertex.steps.edep, pvertex.steps.qedep) else: nil = np.empty(0, dtype=np.float64) ROOT.clear_steps(rvertex) if pvertex.children is not None and len(pvertex.children) > 0: rvertex.children.resize(len(pvertex.children)) any( python_vertex_to_root_vertex(pchild, rchild) for pchild, rchild in zip(pvertex.children, rvertex.children)) else: rvertex.children.resize(0)
def root_vertex_to_python_vertex(vertex): "Returns a chroma.event.Vertex object from a root Vertex object." if len(vertex.step_x): n = len(vertex.step_x) steps = event.Steps(np.empty(n),np.empty(n),np.empty(n),np.empty(n), np.empty(n),np.empty(n),np.empty(n), np.empty(n),np.empty(n),np.empty(n)) ROOT.get_steps(vertex,n,steps.x,steps.y,steps.z,steps.t, steps.dx,steps.dy,steps.dz, steps.ke,steps.edep,steps.qedep) else: steps = None if len(vertex.children) > 0: children = [root_vertex_to_python_vertex(child) for child in vertex.children] else: children = None 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), trackid=vertex.trackid, pdgcode=vertex.pdgcode, steps=steps, children=children)
def __init__(self, filename): self.filename = filename self.file = ROOT.TFile(filename, 'RECREATE') self.T = ROOT.TTree('T', 'Angles') self.angle = ROOT.Gauss_angle() self.T.Branch('angles', self.angle)
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) # 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 return pyev
def root_PMT_Gauss_angle_to_tuple(angle): '''Returns a tuple of (pmt_bin_ind, 1x3 array mean, float sigma) created from the contents of the ROOT Gauss_angle `angle`.''' mean = np.empty((1,3), dtype=np.float32) sigma = np.zeros(1, dtype=np.float32) ROOT.get_Gauss_angle(angle, mean.ravel(), sigma) return angle.pmt_bin_ind, mean, sigma[0]
def write_PMT(self, mean, sigma, pmt_bin_ind): #Write a 1x3 numpy array (mean) and float sigma to a ROOT.Gauss_angle object. mean_shape = np.shape(mean) if mean_shape[0] != 3: print "Error: mean must have three coordinates." return self.angle.pmt_bin_ind = pmt_bin_ind #Fill ROOT.Gauss_angle object by flattening 3D array to 1D w/ ravel; flattens along inner dimensions first ROOT.fill_Gauss_angle(self.angle, mean, sigma) self.T.Fill()
def root_PMT_angles_to_numpy_array(pmt_angles): '''Returns a tuple of (pmt_bin_ind, nX x nY x nZ x 3 numpy array) created from the contents of the ROOT PMT_angles `pmt_angles`.''' det_bins = pmt_angles.detector_bins_x*pmt_angles.detector_bins_y*pmt_angles.detector_bins_z ang_out = np.empty((det_bins,3), dtype=np.float32) ROOT.get_angles(pmt_angles.angles, ang_out.ravel()) #print ang_out #print np.shape(ang_out) #np_ang = np.reshape(ang_out,(pmt_angles.detector_bins_x,pmt_angles.detector_bins_y,pmt_angles.detector_bins_z,3)) #print np.shape(np_ang) #print np_ang[0,0,0] return pmt_angles.pmt_bin_ind, ang_out
def write_event(self, pdf, pmt_bin_ind): "Write 3D numpy array (PDF) to a ROOT.PMT_pdf object." if np.shape(pdf) == (0,): return pdf_shape = np.shape(pdf) if len(pdf_shape) != 3: print "Error: array to write must be 3-dimensional." return self.pdf.pmt_bin_ind = pmt_bin_ind self.pdf.detector_bins_x = pdf_shape[0] self.pdf.detector_bins_y = pdf_shape[1] self.pdf.detector_bins_z = pdf_shape[2] #Fill ROOT.PMT_pdf object by flattening 3D array to 1D w/ ravel; flattens along inner dimensions first ROOT.fill_pdf(self.pdf, np.size(pdf), pdf.ravel()) self.T.Fill()
def testAngularDistributionPolarized(self): # Fully polarized photons self.photons.pol[:] = [1.0, 0.0, 0.0] photons_end = self.sim.simulate([self.photons], keep_photons_end=True, max_steps=1).next().photons_end aborted = (photons_end.flags & (1 << 31)) > 0 self.assertFalse(aborted.any()) # Compute the dot product between initial and final dir rayleigh_scatters = (photons_end.flags & (1 << 4)) > 0 cos_scatter = (self.photons.dir[rayleigh_scatters] * photons_end.dir[rayleigh_scatters]).sum(axis=1) theta_scatter = np.arccos(cos_scatter) h = Histogram(bins=100, range=(0, np.pi)) h.fill(theta_scatter) h = rootify(h) # The functional form for polarized light should be # (1 + \cos^2 \theta)\sin \theta according to GEANT4 physics # reference manual. f = ROOT.TF1("pol_func", "[0]*(1+cos(x)**2)*sin(x)", 0, np.pi) h.Fit(f, 'NQ') self.assertGreater(f.GetProb(), 1e-3)
def __init__(self, filename, detector=None): self.filename = filename self.file = ROOT.TFile(filename, 'RECREATE') if detector is not None: ch_info = ROOT.TTree('CH', 'Chroma channel info') ch_pos = ROOT.TVector3() ch_type = array.array( 'i', [0]) ch_info.Branch('pos',ch_pos) ch_info.Branch('type',ch_type,'type/I') for pos,chtype in zip(detector.channel_index_to_position, detector.channel_index_to_channel_type): ch_pos.SetXYZ(*pos) ch_type[0] = chtype ch_info.Fill() ch_info.Write() self.T = ROOT.TTree('T', 'Chroma events') self.ev = ROOT.Event() self.T.Branch('ev', self.ev)
def write_PMT(self, angles, pmt_bin_ind, nx, ny, nz): #Write nbins x 3 numpy array (angles; one 3D direction per detector bin) to a ROOT.PMT_angles object. if np.shape(angles) == (0,): return angles_shape = np.shape(angles) if angles_shape[0] != nx*ny*nz: print "Error: length of angles array must equal nxbins*nybins*nzbins." return if angles_shape[1] != 3: print "Error: each entry must have three coordinates." return self.angles.pmt_bin_ind = pmt_bin_ind self.angles.detector_bins_x = nx self.angles.detector_bins_y = ny self.angles.detector_bins_z = nz #Fill ROOT.PMT_angles object by flattening 3D array to 1D w/ ravel; flattens along inner dimensions first ROOT.fill_angles(self.angles, angles_shape[0], angles.ravel()) self.T.Fill()
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.photons_beg is not None: photons = pyev.photons_beg ROOT.fill_photons(self.ev.photons_beg, len(photons.pos), photons.pos.ravel(), 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.flags) self.T.Fill()
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 self.ev.primary_vertex.t0 = pyev.primary_vertex.t0 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: nhit = count_nonzero(pyev.channels.hit) if nhit > 0: ROOT.fill_channels(self.ev, nhit, np.arange(len(pyev.channels.t))[pyev.channels.hit].astype(np.uint32), pyev.channels.t, pyev.channels.q, pyev.channels.flags, len(pyev.channels.hit)) else: self.ev.nhit = 0 self.ev.channels.resize(0) self.ev.nchannels = len(pyev.channels.hit) else: self.ev.nhit = 0 self.ev.channels.resize(0) self.ev.nchannels = 0 self.T.Fill()
def __init__(self, filename): '''Open ROOT file named `filename` containing TTree `T`.''' self.f = ROOT.TFile(filename) if hasattr(self.f,'CH'): ch_info = self.f.CH ch_num = ch_info.GetEntries() self.ch_pos = np.empty((ch_num,3),dtype=np.float32) self.ch_type = np.empty((ch_num,),dtype=np.int32) for i in range(ch_num): ch_info.GetEntry(i) ch_info.pos.GetXYZ(self.ch_pos[i]) self.ch_type[i] = ch_info.type self.T = self.f.T self.i = -1
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 if ev.nchannels > 0: 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) else: pyev.channels = None return pyev
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) 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 if ev.nchannels > 0: 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) else: pyev.channels = None return pyev
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.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): python_vertex_to_root_vertex(vertex,self.ev.vertices[i]) if pyev.channels is not None: nhit = count_nonzero(pyev.channels.hit) if nhit > 0: ROOT.fill_channels(self.ev, nhit, np.arange(len(pyev.channels.t))[pyev.channels.hit].astype(np.uint32), pyev.channels.t, pyev.channels.q, pyev.channels.flags, len(pyev.channels.hit)) else: self.ev.nhit = 0 self.ev.channels.resize(0) self.ev.nchannels = len(pyev.channels.hit) else: self.ev.nhit = 0 self.ev.channels.resize(0) self.ev.nchannels = 0 self.T.Fill()
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.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): python_vertex_to_root_vertex(vertex, self.ev.vertices[i]) if pyev.channels is not None: nhit = count_nonzero(pyev.channels.hit) if nhit > 0: ROOT.fill_channels( self.ev, nhit, np.arange(len(pyev.channels.t))[pyev.channels.hit].astype( np.uint32), pyev.channels.t, pyev.channels.q, pyev.channels.flags, len(pyev.channels.hit)) else: self.ev.nhit = 0 self.ev.channels.resize(0) self.ev.nchannels = len(pyev.channels.hit) else: self.ev.nhit = 0 self.ev.channels.resize(0) self.ev.nchannels = 0 self.T.Fill()
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) 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, photons.channel) 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, photons.channel) pyev.photons_end = photons # photon tracks if ev.photon_tracks.size() > 0: photon_tracks = [] for i in range(ev.photon_tracks.size()): photons = make_photon_with_arrays(ev.photon_tracks[i].size()) ROOT.get_photons(ev.photon_tracks[i], photons.pos.ravel(), photons.dir.ravel(), photons.pol.ravel(), photons.wavelengths, photons.t, photons.last_hit_triangles, photons.flags, photons.channel) photon_tracks.append(photons) pyev.photon_tracks = photon_tracks pyev.photon_parent_trackids = np.asarray(ev.photon_parent_trackids).copy() # hits if ev.hits.size() > 0: pyev.hits = {} for hit in ev.hits: photons = make_photon_with_arrays(hit.second.size()) ROOT.get_photons(hit.second, photons.pos.ravel(), photons.dir.ravel(), photons.pol.ravel(), photons.wavelengths, photons.t, photons.last_hit_triangles, photons.flags, photons.channel) pyev.hits[hit.first] = photons # flat_hits if ev.flat_hits.size() > 0: photons = make_photon_with_arrays(ev.flat_hits.size()) ROOT.get_photons(ev.flat_hits, photons.pos.ravel(), photons.dir.ravel(), photons.pol.ravel(), photons.wavelengths, photons.t, photons.last_hit_triangles, photons.flags, photons.channel) pyev.flat_hits = 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, photons.channel) pyev.photons_end = photons # channels if ev.nchannels > 0: hit = np.zeros(ev.nchannels, dtype=np.int32) t = np.zeros(ev.nchannels, dtype=np.float32) q = np.zeros(ev.nchannels, dtype=np.float32) flags = np.zeros(ev.nchannels, dtype=np.uint32) ROOT.get_channels(ev, hit, t, q, flags) pyev.channels = event.Channels(hit.astype(bool), t, q, flags) else: pyev.channels = None return pyev
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 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.photons_beg is not None: photons = pyev.photons_beg if len(photons.pos) > 0: 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, photons.channel) else: self.ev.photons_beg.resize(0) if pyev.photons_end is not None: photons = pyev.photons_end if len(photons.pos) > 0: 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, photons.channel) else: self.ev.photons_end.resize(0) if pyev.photon_tracks is not None: self.ev.photon_tracks.resize(len(pyev.photon_tracks)) for i in range(len(pyev.photon_tracks)): photons = pyev.photon_tracks[i] if len(photons.pos) > 0: ROOT.fill_photons(self.ev.photon_tracks[i], len(photons.pos), photons.pos.ravel(), photons.dir.ravel(), photons.pol.ravel(), photons.wavelengths, photons.t, photons.last_hit_triangles, photons.flags, photons.channel) self.ev.photon_parent_trackids.resize(len(pyev.photon_parent_trackids)) np.asarray(self.ev.photon_parent_trackids)[:] = pyev.photon_parent_trackids else: self.ev.photon_tracks.resize(0) self.ev.photon_parent_trackids.resize(0) if pyev.vertices is not None: self.ev.vertices.resize(len(pyev.vertices)) for i, vertex in enumerate(pyev.vertices): python_vertex_to_root_vertex(vertex,self.ev.vertices[i]) else: self.ev.vertices.resize(0) if pyev.hits is not None: self.ev.hits.clear() for hit in pyev.hits: photons = pyev.hits[hit] if len(photons.pos) > 0: ROOT.fill_photons(self.ev.hits[hit],len(photons.pos), photons.pos.ravel(), photons.dir.ravel(), photons.pol.ravel(), photons.wavelengths, photons.t, photons.last_hit_triangles, photons.flags, photons.channel) else: self.ev.hits.clear() if pyev.flat_hits is not None: photons = pyev.flat_hits if len(photons.pos) > 0: ROOT.fill_photons(self.ev.flat_hits, len(photons.pos), photons.pos.ravel(), photons.dir.ravel(), photons.pol.ravel(), photons.wavelengths, photons.t, photons.last_hit_triangles, photons.flags, photons.channel) else: self.ev.flat_hits.resize(0) if pyev.channels is not None: hit_channels = pyev.channels.hit.nonzero()[0].astype(np.uint32) if len(hit_channels) > 0: ROOT.fill_channels(self.ev, len(hit_channels), hit_channels, len(pyev.channels.hit), pyev.channels.t.astype(np.float32), pyev.channels.q.astype(np.float32), pyev.channels.flags.astype(np.uint32)) else: self.ev.nhit = 0 self.ev.nchannels = 0 self.ev.channels.resize(0) else: self.ev.nhit = 0 self.ev.nchannels = 0 self.ev.channels.resize(0) self.T.Fill()