qam.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
 
 
 
from ROOT import TCanvas, TFile
import argparse
 
 
class QAM():
 
    """
    QAM class.
    """
 
    def __init__(self, fname='input.root'):
        """
        call constructor of base class, required.
        """
 
        super(QAM, self).__init__()
 
        
        self.ff = TFile(fname)
        
        self.histHelixhistHelix = {'ndf': {'all': self.ff.Get('h17;1'),
                                  'up': self.ff.Get('h00;1'),
                                  'down': self.ff.Get('h01;1')},
                          'pval': {'all': self.ff.Get('h18;1'),
                                   'up': self.ff.Get('h02;1'),
                                   'down': self.ff.Get('h03;1')},
                          'd0': {'all': self.ff.Get('h11;1'),
                                 'up': self.ff.Get('h04;1'),
                                 'down': self.ff.Get('h05;1')},
                          'phi0': {'all': self.ff.Get('h12;1'),
                                   'up': self.ff.Get('h06;1'),
                                   'down': self.ff.Get('h07;1')},
                          'omega': {'all': self.ff.Get('h13;1'),
                                    'up': self.ff.Get('h08;1'),
                                    'down': self.ff.Get('h09;1')},
                          'z0': {'all': self.ff.Get('h14;1'),
                                 'up': self.ff.Get('h0a;1'),
                                 'down': self.ff.Get('h0b;1')},
                          'tanl': {'all': self.ff.Get('h15;1'),
                                   'up': self.ff.Get('h0c;1'),
                                   'down': self.ff.Get('h0d;1')},
                          'pt': {'all': self.ff.Get('h16;1'),
                                 'up': self.ff.Get('h0e;1'),
                                 'down': self.ff.Get('h0f;1')}
                          }
 
        
        self.histResohistReso = {'d0': self.ff.Get('h1;1'),
                         'phi0': self.ff.Get('h2;1'),
                         'omega': self.ff.Get('h3;1'),
                         'z0': self.ff.Get('h4;1'),
                         'tanl': self.ff.Get('h5;1'),
                         'pt': self.ff.Get('h6;1')}
 
        
        self.histPullhistPull = {'d0': self.ff.Get('h21;1'),
                         'phi0': self.ff.Get('h22;1'),
                         'omega': self.ff.Get('h23;1'),
                         'z0': self.ff.Get('h24;1'),
                         'tanl': self.ff.Get('h25;1')}
 
        
        self.graphPtgraphPt = {'d0': self.ff.Get('Graph;2'),
                        'phi0': self.ff.Get('Graph;3'),
                        'omega': self.ff.Get('Graph;4'),
                        'z0': self.ff.Get('Graph;5'),
                        'pt': self.ff.Get('Graph;1')}
        
        self.canvascanvas = TCanvas("canvas", "canvas", 800, 800)
        
        self.ndivndiv = 1
        
        self.indexindex = 1
 
    def pull(self, key='d0'):
        '''
        Plot pull distribution
        '''
        self.canvascanvas.cd(self.indexindex)
        self.histPullhistPull[key].Draw()
        self.canvascanvas.Update()
        self.indexindex = self.indexindex + 1
 
    def getMean(self, key='pt'):
        '''
        Getter for mean of helix parameter.
        '''
        h = self.histResohistReso[key]
        m = h.GetMean()
        dm = h.GetMeanError()
        return [m, dm]
 
    def getRms(self, key='pt'):
        '''
        Getter for rms of helix parameter.
        '''
        h = self.histResohistReso[key]
        s = h.GetRMS()
        ds = h.GetRMSError()
        return [s, ds]
 
    def graph(self, key='pt'):
        '''
        Plot graph of resolution as a function of Pt.
        '''
        self.canvascanvas.cd(self.indexindex)
        self.graphPtgraphPt[key].Draw('AP')
        self.canvascanvas.Update()
        self.indexindex = self.indexindex + 1
 
    def reso(self, key='pt'):
        '''
        Plot resolution histogram.
        '''
        self.canvascanvas.cd(self.indexindex)
        self.histResohistReso[key].Draw()
        self.canvascanvas.Update()
        self.indexindex = self.indexindex + 1
 
    def draw(self, key='pt', option='all', gopt=''):
        '''
        Plot histogram of helix parameters etc..
        '''
        self.canvascanvas.cd(self.indexindex)
        self.histHelixhistHelix[key][option].Draw(gopt)
        self.canvascanvas.Update()
        self.indexindex = self.indexindex + 1
 
    def print(self, fname='test.png'):
        """
        Print the current plot.
        """
        self.canvascanvas.Print(fname)
 
    def div(self, i=1, j=1):
        """
        Divide Tcanvas by (i,j).
        """
        self.canvascanvas.Clear()
        self.canvascanvas.Divide(i, j)
        self.ndivndiv = i * j
        self.indexindex = 1
 
 
if __name__ == "__main__":
    # Make the parameters accessible form the outside.
 
    parser = argparse.ArgumentParser()
    parser.add_argument('input', help='Input file to be processed (unpacked CDC data).')
    args = parser.parse_args()
    qam = QAM(args.input)