histograms.py

Go to the documentation of this file.

#!/usr/bin/env python
'''
Make some generator level histograms.

Eg

nuwa.py --output-stats="{'file1':'stats.root'}" \
        [-G $XMLDETDESCROOT/DDDB/dayabay-radslabs.xml]
        -n100 -o blah.root \
        -m 'GenTools.Opts --helper=radslab' \
        -m GenTools.histograms


<bv@bnl.gov> Sat Jul 16 02:04:47 2011
'''

import PyCintex
Gaudi = PyCintex.makeNamespace('Gaudi')

from DybPython.DybPythonAlg import DybPythonAlg
from DybPython.Util import hists
from GaudiPython import SUCCESS, FAILURE, loaddict

from DybPython.Util import irange
import GaudiKernel.SystemOfUnits as units
import math

from ROOT import TH2F

class HistBase(DybPythonAlg):
    '''
    Basic histogram handling common to all algs
    '''

    class PerSite:
        def __init__(self,nick,site_de,ad_des,keeper):
            self.nick = nick
            self.site_de = site_de
            self.ad_de = ad_des
            self.keeper = keeper
            return
        pass

    def __init__(self,name,prefix):
        DybPythonAlg.__init__(self,name)
        self.prefix = prefix
        self.sites = []         # PerSite infor
        return

    def initialize(self):
        print 'Calling DybPythonAlg.initialize(self)'
        sc = DybPythonAlg.initialize(self)
        if sc.isFailure(): return sc

        print 'HistBase.initialize()'

        statSvc = self.svc('IStatisticsSvc','StatisticsSvc')

        for nads,nick in [(2,'db'),(2,'la'),(4,'far')]:
            keeper = hists.Keeper(statSvc,self.prefix+nick)
            site_de_name = '/dd/Structure/Sites/%s-rock' % nick
            site_de = self.getDet(site_de_name)
            ad_des = []
            for iad in range(nads):
                iad += 1
                ad_de_name = '/dd/Structure/AD/%s-oil%d' % (nick,iad)
                ad_de = self.getDet(ad_de_name)
                ad_des.append(ad_de)
                continue                
            self.sites.append(HistBase.PerSite(nick,site_de,ad_des,keeper))
            continue

        return SUCCESS

    def site_ad_by_point(self, gaudi_point):
        '''
        Return tuple of (site,ad_de) that contains the given global gaudi point.
        If not in a site, both are None.  If not in an AD ad_de is None.
        '''

        xyz  = [gaudi_point.x(),gaudi_point.y(),gaudi_point.z()]
        xyz_m = map(lambda x: x/units.meter, xyz)
        for site in self.sites:
            site_gi = site.site_de.geometry()
            local_point = site_gi.toLocal(gaudi_point)
            if not site_gi.isInside(gaudi_point):
                continue
            for ad_de in site.ad_de:
                if ad_de.geometry().isInside(gaudi_point):
                    return (site,ad_de)
                continue
            return (site,None)
            continue
        return (None,None)

class HistVerts(HistBase):
    def __init__(self,name='GenVertex',prefix='/file1/gen'):
        HistBase.__init__(self,name,prefix)
        return

    def initialize(self):
        print 'Calling HistBase.initialize(self)'
        sc = HistBase.initialize(self)
        if sc.isFailure(): return sc

        binsPerMeter = 10
        for site in self.sites:
            k = site.keeper
            nick = site.nick
            
            #print 'Booking histograms for "%s"' % nick
            # Pool
            pool_height = 11.0
            pool_xsize = 17.0
            pool_ysize = 11.0
            nads = 2
            if nick == 'far':
                pool_ysize = 17.0
                nads = 4
            k.poolXY = TH2F('poolXY','%s Pool X-Y Vertices (m)'%nick.upper(),
                            pool_xsize*binsPerMeter,-pool_xsize/2,pool_xsize/2,
                            pool_ysize*binsPerMeter,-pool_ysize/2,pool_ysize/2)

            k.poolXZ = TH2F('poolXZ','%s Pool X-Z Vertices (m)'%nick.upper(),
                            pool_xsize*binsPerMeter,-pool_xsize/2,pool_xsize/2,
                            pool_height*binsPerMeter,0,-pool_height)

            k.poolYZ = TH2F('poolYZ','%s Pool Y-Z Vertices (m)'%nick.upper(),
                            pool_ysize*binsPerMeter,-pool_ysize/2,pool_ysize/2,
                            pool_height*binsPerMeter,0,-pool_height)

            # ADs
            ad_radius = 5.0
            ad_height = 5.0
            for iad in range(nads):
                iad += 1
                ad = 'AD%d' % iad
                #print 'Booking AD #%d/%d' % (iad,nads)

                nam = ad+'_XY'
                nbinsx = 2*ad_radius*binsPerMeter
                nbinsy = 2*ad_radius*binsPerMeter
                #print 'TH2F %s with %d * %d = %d' % (nam,nbinsx,nbinsy,nbinsx*nbinsy)
                h1 = TH2F(nam,'%s %s X-Y Vertices (m)'%(nick.upper(), ad.upper()),
                         nbinsx,-ad_radius,ad_radius,
                         nbinsy,-ad_radius,ad_radius)
                k.__setattr__(nam,h1)

                nam = ad+'_R2Z'
                nbinsx = 2*ad_radius**2*binsPerMeter/10
                nbinsy = ad_height*binsPerMeter
                #print 'TH2F %s with %d * %d = %d' % (nam,nbinsx,nbinsy,nbinsx*nbinsy)
                h2 = TH2F(nam,'%s %s R2Z Vertices (m2 v m)'%(nick.upper(), ad.upper()),
                         nbinsx,0,ad_radius**2,
                         nbinsy,-ad_height/2.0,ad_height/2.0)
                k.__setattr__(nam,h2)

                continue        # loop over ADs
            continue            # loop over site
        return SUCCESS          # initialize()

    def execute(self):
        evt = self.evtSvc()
        path = '/Event/Gen/GenHeader'
        gh = evt[path]
        event = gh.event()

        if event is None:
            self.error('Failed to get GenHeader from "%s"' % path)
            return FAILURE


        for vtx in irange(event.vertices_begin(),event.vertices_end()):
            
            pos = vtx.position()
            gpoint = Gaudi.XYZPoint(pos.x(),pos.y(),pos.z())
            site,ad = self.site_ad_by_point(gpoint)
            if not site:
                self.warning('Found vertex not in site: (%s,%s,%s) (meter)' % \
                                 (pos.x()/units.meter,
                                  pos.y()/units.meter,
                                  pos.z()/units.meter))
                continue

            self.fill_vertex(gpoint,site,ad)
            continue
        return SUCCESS

    def fill_vertex(self,gpoint,site,ad_de = None):
        '''
        Fill vertex histograms using given gpoint using the site and ad_de
        '''
        k = site.keeper
        site_point = site.site_de.geometry().toLocal(gpoint)
        sx_m = site_point.x()/units.meter
        sy_m = site_point.y()/units.meter
        sz_m = site_point.z()/units.meter
        k.poolXY.Fill(sx_m,sy_m)
        k.poolXZ.Fill(sx_m,sz_m)
        k.poolYZ.Fill(sy_m,sz_m)

        if not ad_de: return

        ad_point = ad_de.geometry().toLocal(gpoint)

        adx_m = ad_point.x()/units.meter
        ady_m = ad_point.y()/units.meter
        adz_m = ad_point.z()/units.meter

        iad = site.ad_de.index(ad_de) + 1
        k.__getattr__('AD%d_XY'%iad).Fill(adx_m,ady_m)
        k.__getattr__('AD%d_R2Z'%iad).Fill(adx_m**2+ady_m**2,adz_m)
        return

def configure(argv = None):
    return

def run(app):
    app.addAlgorithm(HistVerts())
    return