Ejemplo n.º 1
0
    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)
Ejemplo n.º 2
0
    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)
Ejemplo n.º 3
0
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)
Ejemplo n.º 4
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)
Ejemplo n.º 5
0
    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)
Ejemplo n.º 6
0
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
Ejemplo n.º 7
0
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]
Ejemplo n.º 8
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()
Ejemplo n.º 9
0
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
Ejemplo n.º 10
0
 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()
Ejemplo n.º 11
0
    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)
Ejemplo n.º 12
0
 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)
Ejemplo n.º 13
0
 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()
Ejemplo n.º 14
0
    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()
Ejemplo n.º 15
0
    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()
Ejemplo n.º 16
0
 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
Ejemplo n.º 17
0
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
Ejemplo n.º 18
0
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
Ejemplo n.º 19
0
    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()
Ejemplo n.º 20
0
    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()
Ejemplo n.º 21
0
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
Ejemplo n.º 22
0
 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
Ejemplo n.º 23
0
    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()