Ejemplo n.º 1
0
 def MakeROOTPrimitive(self):
     result = TList()
     result.SetName(self.triggername)
     self.__eventcount.SetName("events")
     self.__projectedSpectrum.SetName("spectrum")
     result.Add(self.__eventcount)
     result.Add(self.__projectedSpectrum)
     return result
Ejemplo n.º 2
0
 def MakeROOTPrimitive(self, name):
     result = TList()
     result.SetName(name)
     if self.__MCtruth:
         result.Add(self.__MCtruth)
     for hist in self.__triggers.itervalues():
         result.Add(hist)
     return result
Ejemplo n.º 3
0
 def ROOTify(self):
     """
     Create simple primitive ROOT object structure
     """
     outputlist = TList()
     outputlist.SetName("JetSpectraPt%03d" % (int(self.__jetpt)))
     for spec in self.__spectra:
         outputlist.Add(spec)
     return outputlist
Ejemplo n.º 4
0
def dict2rootmap(dict):
    '''assumes keys are strings, and vals are lists of strings'''
    map = TMap()
    for (key, entries) in dict.iteritems():
        rootkey = TObjString(key)
        rootval = TList()
        for entry in entries:
            rootval.Add(TObjString(entry))
        map.Add(rootkey, rootval)
    return map
Ejemplo n.º 5
0
def DrawUL(gUL):
    alist = TList()
    for i in range(gUL.GetN()):
        a = TArrow(gUL.GetX()[i],
                   gUL.GetY()[i],
                   gUL.GetX()[i],
                   gUL.GetY()[i] / 3., 0.01, ">")
        a.SetLineColor(gUL.GetLineColor())
        a.Draw()
        alist.Add(a)

    return alist
Ejemplo n.º 6
0
def hadd(input_file_names, result_file_name):
    """
    The top function for merging files.

    :param str result_file_name: path to the output file
    :param list[str] source_list: list of paths to the input files
    """
    target = TFile.Open(result_file_name, "RECREATE")
    file_list = TList()
    for file_name in input_file_names:
        file_list.Add(TFile.Open(file_name))
    merge_root_file(target, file_list)
Ejemplo n.º 7
0
 def GetRootPrimitive(self, name):
     """
     Convert object to a root primitive (so that it can be wirtten to a rootfile)
     """
     result = TList()
     result.SetName(name)
     events = self._events.GetROOTHisto()
     events.SetName("events")
     result.Add(events)
     spectrum = self.__spectrum.GetRootPrimitive()
     spectrum.SetName("spectrum")
     result.Add(spectrum)
     return result
Ejemplo n.º 8
0
    def WriteToFile(self, outputname):
        outputlists = []
        for categ in self.__categories.keys():
            mylist = TList()
            mylist.SetName(categ)
            for entry in self.__categories[categ].itervalues():
                mylist.Add(entry)
            outputlists.append(mylist)

        outputfile = TFile(outputname, "RECREATE")
        outputfile.cd()
        for myobject in outputlists:
            myobject.Write(myobject.GetName(), TObject.kSingleKey)
        outputfile.Close()
Ejemplo n.º 9
0
def CreateData(fileList, filename):
    pathList = fileList  #['test1.root', 'test2.root']
    signalTreeList = TList()
    backgroundTreeList = TList()
    outputFile = TFile(filename, 'recreate')
    pyfilelist = []
    pySignalTreelist = []
    pyBackgroundTreelist = []

    for path in pathList:
        print("Path", path)
        inputFile = TFile(path, 'read')
        pyfilelist.append(inputFile)  # Make this TFile survive the loop!
        inputSignalTree = inputFile.Get('Signal')
        inputBackgroundTree = inputFile.Get('Background')
        pySignalTreelist.append(
            inputSignalTree)  # Make this TTree survive the loop!
        pyBackgroundTreelist.append(
            inputBackgroundTree)  # Make this TTree survive the loop!
        outputSignalTree = inputSignalTree.CloneTree(
        )  #instead of extensive processing
        outputBackgroundTree = inputBackgroundTree.CloneTree(
        )  #instead of extensive processing
        signalTreeList.Add(inputSignalTree)
        backgroundTreeList.Add(inputBackgroundTree)

    outputFile.cd()
    outputSignalTree = TTree.MergeTrees(signalTreeList)
    outputBackgroundTree = TTree.MergeTrees(backgroundTreeList)
    outputFile.Write()
    outputFile.Close()
Ejemplo n.º 10
0
 def ProcessBin(self, mybin):
     results = BinContent()
     bindata = self._inputcol.GetData(mybin)
     for trigger in [
             "MinBias", "EMCJHigh", "EMCJLow", "EMCGHigh", "EMCGLow"
     ]:
         print "Doing trigger %s" % (trigger)
         jetcont = bindata.GetData(trigger).GetJetContainer()
         outputcont = TList()
         outputcont.SetName(trigger)
         for jetpt in jetcont.GetListOfJetPts():
             print "Inspecting jet pt %f" % (jetpt)
             jetdat = JetData(jetpt, trigger)
             projectedRec = jetcont.MakeProjectionRecKine(
                 jetpt, 0, "projectedPtRec")
             projectedMC = jetcont.MakeProjectionMCKine(
                 jetpt, 0, "projectedPtMC")
             jetdat.AddSpectrum(projectedRec, False)
             jetdat.AddSpectrum(projectedMC, True)
             outputcont.Add(jetdat.ROOTify())
         results.AddTrigger(trigger, outputcont)
     return results
Ejemplo n.º 11
0
def MergeHists(listOfHists):
    '''
    Method to merge histos

    Parameters
    ----------
    - listOfHists: python list of histos

    Returns
    ----------
    - hMerged: merged histo

    '''
    listMerge = TList()
    for iHist, hist in enumerate(listOfHists):
        if iHist == 0:
            hMerged = hist.Clone()
        else:
            listMerge.Add(hist)
    hMerged.Merge(listMerge)

    return hMerged
 def Add(self, other):
     if self.__outputtype == "TDirectoryFile":
         othercontent = other.content()
         logging.info("self scaled %s, has content: %s" %
                      (self.name(), "Yes" if self.hascontent() else "No"))
         logging.info("other scaled %s, has content: %s" %
                      (other.name(), "Yes" if other.hascontent() else "No"))
         for c in self.__content.GetListOfKeys():
             for o in othercontent.GetListOfKeys():
                 if c.GetName() == o.GetName():
                     others = TList()
                     others.append(o.ReadObj())
                     c.ReadObj().Merge(others)
     else:
         others = TList()
         others.append(other.content())
         self.__content.Merge(others)
Ejemplo n.º 13
0
 def MakeOutputList(self):
     result = TList()
     result.SetName("Results%d" % (self.__limit))
     rawspectrum = deepcopy(self.__datamodel.GetRawSpectrum())
     rawspectrum.SetName("rawspectrum")
     fitcurve = deepcopy(self.__datamodel.GetFitCurve())
     fitcurve.SetName("parameterisation")
     ptreach = self.__points.MakeLimitCurve(None, direction="central")
     ptreach.SetName("ptreach")
     result.Add(rawspectrum)
     result.Add(fitcurve)
     result.Add(ptreach)
     return result
Ejemplo n.º 14
0
import cuts as ct
import ROOT as rt
from ROOT import TFile, TTree, TCanvas, TH1F, TList, TH2F, TH3F, TMath, TF1, TStyle, gStyle, TRefArray, TClonesArray, TObjArray, gPad, TPaveText, TLegend, TString, TObject, gROOT, TFormula, TEllipse, TDirectory, TLorentzVector
from ROOT import TMath as mt
from datetime import date
# this code creates the histograms from the analysis results and also applies some additional cuts
# however I am still new to the pyroot so I might be defining histogram more than I needed
# my goal was to create just one set of histogram and clear it after every time I store them in root file for different level of cuts
# Different directory are created to store all the histogram at different level of cuts

#Opening the file>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
filename = "PolarizationSample"
infilename = filename + ".root"
file = TFile.Open(infilename)
#file = TFile.Open("AnalysisResults.root")
list = TList()
list = file.Get("Polar")
tree = list.FindObject("result")
daughter = TLorentzVector()
#daughter =
#tree2 = file.Get("tree2")

#tree2.AddFriend(tree)
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
text = TPaveText(.05, .1, .95, .8)
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

#this is to use to run in batch mode  to avoid printing all the plots
gROOT.SetBatch(rt.kTRUE)
#this create the master canvas color and the entries to print
#gStyle.SetCanvasColor(22)
Ejemplo n.º 15
0
def main(options, args):
    gROOT.Reset()

    #load our super special Polarization PDF
    gROOT.ProcessLine('.L RooPolarizationPdf.cxx+')
    gROOT.ProcessLine('.L RooPolarizationConstraint.cxx+')

    #setup integration
    intConf = ROOT.RooAbsReal.defaultIntegratorConfig()
    #intConf.Print('v')
    #intConf.method1D().setLabel('RooAdaptiveGaussKronrodIntegrator1D')

    output = TFile.Open('PolVsPt' + options.fitFrame + '.root', 'recreate')
    output.cd()

    lth_pt_p = TGraphAsymmErrors()
    lphi_pt_p = TGraphAsymmErrors()
    lthphi_pt_p = TGraphAsymmErrors()
    lthtilde_pt_p = TGraphAsymmErrors()
    f_pt_p = TGraphAsymmErrors()

    lth_pt_p_list = TList()
    lphi_pt_p_list = TList()
    lthphi_pt_p_list = TList()

    lth_pt_p.SetName('lth_pt_p')
    lphi_pt_p.SetName('lphi_pt_p')
    lthphi_pt_p.SetName('lthphi_pt_p')
    lthtilde_pt_p.SetName('lthtilde_pt_p')
    f_pt_p.SetName('f_pt_p')

    lth_pt_np = TGraphAsymmErrors()
    lphi_pt_np = TGraphAsymmErrors()
    lthphi_pt_np = TGraphAsymmErrors()
    lthtilde_pt_np = TGraphAsymmErrors()
    f_pt_np = TGraphAsymmErrors()

    lth_pt_np_list = TList()
    lphi_pt_np_list = TList()
    lthphi_pt_np_list = TList()

    lth_pt_np.SetName('lth_pt_np')
    lphi_pt_np.SetName('lphi_pt_np')
    lthphi_pt_np.SetName('lthphi_pt_np')
    lthtilde_pt_np.SetName('lthtilde_pt_np')
    f_pt_np.SetName('f_pt_np')

    total_nonzero_points = 0

    for f in args:

        infile = TFile.Open(f)

        #ws = output.Get(f[:f.find('.root')].split('-')[0])
        frame = f[:f.find('.root')].split('-')[1]
        rapBin = f[:f.find('.root')].split('-')[2].split('_')[0]
        ptBin = f[:f.find('.root')].split('-')[2].split('_')[1]

        lth_p = infile.Get('lth_p_rap' + rapBin)

        for point in range(lth_p.GetN()):
            x = ROOT.Double(0)
            y = ROOT.Double(0)
            exl = float(0)
            exh = float(0)
            eyl = float(0)
            eyh = float(0)
            lth_p.GetPoint(point, x, y)
            if x != 0.0 and y != 0.0:

                infile.Get('lth_p_rap' + rapBin).GetPoint(point, x, y)
                lth_pt_p.SetPoint(total_nonzero_points, x, y)

                infile.Get('lphi_p_rap' + rapBin).GetPoint(point, x, y)
                lphi_pt_p.SetPoint(total_nonzero_points, x, y)

                infile.Get('lthphi_p_rap' + rapBin).GetPoint(point, x, y)
                lthphi_pt_p.SetPoint(total_nonzero_points, x, y)

                infile.Get('lthtilde_p_rap' + rapBin).GetPoint(point, x, y)
                lthtilde_pt_p.SetPoint(total_nonzero_points, x, y)

                infile.Get('f_p_rap' + rapBin).GetPoint(point, x, y)
                f_pt_p.SetPoint(total_nonzero_points, x, y)

                exh = infile.Get('lth_p_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('lth_p_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('lth_p_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lth_p_rap' + rapBin).GetErrorYlow(point)
                lth_pt_p.SetPointError(total_nonzero_points, exl, exh, eyl,
                                       eyh)

                exh = infile.Get('lphi_p_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('lphi_p_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('lphi_p_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lphi_p_rap' + rapBin).GetErrorYlow(point)
                lphi_pt_p.SetPointError(total_nonzero_points, exl, exh, eyl,
                                        eyh)

                exh = infile.Get('lthphi_p_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('lthphi_p_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('lthphi_p_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lthphi_p_rap' + rapBin).GetErrorYlow(point)
                lthphi_pt_p.SetPointError(total_nonzero_points, exl, exh, eyl,
                                          eyh)

                exh = infile.Get('lthtilde_p_rap' +
                                 rapBin).GetErrorXhigh(point)
                exl = infile.Get('lthtilde_p_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('lthtilde_p_rap' +
                                 rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lthtilde_p_rap' + rapBin).GetErrorYlow(point)
                lthtilde_pt_p.SetPointError(total_nonzero_points, exl, exh,
                                            eyl, eyh)

                exh = infile.Get('f_p_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('f_p_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('f_p_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('f_p_rap' + rapBin).GetErrorYlow(point)
                f_pt_p.SetPointError(total_nonzero_points, exl, exh, eyl, eyh)

                infile.Get('lth_np_rap' + rapBin).GetPoint(point, x, y)
                lth_pt_np.SetPoint(total_nonzero_points, x, y)

                infile.Get('lphi_np_rap' + rapBin).GetPoint(point, x, y)
                lphi_pt_np.SetPoint(total_nonzero_points, x, y)

                infile.Get('lthphi_np_rap' + rapBin).GetPoint(point, x, y)
                lthphi_pt_np.SetPoint(total_nonzero_points, x, y)

                infile.Get('lthtilde_np_rap' + rapBin).GetPoint(point, x, y)
                lthtilde_pt_np.SetPoint(total_nonzero_points, x, y)

                infile.Get('f_np_rap' + rapBin).GetPoint(point, x, y)
                f_pt_np.SetPoint(total_nonzero_points, x, y)

                exh = infile.Get('lth_np_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('lth_np_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('lth_np_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lth_np_rap' + rapBin).GetErrorYlow(point)
                lth_pt_np.SetPointError(total_nonzero_points, exl, exh, eyl,
                                        eyh)

                exh = infile.Get('lphi_np_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('lphi_np_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('lphi_np_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lphi_np_rap' + rapBin).GetErrorYlow(point)
                lphi_pt_np.SetPointError(total_nonzero_points, exl, exh, eyl,
                                         eyh)

                exh = infile.Get('lthphi_np_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('lthphi_np_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('lthphi_np_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lthphi_np_rap' + rapBin).GetErrorYlow(point)
                lthphi_pt_np.SetPointError(total_nonzero_points, exl, exh, eyl,
                                           eyh)

                exh = infile.Get('lthtilde_np_rap' +
                                 rapBin).GetErrorXhigh(point)
                exl = infile.Get('lthtilde_np_rap' +
                                 rapBin).GetErrorXlow(point)
                eyh = infile.Get('lthtilde_np_rap' +
                                 rapBin).GetErrorYhigh(point)
                eyl = infile.Get('lthtilde_np_rap' +
                                 rapBin).GetErrorYlow(point)
                lthtilde_pt_np.SetPointError(total_nonzero_points, exl, exh,
                                             eyl, eyh)

                exh = infile.Get('f_np_rap' + rapBin).GetErrorXhigh(point)
                exl = infile.Get('f_np_rap' + rapBin).GetErrorXlow(point)
                eyh = infile.Get('f_np_rap' + rapBin).GetErrorYhigh(point)
                eyl = infile.Get('f_np_rap' + rapBin).GetErrorYlow(point)
                f_pt_np.SetPointError(total_nonzero_points, exl, exh, eyl, eyh)

                total_nonzero_points += 1

    output.cd()

    lth_pt_p.Write()
    lphi_pt_p.Write()
    lthphi_pt_p.Write()
    lthtilde_pt_p.Write()
    f_pt_p.Write()

    lth_pt_np.Write()
    lphi_pt_np.Write()
    lthphi_pt_np.Write()
    lthtilde_pt_np.Write()
    f_pt_np.Write()

    output.Close()
Ejemplo n.º 16
0
    format(home),
    "{0}/scratch0/ntuple/background_analysis/PFAnalysis_QCDEM3080_4.root".
    format(home),
    "{0}/scratch0/ntuple/background_analysis/PFAnalysis_QCDEM3080_5.root".
    format(home),
    "{0}/scratch0/ntuple/background_analysis/PFAnalysis_QCDEM3080_6.root".
    format(home),
    "{0}/scratch0/ntuple/background_analysis/PFAnalysis_QCDEM3080_7.root".
    format(home),
    "{0}/scratch0/ntuple/background_analysis/PFAnalysis_QCDEM3080_8.root".
    format(home),
    "{0}/scratch0/ntuple/background_analysis/PFAnalysis_QCDEM3080_9.root".
    format(home),
]

tree_list = TList()
tree_name = "Summary/tree"
output_file_name = "~/scratch0/PFAnalysis_QCDEM3080_0.root"

files = [
    TFile(file_name) for file_name in file_list if os.path.exists(file_name)
]
print("{0} files found.".format(len(files)))
trees = [file.Get(tree_name) for file in files]
for tree in trees:
    if tree:
        tree_list.Add(tree)

merged_tree = TTree.MergeTrees(tree_list)
output_file = TFile(output_file_name, "recreate")
merged_tree.Write()
Ejemplo n.º 17
0
#!/usr/bin/python
import code
import ROOT as rt
from ROOT import TFile, TTree , TCanvas, TH1F, TList , TH2F ,TMath ,TF1, TStyle ,gStyle , TRefArray, TClonesArray, TObjArray, gPad , TLegend ,TString
file = TFile.Open("Incohernt_Phi_PID.root")
#list = file.Get("")
PhiMass = TCanvas("PhiMass","Phi_Mass",600,300,700,527)
Sigma_Radius = TCanvas("Sigma_Radius","Sigma_Radius",600,300)
SigmavsSigma = TCanvas("SigmavsSigma","SigmaVsSigama",600,300)
dEdX = TCanvas("dEdX","dEdX",600,300)
test = TCanvas("test","dEdX",600,300)
list = TList()
list = file.Get("UPCPhiTaskTest/UPCPhiWithTrigger")
tree = list.FindObject("scatterplot")
#MassPlot for Incoherent Candidates
HighPt_Mass = TH1F("HighPt_Mass","Mass Dist of Inco #Phi with TOF",100, 0 ,10)

HighPt_Mass.SetFillStyle(4050)
HighPt_Mass.SetYTitle("Number of Events")
HighPt_Mass.SetXTitle("Phi Mass(MeV/C^2)")
HighPt_Mass.SetTitleOffset(1.5, "Y")
HighPt_Mass.SetTitleOffset(1.3, "X")
HighPt_Mass.SetLineColor(rt.kGreen)
HighPt_Mass.SetMarkerStyle(rt.kCircle)
HighPt_Mass.SetLineWidth(2)
HighPt_Mass.SetOption("E1")
Pt_mass = HighPt_Mass.Clone()
Pt_mass.SetLineColor(rt.kRed)
Pt_mass.SetLineWidth(4)
#Mass Plot for Coherent Candidates
LowPt_Mass =  TH1F("LowPt_Mass","Mass Dist of Incoherent #Phi without TOF info",100, 0 ,2)
Ejemplo n.º 18
0
def recursiveMerge(target,
                   infile,
                   path='',
                   cache={'TOTALLUMI': 0},
                   cutflow=True):
    l = infile.GetDirectory(path)
    keys = l.GetListOfKeys()
    cycles = {}

    #print("keys in input file: \n\n{0}\n\n".format(keys.ls()))

    for entry in range(keys.GetEntries()):
        name = keys.At(entry).GetName() + ";" + str(keys.At(entry).GetCycle())
        if path:
            cachename = path + "/" + name
        else:
            cachename = name
        obj = l.Get(name)

        if type(obj) == TDirectoryFile:
            #print("TDirectory obj name: {0}".format(obj.GetName()))
            targetpath = keys.At(entry).GetName()
            if not target.Get(targetpath):
                target.mkdir(targetpath)
            recursiveMerge(target, infile, path + "/" + obj.GetName(), cache)
        elif type(obj) == TTree:
            #print("TTree obj name: {0} - cachename: {1} ".format(obj.GetName(), cachename))
            cyclename, cyclenumber = cachename.split(';')
            if cyclename in cycles: continue
            #print("cyclename: {0} - cyclenumber: {1}".format(cyclename, cyclenumber))
            cycles[cyclename] = cyclenumber
            if not cyclename in cache:
                #print("adding cyclename {0} to cache (via TTree::CloneTree())".format(cyclename))
                target.cd(path)
                cache[cyclename] = obj.CloneTree()
            else:
                objcached = cache[cyclename]
                col = TObjArray()
                col.Add(obj)
                #print("merging TTree obj to cached object")
                objcached.Merge(col)
        elif issubclass(obj.__class__, TH1):
            #print("TH1 obj name: {0}".format(obj.GetName()))
            if not cutflow and keys.At(entry).GetName() == "CutFlow":
                continue
            if not cachename in cache:
                target.cd(path)
                cache[cachename] = obj.Clone()
            else:
                objcached = cache[cachename]
                col = TObjArray()
                col.Add(obj)
                objcached.Merge(col)
        elif type(obj) == TObjString:
            #print("TObjString obj name: {0}".format(obj.GetName()))
            if obj:
                target.cd(path)
                objnew = TObjString(obj.GetString().Data())
                objnew.Write(keys.At(entry).GetName())
                cache['TOTALLUMI'] += 1
        elif issubclass(obj.__class__, TList):
            #print("TList obj name: {0}".format(obj.GetName()))
            if obj:
                target.cd(path)
                objnew = TList(obj)
                objnew.Write(keys.At(entry).GetName())  # not working...
        else:
            print "UNKNOWN OBJECT", name, "OF TYPE", type(obj)
Ejemplo n.º 19
0
 def GetROOTPrimitive(self):
     output = TList()
     output.SetName(self.__name)
     for rootobject in self.__rootobjects:
         output.Add(rootobject.GetROOTPrimitive())
     return output
Ejemplo n.º 20
0
def run_analysis(options,args):    
    tm = tree_manager()
    pu_weight_sum = 0.0
    selected_events = []
    metaInfoTrees = []
    nEvents_total = int(0)
    pwd = ROOT.gDirectory.GetPath()
    for kFile,input_file in enumerate(args):
        print
        print 'processing input file %i/%i: %s'%(kFile+1,
                                                 len(args),
                                                 input_file)
        in_file = TFile.Open(input_file,'read')
        ROOT.gDirectory.cd(pwd)
        
        leptonType = options.leptonType    
        
        tree = None
        nEvents_sample = 0    
        specific = None
        treeName = 'Ntuple'
        if not options.allBranches:    
            if leptonType == 'muon':                
                #specific = muonBranches+commonBranches
                mmg = in_file.Get('mmg')
                tree = mmg.Get('final').Get(treeName)
                metaInfoTrees.append(mmg.Get('metaInfo').CloneTree())
                nEvents_sample = mmg.Get('eventCount').GetBinContent(1)
                nEvents_total += mmg.Get('skimCounter').GetBinContent(1)
                mmg = None
            elif leptonType == 'electron':                 
                #specific = electronBranches+commonBranches
                eeg = in_file.Get('eeg')
                tree = eeg.Get('final').Get(treeName)
                metaInfoTrees.append(eeg.Get('metaInfo').CloneTree())
                nEvents_sample = eeg.Get('eventCount').GetBinContent(1)
                nEvents_total += eeg.Get('skimCounter').GetBinContent(1)
                eeg = None
            else:             
                raise Exception('invalid lepton type: %s'%options.leptonType)
                
        total_events = tree.GetEntriesFast()    
        tick_size = total_events/100.0
        if tick_size < 1: tick_size = 1.0
        
        #tm.importTree(treeName,tree,specific)

        
        #tm.cloneTree(treeName,'%s_zs'%treeName,specific)
        #tm.cloneTree(treeName,'%s_zgs'%treeName,specific)
        #tm.cloneTree(treeName,'%s_zgs_nosihih'%treeName,specific)
        
        #setup process dependent stuff
        cuts = setupCuts(options)
        #lepton_mass = lepton_masses[options.leptonType]
        #z_info = z_infos[options.leptonType]
        #lepton_info = lepton_infos[options.leptonType]

        #setup pu-reweighing
        #pu_weight = pu_weight_nominal 
        #if options.isSummer11:
        #    pu_weigt = pu_weight_summer11
        
        procWeight = 1.0
        if options.datType != 'data':
            procWeight = (options.crossSection *
                          run_lumi[options.leptonType][options.runYear]\
                                  [options.runType])

        # setup the corrector (this links the appropriate four momenta
        # into a common naming scheme
        correct = setup_corrections(options.runYear   , options.runType,
                                    options.leptonType, options.datType,
                                    options.leptonCor , options.gamCor,
                                    options.vanilla                       )
        
        ievent = long(0)        
        for event in tree:
            ievent+=1
            #print ievent,total_events,fmod(ievent/total_events,0.01)
            if( not (ievent+1)%int(tick_size) or
                ievent+1 == total_events ):  
                sys.stdout.write('\r%3.0f%s complete! (%i/%i)'%(((ievent+1)/
                                                                 tick_size),
                                                                '%',
                                                                ievent+1,
                                                                total_events))
                sys.stdout.flush()
            
            # setup the common event momentum
            # ell1 = lepton1, ell2 = lepton2
            # gam = photon, Z = dilepton, Zg = Z+photon        
            correct(event)
            
            run_idx = getRunIndex(event.run[0],options.runYear,
                                  options.runType,options.datType,
                                  options.leptonType)
            setattr(event,'procWeight',procWeight)
            setattr(event,'puWeight',1.0)
            if options.datType != 'data':            
                setattr(event,'eventFraction',float(ievent+1)/total_events)
                event.puWeight = pu_S10_CD_reweight(event.nTruePU[0])
                
            #selected_z = []
            #selected_pho_nosihih = []
            #selected_pho = []        
            #bad_leptons = []
            
            if options.datType == 'data':
                # kill run 170722
                # kill the obvious pile up combinatorial event
                if ( event.run[0] == 170722 or
                     (event.run[0]   == 166512 and
                      event.lumi[0] == 1678 and
                      event.evt[0] == 1822682238) ):
                    continue
        
            bestLLG = None
            bestZdiff = -1
            
            for i in range(event.N_PATFinalState):            

                cuts.getCutflow('trigger')(event,i)
                if options.exlProc and not cuts.getCutflow('trigger') :
                    continue        
            
                cuts.getCutflow('pho')(event,i)            
                cuts.getCutflow('leptons')(event,i) 
                if ( options.exlProc and
                     not cuts.getCutflow('leptons') and
                     not cuts.getCutflow('pho') ):
                    continue

                cuts.getCutflow('z')(event,i)
                if options.exlProc and not cuts.getCutflow('z'):
                    continue

                if ( cuts.getCutflow('trigger') and
                     cuts.getCutflow('leptons') and
                     cuts.getCutflow('pho') and
                     cuts.getCutflow('z') ):
                    thisZdiff = abs(Z_POLE-event.Z[i].M())
                    if( cuts.getCut('mindr')[0](event.ell1[i],
                                                event.ell2[i],
                                                event.gam[i]) and
                        (thisZdiff < bestZdiff or bestLLG is None) ):
                        bestZdiff = thisZdiff
                        bestLLG = i
                    
            #event object selection done        
            if options.exlProc and bestLLG is None:
                continue
             
            bestZ = bestLLG
        
            if bestLLG is not None:
                setattr(event,'ell1SF',1.0)
                setattr(event,'ell2SF',1.0)
                #if run_idx != -1:
                #    event.ell1SF = leptonSF_nominal(event.nGoodVtx,
                #                     getattr(event,lepton_info['pt'])[idx1],
                #                    getattr(event,lepton_info['eta'])[idx1],
                #                                    run_idx,
                #                                    options.leptonType)
                #    event.ell2SF = leptonSF_nominal(event.nGoodVtx,
                #                      getattr(event,lepton_info['pt'])[idx2],
                #                     getattr(event,lepton_info['eta'])[idx2],
                #                                    run_idx,
                #                                    options.leptonType)
                setattr(event,'bestZ',bestLLG)
                outTrees.bestZTree(event,tm)
                #tm.fillTree('%s_zs'%treeName,{})
            
            #bestPhoNoSihih = None
            #bestPhoNoSihihPt = -1
            #for idxph in selected_pho_nosihih:
            #    pho.SetPtEtaPhiM(event.phoCorEt[idxph],
            #                     event.phoEta[idxph],
            #                     event.phoPhi[idxph],
            #                     0.0)
            #    if bestPhoNoSihih is None or pho.Pt() > bestPhoNoSihihPt:
            #        if ( bestZ is not None and
            #             cuts.getCut('mindr')[0](event.bestZLeg1,
            #                                     event.bestZLeg2,
            #                                     pho) ):
            #            bestPhoNoSihih = idxph
            #            bestPhoNoSihihPt = pho.Pt()
            #        elif( bestZ is None ):
            #            bestPhoSihih = idxph
            #            bestPhoSihihPt = pho.Pt()
            
            bestPho = bestLLG #and bestPhoNoSihih is None: (below)
            if  options.exlProc and bestPho is None:
                continue

            if bestPho is not None:            
                setattr(event,'phoSF',1.0)
                #if run_idx != -1:
                #    event.phoSF = phoSF_nominal(event.nGoodVtx,
                #                                event.phoEt[bestPho],
                #                                event.phoEta[bestPho],
                #                                run_idx)            
                setattr(event,'bestPho',bestPho)
                
                #if bestPhoNoSihih is not None:            
                #    setattr(event,'phoNoSihihSF',event.phoSF)                
                #    setattr(event,'bestPhoNoSihihIdx',bestPho)
                #    setattr(event,'bestPhoNoSihih',pho)
                #elif bestPhoNoSihih is not None:
                #    pho.SetPtEtaPhiM(event.phoCorEt[bestPhoNoSihih],
                #                     event.phoEta[bestPhoNoSihih],
                #                     event.phoPhi[bestPhoNoSihih],
                #                     0.0)
                #    setattr(event,'phoNoSihihSF',1.0)
                #    if run_idx != -1:
                #        event.phoNoSihihSF = phoSF_nominal(event.nGoodVtx,
                #                                 event.phoEt[bestPhoNoSihih],
                #                                 event.phoEta[bestPhoNoSihih],
                #                                           run_idx)
                #    setattr(event,'bestPhoNoSihihIdx',bestPhoNoSihih)
                #    setattr(event,'bestPhoNoSihih',pho)
        
            if bestLLG is not None:
                #if bestPhoNoSihih is not None:
                #    thezg = event.bestZ + event.bestPhoNoSihih
                #    setattr(event,'bestZGNoSihih',thezg)
                #    outTrees.bestZGTreeNoSihih(event,tm)
                #    tm.fillTree('EventTree_zgs_nosihih',{})
                hzg_r94cat = hzg_4cat_r9based[leptonType](event,bestLLG)
                setattr(event,'bestZG_r94cat',hzg_r94cat)
                hzg_r94cat_mod = \
                               hzg_4cat_r9based_mod[leptonType](event,bestLLG)
                setattr(event,'bestZG_r94cat_mod',hzg_r94cat_mod)
                thezg = event.Zg[bestLLG]            
                selected_events.append((event.run[bestLLG],
                                        event.lumi[bestLLG],
                                        event.evt[bestLLG]))
                setattr(event,'bestZG',thezg)
                outTrees.bestZGTree(event,tm)
                pu_weight_sum += event.puWeight
                #tm.fillTree('%s_zgs'%treeName,{})
            tree = None
        in_file.Close()
        del in_file
        
    #make a nice file name
    input_file = args[0]
    nameparts = input_file[input_file.rfind('/')+1:]
    nameparts = nameparts.split('.')
    
    #output selected event numbers to file if needed
    print
    print 'Selected %i (%.2f) events after processing!'%(len(selected_events),
                                                         pu_weight_sum)

    try:
        os.makedirs(options.prefix)
        print 'Created prefix directory: %s'%options.prefix
    except os.error:
        print 'Prefix directory: %s already exists'%options.prefix
    
    if options.dumpSelectedEvents:
        #sort events
        selected_events.sort(key=lambda event: event[2])
        selected_events.sort(key=lambda event: event[1])
        selected_events.sort(key=lambda event: event[0])
        evf = open(options.prefix +
                   '%s_%s_%s_processed.%s'%('_'.join(nameparts[:-1]),
                                            options.runType,
                                            options.leptonType,
                                            'txt'),
                   'w')
        for ev in selected_events:
            evf.write('Run, Lumi, Event #:\t%i\t%i\t%i\n'%ev)
        evf.close()
        
    #push all of our output trees to a file
    outFileName = ''
    if options.vanilla:
        outFileName = '%s_%s_%s_%s_%s_processed.%s'\
                      %('_'.join(nameparts[:-1]),
                        options.runType,
                        options.datType,
                        options.leptonType,
                        'vanilla',
                        nameparts[-1])
    else:
        outFileName = '%s_%s_%s_%s_%s_%s_processed.%s'\
                      %('_'.join(nameparts[:-1]),
                        options.runType,
                        options.datType,
                        options.leptonType,
                        options.leptonCor,
                        options.gamCor,
                        nameparts[-1])    
    
    hEventCount = TH1I('eventCount','Total Events Processed',1,0,1)
    hEventCount.SetBinContent(1,nEvents_total)
    metaTList = TList()
    for tree in metaInfoTrees:
        metaTList.Add(tree)
    metaTree = TTree.MergeTrees(metaTList)    
    outf = TFile.Open(options.prefix + outFileName,'RECREATE')
    outf.cd()    
    metaTree.Write()    
    hEventCount.Write()
    tm.write()
    outf.Close()
    ROOT.gDirectory.cd(pwd)
Ejemplo n.º 21
0
def init_pdfs(fit, params, prefix, signal):
    '''
    Initialize RooFit PDFs from parameter dictionary

    :param MassFit fit: - MassFit instance
    :param dict params: - parameter dictionary
    :param str prefix: - pdf name prefix
    :param bool signal: - signal pdf
    :return: (pdf, components) with pdf - Composit pdf, components - list of component pdfs
    :rtype: tuple
    '''
    # Build recursive composit pdf:
    # M(x) = f_sig*S(x) + (1-f_sig){f_peak*P(x) + (1-f_peak)*B(x)}
    #
    # This avoids problems with sum(fractions) > 1
    # See chapter 3 of RooFit Users Manual
    #
    # Store all auto generated obejcts in TLists. Otherwise python will destroy
    # them as soon as this function ends

    _pdf = None  #: Composit PDF to return
    _components = TList()  #: List of component PDFs
    _components_all = TList()  #: List of component PDFs
    _fracs = TList()  #: List of fractions
    _funcs = TList()  #: List of formula variables
    _params = { prefix+"_"+fit.mass_name: fit.mass }
    _i = 1
    _n_comps = len(params["components"])
    for _comp in params["components"]:
        # Create parameters
        _comp_params = TList()
        _comp_constrs = TList()
        for _par in _comp["params"]:
            _name = prefix + "_" + _par["name"]
            if _name in _params:
                _comp_params.append(_params[_name])
            else:
                if "random" in _par and _par["random"] == True:
                    # Generate starting value and validity range
                    _min = random.uniform(_par["min"], _par["value"])
                    _max = random.uniform(_par["value"], _par["max"])
                    _value = random.uniform(_min, _max)
                    _par["random"] = False
                    # Store obtained values
                    _par["min"] = _min
                    _par["max"] = _max
                    _par["value"] = _value
                    _var = RooRealVar(
                            _name, _par["title"], _value,
                            _min, _max, _par["units"]
                            )
                else:
                    _var = RooRealVar(
                            _name, _par["title"], _par["value"],
                            _par["min"], _par["max"], _par["units"]
                            )
                fit.parameters.push_back(_var)
                _comp_params.Add(_var)
                _params[_name] = _var
                if "constraint" in _par:
                    _constr = RooGaussian(
                            _name+"_constr",
                            _par["title"]+"_constr",
                            _var,
                            RooFit.RooConst(_par["constraint"]["mean"]),
                            RooFit.RooConst(_par["constraint"]["sigma"])
                            )
                    fit.constraints.push_back(_constr)
                    _comp_constrs.Add(_constr)
        # Store mean and sigma of core part of the signal PDF
        if _i == 1 and signal == True:
            fit.sig_mass = _comp_params.At(0)
            fit.sig_sigma = _comp_params.At(1)
        _name = prefix+"_"+_comp["name"]
        # Initialize component
        if _comp["type"] not in supported_pdfs:
            raise Exception("Unknown PDF type: %s" % _comp["type"])
        if _comp["type"] == "RooAddPdf":
            (_sum_pdf, _sum_components, _sum_params) = init_pdfs(fit, _comp, _comp['name'], signal)
            _components_all.AddAll(_sum_components)
            _components.Add(_sum_pdf)
            _params.update(_sum_params)
        elif _comp["type"] == "RooGaussian":
            _components.Add(RooGaussian(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1))
                    )
        elif _comp["type"] == "RooParamGaussian":
            _v = RooFormulaVar(
                    _name+"_variable",
                    "",
                    "@0*@1",
                    RooArgList(
                        _comp_params.At(1),
                        _comp_params.At(2))
                    )
            _funcs.Add(_v)
            _components.Add(RooGaussian(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _v)
                    )
        elif _comp["type"] == "RooParamGaussianOffset":
            _m = RooFormulaVar(
                    _name+"_var_mean",
                    "",
                    "@0+@1",
                    RooArgList(
                        _comp_params.At(0),
                        _comp_params.At(1))
                    )
            _s = RooFormulaVar(
                    _name+"_var_sigma",
                    "",
                    "@0*@1",
                    RooArgList(
                        _comp_params.At(2),
                        _comp_params.At(3))
                    )
            _funcs.Add(_m)
            _funcs.Add(_s)
            _components.Add(RooGaussian(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _m,
                    _s)
                    )
        elif _comp["type"] == "RooPolynomial":
            if len(_comp_params):
                _components.Add(RooPolynomial(
                        _name,
                        _comp["title"],
                        fit.mass,
                        RooArgList(_comp_params))
                        )
            else:
                _components.Add(RooPolynomial(
                        _name,
                        _comp["title"],
                        fit.mass)
                        )
        elif _comp["type"] == "RooChebychev":
            _components.Add(RooChebychev(
                    _name,
                    _comp["title"],
                    fit.mass,
                    RooArgList(_comp_params))
                    )
        elif _comp["type"] == "RooExponential":
            _components.Add(RooExponential(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0))
                    )
        elif _comp["type"] == "RooVoigtian":
            _components.Add(RooVoigtian(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2))
                    )
        elif _comp["type"] == "RooCBShape":
            _components.Add(RooCBShape(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3))
                    )
        elif _comp["type"] == "RooParamCBShape":
            _v = RooFormulaVar(
                    _name+"_variable",
                    "",
                    "@0*@1",
                    RooArgList(
                        _comp_params.At(1),
                        _comp_params.At(2))
                    )
            _funcs.Add(_v)
            _components.Add(RooCBShape(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _v,
                    _comp_params.At(3),
                    _comp_params.At(4))
                    )
        elif _comp["type"] == "RooDoubleCBShape":
            _components.Add(RooDoubleCBShape(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3),
                    _comp_params.At(4),
                    _comp_params.At(5))
                    )
        elif _comp["type"] == "RooCutCBShape":
            _components.Add(RooCutCBShape(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3),
                    _comp_params.At(4))
                    )
        elif _comp["type"] == "RooArgusBG":
            _components.Add(RooArgusBG(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1))
                    )
        elif _comp["type"] == "RooRArgusBG":
            _v = RooFormulaVar(
                    _name+"_variable",
                    "",
                    "2*@0 - @1",
                    RooArgList(
                        _comp_params.At(0),
                        fit.mass)
                    )
            _funcs.Add(_v)
            _components.Add(RooArgusBG(
                    _name,
                    _comp["title"],
                    _v,
                    _comp_params.At(0),
                    _comp_params.At(1))
                    )
        elif _comp["type"] == "RooIpatia":
            _components.Add(RooIpatia(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3),
                    _comp_params.At(4),
                    _comp_params.At(5),
                    _comp_params.At(6))
                    )
        elif _comp["type"] == "RooIpatia2":
            _components.Add(RooIpatia2(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3),
                    _comp_params.At(4),
                    _comp_params.At(5),
                    _comp_params.At(6),
                    _comp_params.At(7),
                    _comp_params.At(8))
                    )
        elif _comp["type"] == "RooCassandra":
            _components.Add(RooCassandra(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3))
                    )
        elif _comp["type"] == "RooAsymCassandra":
            _components.Add(RooAsymCassandra(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3),
                    _comp_params.At(4))
                    )
        elif _comp["type"] == "RooDecay":
            _r = RooGaussModel(
                    _name+"_resolution",
                    _comp["title"]+"_resolution",
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(2)
                    )
            _components.Add(RooDecay(
                _name,
                _comp["title"],
                fit.mass,
                _comp_params.At(1),
                _r,
                RooDecay.DoubleSided)
                )
            _funcs.Add(_r)
        elif _comp["type"] == "RooDecayLeft":
            _r = RooGaussModel(
                    _name+"_resolution",
                    _comp["title"]+"_resolution",
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(2)
                    )
            _components.Add(RooDecay(
                _name,
                _comp["title"],
                fit.mass,
                _comp_params.At(1),
                _r,
                RooDecay.Flipped)
                )
            _funcs.Add(_r)
        elif _comp["type"] == "RooDecayRight":
            _r = RooGaussModel(
                    _name+"_resolution",
                    _comp["title"]+"_resolution",
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(2)
                    )
            _components.Add(RooDecay(
                _name,
                _comp["title"],
                fit.mass,
                _comp_params.At(1),
                _r,
                RooDecay.SingleSided)
                )
            _funcs.Add(_r)
        elif _comp["type"] == "RooExpGaussExp":
            _components.Add(RooExpGaussExp(
                    _name,
                    _comp["title"],
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(1),
                    _comp_params.At(2),
                    _comp_params.At(3))
                    )
        elif _comp["type"] == "RooHistPdf":
            if 'config' not in _comp:
                raise Exception("Missing 'config' section for RooHistPdf.")
            if 'file' not in _comp['config']:
                raise Exception("Missing 'file' key in the 'config' section for RooHistPdf.")
            if 'name' not in _comp['config']:
                raise Exception("Missing 'name' key in the 'config' section for RooHistPdf.")
            _f = TFile(_comp['config']['file'])
            _pdf = _f.Get(_comp['config']['name'])
            _components.Add(_pdf)
        elif _comp["type"] == "Conv_RooDecay_RooBreitWigner":
            _r1 = RooGaussModel(
                    _name+"_resolution_1",
                    _comp["title"]+"_resolution_1",
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(2)
                    )
            _f1 = RooDecay(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(1),
                _r1,
                RooDecay.DoubleSided
                )
            _f2 = RooBreitWigner(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(3),
                _comp_params.At(4)
                )
            _funcs.Add(_r1)
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooNumConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f1,
                _f2
                )
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooDecay_RooBreitWigner":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _r1 = RooGaussModel(
                    _name+"_resolution_1",
                    _comp["title"]+"_resolution_1",
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(2)
                    )
            _f1 = RooDecay(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(1),
                _r1,
                RooDecay.DoubleSided
                )
            _f2 = RooBreitWigner(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(3),
                _comp_params.At(4)
                )
            _funcs.Add(_r1)
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f1,
                _f2
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooDecay_RooCBShape":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _r1 = RooGaussModel(
                    _name+"_resolution_1",
                    _comp["title"]+"_resolution_1",
                    fit.mass,
                    _comp_params.At(0),
                    _comp_params.At(2)
                    )
            _f1 = RooDecay(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(1),
                _r1,
                RooDecay.DoubleSided
                )
            _f2 = RooCBShape(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(3),
                _comp_params.At(4),
                _comp_params.At(5),
                _comp_params.At(6)
                )
            _funcs.Add(_r1)
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f1,
                _f2
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooCassandra_RooGaussian":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _f1 = RooCassandra(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2)
                )
            _f2 = RooGaussian(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(3),
                _comp_params.At(4)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f1,
                _f2
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooAsymCassandra_RooGaussian":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _f1 = RooAsymCassandra(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2),
                _comp_params.At(3),
                _comp_params.At(4)
                )
            _f2 = RooGaussian(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(5),
                _comp_params.At(6)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            #_pdf = RooFFTConvPdf(
            #    _name,
            #    _comp["title"],
            #    fit.mass,
            #    _f1,
            #    _f2
            #    )
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f2,
                _f1
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooAsymCassandra_RooDoubleGaussian":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _v = RooFormulaVar(
                    _name+"_variable",
                    "",
                    "@0*@1",
                    RooArgList(
                        _comp_params.At(6),
                        _comp_params.At(7))
                    )
            _funcs.Add(_v)
            _f1 = RooAsymCassandra(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2),
                _comp_params.At(3),
                _comp_params.At(4)
                )
            _f2 = RooGaussian(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(5),
                _comp_params.At(6)
                )
            _f3 = RooGaussian(
                _name+"_comp_3",
                _comp["title"]+"_comp_3",
                fit.mass,
                _comp_params.At(5),
                _v
                )
            _f4 = RooAddPdf(
               _name,
               _comp["title"],
               RooArgList(_f2, _f3),
               RooArgList(_comp_params.At(8))
               )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _funcs.Add(_f3)
            _funcs.Add(_f4)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f4,
                _f1
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        elif _comp["type"] == "Conv_RooAsymCassandra_RooGaussian":
            _f1 = RooAsymCassandra(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2),
                _comp_params.At(3),
                _comp_params.At(4)
                )
            _f2 = RooGaussian(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(5),
                _comp_params.At(6)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooNumConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f2,
                _f1
                )
            # Does not work
            #_pdf.setConvolutionWindow(_comp_params.At(5), _comp_params.At(6), 6.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooAsymCassandra3_RooGaussian":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _f1 = RooAsymCassandra3(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2),
                _comp_params.At(3),
                _comp_params.At(4),
                _comp_params.At(5),
                _comp_params.At(6)
                )
            _f2 = RooGaussian(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(7),
                _comp_params.At(8)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f2,
                _f1
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooCassandra_RooCBShape":
            # Define sampling frequency
            fit.mass.setBins(50000,"fft") ;
            _f1 = RooCassandra(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2)
                )
            _f2 = RooCBShape(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(3),
                _comp_params.At(4),
                _comp_params.At(5),
                _comp_params.At(6)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f1,
                _f2
                )
            _pdf.setBufferFraction(100.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooCassandra3_RooCBShape":
            # Define sampling frequency
            fit.mass.setBins(50000,"fft") ;
            _f1 = RooCassandra3(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2),
                _comp_params.At(3)
                )
            _f2 = RooCBShape(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(4),
                _comp_params.At(5),
                _comp_params.At(6),
                _comp_params.At(7)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f1,
                _f2
                )
            _pdf.setBufferFraction(100.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooHistPdf_RooGaussian":
            # Load reference ROOT tree
            _ref_full_ch = TChain(_comp['config']['tree'])
            _ref_full_ch.Add(_comp['config']['files'])
            # Run selection
            _ref_ch = _ref_full_ch.CopyTree(_comp['config']['selection'])
            _bins = fit.mass.getBins()
            fit.mass.setBins(_comp['config']['bins'])
            # Create hist PDF
            _data = RooDataSet(
                    'ds_' + _name,
                    '',
                    RooArgSet(fit.mass),
                    RooFit.Import(_ref_ch)
                    )
            _hist = _data.binnedClone()
            _f1 = RooHistPdf(
                    name+'_histpdf',
                    _comp['title']+'_histpdf',
                    RooArgSet(fit.mass),
                    _hist,
                    2  # Order of interpolation function
                    )
            fit.mass.setBins(_bins)

            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _f2 = RooGaussModel(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1)
                )
            _funcs.Add(_hist)
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f2,
                _f1
                )
            _pdf.setBufferFraction(5.0)
            _components.Add(_pdf)
        elif _comp["type"] == "Conv_RooHistPdf_RooGaussian":
            if 'config' not in _comp:
                raise Exception("Missing 'config' section for RooHistPdf.")
            if 'file' not in _comp['config']:
                raise Exception("Missing 'file' key in the 'config' section for RooHistPdf.")
            if 'name' not in _comp['config']:
                raise Exception("Missing 'name' key in the 'config' section for RooHistPdf.")
            _file = TFile(_comp['config']['file'])
            _f1 = _file.Get(_comp['config']['name'])
            _f2 = RooGaussian(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooNumConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f2,
                _f1
                )
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_RooDoubleCBShape_RooGaussian":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _f1 = RooDoubleCBShape(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2),
                _comp_params.At(3),
                _comp_params.At(4),
                _comp_params.At(5)
                )
            _f2 = RooGaussian(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(6),
                _comp_params.At(7)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f1,
                _f2
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        elif _comp["type"] == "FFT_2xCBShape_RooGaussian":
            # Define sampling frequency
            fit.mass.setBins(10000,"fft") ;
            _cb_frac = RooRealVar(
                    _name+"_cb_frac",
                    _name+"_cb_frac",
                    0.1, 0.0, 1.0)
            _f1 = RooCBShape(
                _name+"_comp_1",
                _comp["title"]+"_comp_1",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(2),
                _comp_params.At(3)
                )
            _f2 = RooCBShape(
                _name+"_comp_2",
                _comp["title"]+"_comp_2",
                fit.mass,
                _comp_params.At(0),
                _comp_params.At(1),
                _comp_params.At(4),
                _comp_params.At(5)
                )
            _f3 = RooAddPdf(
                _name+"_comp_3",
                _comp["title"]+"_comp_3",
                RooArgList(_f1, _f2),
                RooArgList(_cb_frac),
                kTRUE)
            _f4 = RooGaussian(
                _name+"_comp_4",
                _comp["title"]+"_comp_4",
                fit.mass,
                _comp_params.At(6),
                _comp_params.At(7)
                )
            _funcs.Add(_f1)
            _funcs.Add(_f2)
            _funcs.Add(_f3)
            _funcs.Add(_f4)
            _funcs.Add(_cb_frac)
            _pdf = RooFFTConvPdf(
                _name,
                _comp["title"],
                fit.mass,
                _f3,
                _f4
                )
            _pdf.setBufferFraction(50.0)
            _components.Add(_pdf)
        else:
            raise Exception("Unknown PDF type: "+_comp["type"])

        if _i < _n_comps:
            _f_v = 0.1
            _f_min = 0.0
            _f_max = 1.0
            if "fraction" in _comp:
                _f_v = _comp["fraction"]["value"]
                _f_min = _comp["fraction"]["min"]
                _f_max = _comp["fraction"]["max"]
            _f = RooRealVar(
                    prefix+"_frac_"+_comp["name"],
                    prefix+" component "+_comp["name"]+" fraction",
                    _f_v, _f_min, _f_max)
            if "fraction" in _comp and "constraint" in _comp["fraction"]:
                _constr = RooGaussian(
                        _name+"_frac_constr",
                        _name+"_frac_constr",
                        _f,
                        RooFit.RooConst(_comp["fraction"]["constraint"]["mean"]),
                        RooFit.RooConst(_comp["fraction"]["constraint"]["sigma"])
                        )
                fit.constraints.push_back(_constr)
                _comp_constrs.Add(_constr)
            _fracs.Add(_f)

        _i += 1

    if _n_comps == 1:
        _pdf = _components.At(0)
    else:
        # Make sure we are creating a recursive composit PDF
        _pdf = RooAddPdf(
            prefix+"_pdf",
            prefix+" PDF",
            RooArgList(_components),
            RooArgList(_fracs),
            kTRUE)

    _components_all.AddAll(_components)
    return (_pdf, _components_all, _params)
Ejemplo n.º 22
0
def splitData():
    """ Sometimes this segfaults.  No idea why.  Just re-run and it's fine. """

    f1 = TFile("./data/malbek_data.root")
    tree0 = f1.Get("malbek_wrt")

    axPeaks = [1.739, 1.836, 2.307, 2.464]

    # declare as arrays s/t ROOT branches work.  access values with e.g. ene[0] = 1.
    ene0, wt0, shift = array('d', [0.]), array('d', [0.]), array('d', [0.])

    tree0.SetBranchAddress("energy_keV", ene0)
    tree0.SetBranchAddress("weight", wt0)

    tList = TList()
    tVec = [0., 0., 0., 0.]
    tmp = TFile("./data/malbek_splitData.root", "RECREATE")

    for i in range(len(axPeaks)):
        ene, wt = array('d', [0.]), array('d', [0.])
        shift = axPeaks[len(axPeaks) - 1] - axPeaks[i]
        tVec[i] = TTree("t%d" % i, "t%d" % i)
        tVec[i].Branch("energy_keV", ene, "energy_keV/D")
        tVec[i].Branch("weight", wt, "weight/D")
        for j in range(tree0.GetEntries()):
            tree0.GetEntry(j)
            ene[0], wt[0] = ene0[0], wt0[0]
            ene[0] += shift
            tVec[i].Fill()

        print "axPeaks: %d  tree %d - %d entries  shift %.3f keV" % (
            len(axPeaks), len(tVec), tVec[i].GetEntries(), shift)
        tList.Add(tVec[i])
        tVec[i].Write()

    tree1 = ROOT.TTree.MergeTrees(tList)
    tree1.SetName("mergeTree")
    tree1.SetTitle("mergeTree")
    print "Trees merged, with %d entries total." % tree1.GetEntries()
    tree1.Write()

    # -- create diagnostic shifted & merged spectrum --
    malbekExposure = 89.5
    binSize = 0.04
    eLo, eHi = 0.8, 5.
    nBins = int((eHi - eLo) / binSize + 0.5)

    h0 = H1D(tree0, nBins, eLo, eHi, "energy_keV", "weight")
    h0.SetLineColor(ROOT.kBlack)
    h0.SetLineWidth(2)
    h0.Scale(1. / malbekExposure)

    h1 = H1D(tVec[0], nBins, eLo, eHi, "energy_keV", "weight")
    h1.SetLineColorAlpha(ROOT.kBlue, 0.5)
    h1.Scale(1. / malbekExposure)

    h2 = H1D(tVec[1], nBins, eLo, eHi, "energy_keV", "weight")
    h2.SetLineColorAlpha(ROOT.kGreen, 0.5)
    h2.Scale(1. / malbekExposure)

    h3 = H1D(tVec[2], nBins, eLo, eHi, "energy_keV", "weight")
    h3.SetLineColorAlpha(ROOT.kMagenta, 0.5)
    h3.Scale(1. / malbekExposure)

    h4 = H1D(tVec[3], nBins, eLo, eHi, "energy_keV", "weight")
    h4.SetLineColor(ROOT.kCyan)
    h4.Scale(1. / malbekExposure)

    hAdd = TH1D("hAdd", "", nBins, eLo, eHi)
    hAdd.Add(h1)
    hAdd.Add(h2)
    hAdd.Add(h3)
    hAdd.Add(h4)
    hAdd.SetLineColor(ROOT.kBlack)
    hAdd.SetLineWidth(2)
    hAdd.SetMinimum(0)
    hAdd.SetMaximum(2)
    ymax = hAdd.GetMaximum()
    hAdd.GetXaxis().SetTitle("Energy (keV)")
    hAdd.GetYaxis().SetTitle("Counts / kg-d")

    c = TCanvas("c", "Bob Ross's Canvas", 1100, 800)
    hAdd.Draw("hist")
    h0.Draw("hist same")
    h1.Draw("hist same")
    h2.Draw("hist same")
    h3.Draw("hist same")
    h4.Draw("hist same")

    # Draw lines around the gaussian fit ROI, using 2 sigma of the energy resolution
    fitWin = 3 * getSigma(axPeaks[3])
    print "Malbek resolution at %.2f keV is %.2f" % (axPeaks[3],
                                                     getSigma(axPeaks[3]))
    print "Set 3-sigma fit region to %.2f - %.2f" % (axPeaks[3] - fitWin,
                                                     axPeaks[3] + fitWin)

    l1 = ROOT.TLine(axPeaks[3], 0., axPeaks[3], ymax)
    l1.SetLineColor(ROOT.kRed)
    l1.SetLineWidth(2)
    l1.Draw("same")

    l2 = ROOT.TLine(axPeaks[3] + fitWin, 0., axPeaks[3] + fitWin, ymax)
    l2.SetLineColorAlpha(ROOT.kRed, 0.5)
    l2.SetLineWidth(2)
    l2.Draw("same")

    l3 = ROOT.TLine(axPeaks[3] - fitWin, 0., axPeaks[3] - fitWin, ymax)
    l3.SetLineColorAlpha(ROOT.kRed, 0.5)
    l3.SetLineWidth(2)
    l3.Draw("same")

    leg = ROOT.TLegend(0.6, 0.6, 0.85, 0.85)
    leg.AddEntry(hAdd, "Shifted+Summed", "l")
    leg.AddEntry(h4, "peak-%.3f (unshifted)" % (axPeaks[3]), "l")
    leg.AddEntry(h3, "peak-%.3f" % (axPeaks[2]), "l")
    leg.AddEntry(h2, "peak-%.3f" % (axPeaks[1]), "l")
    leg.AddEntry(h1, "peak-%.3f" % (axPeaks[0]), "l")
    leg.AddEntry(l1, "3-sigma fit region", "l")
    leg.Draw("same")

    c.Print("./plots/shiftSpec.pdf")
Ejemplo n.º 23
0
def stackHistograms(t_plots, t_processList, l_histoList, l_legendList, l_textList, l_axisList, saveAs,l_ratioList,  saveString, tagText):

    ''' 
    PURPOSE:

    Function to create stacked plots and save them to the local disk.
    This function will:
        * Create a stacked plot
        * Create a ratio plot
        * Allow for Blinding
        * Save the canvas as requested
    '''
    # Release the variables from the lists
    h_st, location_suf, name, c_x, c_y, v_rebin, xMin, xMax, doBlinding, b_low, b_high = l_histoList
    l_st, l_x1, l_y1, l_x2, l_y2, l_textSize = l_legendList
    t_st, plotText, t_textSize, t_alignment, t_start, t_gap, t_x = l_textList
    a_st, xLabel, yLabel, xOffset, yOffset, a_labelSize, a_textSize, y_scale = l_axisList
    r_st, doRatio, r_yLabel, r_xOffset, r_yOffset, r_labelSize, r_textSize = l_ratioList
    len_plots = len(t_plots)

    plotText = plotText + tagText
    
    l_overlay = []
    l_data    = []
    sh        = THStack('sh', '')
    t_legend = []
    listMC = TList()
    for i in xrange(len_plots):
        process, location_pre, v_colour, v_scale, ovack, sigrouta, l_marker, sel = t_processList[i]
        histogram = rebin(t_plots[i][0], v_rebin)
        histogram.Scale(float(v_scale))
        histogram.SetLineWidth(1)
        histogram.SetLineColor(int(v_colour))
        l_legendLine = [t_plots[i][0], process, l_marker]
        t_legend.append(l_legendLine)
        
        if sigrouta == "DATA" and int(doBlinding):
            histogram = blindingPlots(histogram, b_low, b_high)
        if ovack == "STACK":
            if sigrouta == "BACKGROUND":
                listMC.Add(histogram)
            histogram.SetFillColor(int(v_colour))
            sh.Add(histogram)
        elif ovack == "OVERLAY" and sigrouta != "DATA":
            l_overlay.append(histogram)
        elif ovack == "OVERLAY" and sigrouta == "DATA":
            histogram.SetMarkerStyle(8)
            histogram.SetMarkerColor(1)
            l_data.append(histogram)
        else:
            print "\n%s or %s are not a valid options" % (sigrouta, ovack)
            sys.exit(1)
    mergeMC = histogram.Clone('mergeMC')
    mergeMC.Reset()
    mergeMC.Merge(listMC)

    gROOT.SetBatch(1)
    c = TCanvas('c', 'c', int(c_x), int(c_y))
    gStyle.SetOptStat(0)
    if int(doRatio):
        m_Pad = TPad('m_pad', 'm_pad', 0, 0.30, 1, 1)
        r_Pad = TPad('r_pad', 'r_pad', 0, 0, 1, 0.30)
        m_Pad.Draw()
        r_Pad.Draw()
        m_Pad.SetTopMargin(0.05)
        m_Pad.SetBottomMargin(0.01)
        r_Pad.SetTopMargin(0)
        r_Pad.SetBottomMargin(0.3)
        m_Pad.cd()
        m_Pad.SetTicks(1)
        m_Pad.Update()
    else:
        c.SetTicks(1)
        c.Update
    
    # Draw our plots we defined earlier. First the stacked plot, then any overlaid plots and finally data
    sh.Draw()
    if l_overlay:
        for i in xrange(len(l_overlay)):
            l_overlay[i].Draw('hist same')
    h_MCerr = getStatUncert(mergeMC)
    h_MCerr.SetMarkerSize(0)
    h_MCerr.SetFillColor(921)
    h_MCerr.SetFillStyle(3344)
    h_MCerr.Draw('e2s same')
    t_legend.append([h_MCerr, 'Stat Uncert.', 'f']) 

    if l_data:
        h_data = mergeData(l_data, doBlinding, b_low, b_high)
        h_data.Draw('epsames')
    
    ## Next up is to add the styling:
    ## 1, Axis labels
    ## 2. Text on the plot
    ## 3. Legend
    if not int(doRatio):
        setupAxes(sh, xLabel, yLabel, xOffset, 0.8, a_textSize, xMin, xMax, y_scale)
    else:
        sh.SetMinimum(0.001)
        setupAxes(sh, "", yLabel, -999, yOffset, a_textSize, xMin, xMax, y_scale)

    setupLegend(t_legend, l_x1, l_y1, l_x2, l_y2, l_textSize)
    setupTextOnPlot(plotText, t_textSize, t_alignment, t_x, t_start, t_gap)
    
    if int(doRatio):
        r_Pad.cd()
        r_Pad.SetTicks(1)
        r_Pad.Update()
        l_mc   = TList()
        l_data = TList()
        for i in xrange(len_plots):
            l_histogram = rebin(t_plots[i][0], v_rebin)
            process, location_pre, v_colour, v_scale, ovack, sigrouta, l_marker, tag = t_processList[i] 
            if sigrouta == "SIGNAL":
                l_mc.Add(l_histogram)
            elif sigrouta == "BACKGROUND":
                l_mc.Add(l_histogram)
            elif sigrouta == "DATA":
                l_data.Add(l_histogram)
            elif sigrouta == "OVERLAY":
                pass
            else:
                print "\nOption %s not available" % (sigrouta)

        mh_MC = l_histogram.Clone('mh_MC')
        mh_MC.Merge(l_mc)

        h_statMC = getRelError(mh_MC)
        h_statMC.SetMarkerSize(0)
        h_statMC.SetFillColor(921)
        h_statMC.SetFillStyle(3344)
        h_statMC.Draw('e2s')
        setupAxes(h_statMC, xLabel, r_yLabel, r_xOffset, r_yOffset, r_textSize, xMin, xMax, 1)
        h_statMC.GetYaxis().SetRangeUser(-0.38, 0.38)
       
        mh_data = l_histogram.Clone('mh_data')
        mh_data.Merge(l_data)
        h_dataMC = getRatio(mh_data, mh_MC, doBlinding, b_low, b_high)
        h_dataMC.Draw('epsames')
 
        h_line = TLine(float(xMin), 0, float(xMax), 0)
        h_line.SetLineWidth(2)
        h_line.SetLineColor(1)
        h_line.SetLineStyle(2)
        h_line.Draw('same')
    
    savePlots(c, saveString, saveAs)
Ejemplo n.º 24
0
    def fhadd(self, prefix="", force=False, verbose=False, slow=True):
        """ taken from https://root.cern.ch/phpBB3/viewtopic.php?t=14881
        This function will merge objects from a list of root files and write them    
        to a target root file. The target file is newly created and must not
        exist, or if -f ("force") is given, must not be one of the source files.
        
        IMPORTANT: It is required that all files have the same content!

        Fast but memory hungry alternative to ROOT's hadd.
        
        Arguments:

        target -- name of the target root file
        sources -- list of source root files
        classname -- restrict merging to objects inheriting from classname
        force -- overwrite target file if exists
        """

        target = prefix + self.Name + ".root"
        sources = [j.Output for j in self.Jobs]

        TH1.AddDirectory(False)
        # check if target file exists and exit if it does and not in force mode
        if not force and os.path.exists(target):
            raise RuntimeError("target file %s exists" % target)

        # open the target file
        print "fhadd Target file:", target
        outfile = TFile(target, "RECREATE")

        # open the seed file - contents is looked up from here
        seedfilename = sources[0]
        print "fhadd Source file 1", seedfilename
        seedfile = TFile(seedfilename)

        # get contents of seed file
        print "looping over seed file"
        contents = self.loop(seedfile)
        print "done %d objects are ready to be merged" % len(contents)
        if (verbose):
            for c in contents:
                print c

        # open remaining files
        otherfiles = []
        for n, f in enumerate(sources[1:]):
            print "fhadd Source file %d: %s" % (n + 2, f)
            otherfiles.append(TFile(f))

        # loop over contents and merge objects from other files to seed file objects
        for n, (path, hname) in enumerate(contents):

            print "fhadd Target object: %s" % os.path.join(path, hname)
            obj_path = os.path.join(path, hname)
            obj_ = seedfile.Get(obj_path[1:])

            outfile.cd('/')
            # create target directory structure
            for d in path.split('/')[1:]:
                directory = gDirectory.GetDirectory(d)
                if not directory:
                    gDirectory.mkdir(d).cd()
                else:
                    gDirectory.cd(d)
            obj = None
            if obj_.InheritsFrom("TTree"):
                obj = obj_.CloneTree()
            else:
                obj = obj_.Clone()

            # merge objects
            l = TList()
            for o in [of.Get(obj_path[1:]) for of in otherfiles]:
                l.Add(o)
            obj.Merge(l)

            # delete objects if in slow mode
            if slow:
                print "Deleting %d object(s)", l.GetEntries()
                l.Delete()

            # write object to target
            obj.Write(obj.GetName(), TObject.kOverwrite)

        print "Writing and closing file"

        # let ROOT forget about open files - prevents deletion of TKeys
        for f in [outfile, seedfile] + otherfiles:
            gROOT.GetListOfFiles().Remove(f)

        outfile.Write()
        outfile.Close()

        for f in [seedfile] + otherfiles:
            f.Close()
Ejemplo n.º 25
0
 def GetRootPrimitive(self, listname):
     """
     Make root primitives (for root IO)
     """
     result = TList()
     result.SetName(listname)
     tracklist = TList()
     tracklist.SetName("trackContainers")
     for name, tc in self.__trackContainers.iteritems():
         tracklist.Add(tc.GetRootPrimitive("trackcontainer_%s" % (name)))
     clusterlist = TList()
     clusterlist.SetName("clusterContainers")
     for name, cc in self.__clusterContainers.iteritems():
         clusterlist.Add(cc.GetRootPrimitive("clustercontainer_%s" %
                                             (name)))
     result.Add(tracklist)
     result.Add(clusterlist)
     return result
Ejemplo n.º 26
0
    def test2Lists(self):
        """Test list/TList behavior and compatibility"""

        l = TList()
        l.Add(TObjString('a'))
        l.Add(TObjString('b'))
        l.Add(TObjString('c'))
        l.Add(TObjString('d'))
        l.Add(TObjString('e'))
        l.Add(TObjString('f'))
        l.Add(TObjString('g'))
        l.Add(TObjString('h'))
        l.Add(TObjString('i'))
        l.Add(TObjString('j'))

        self.assertEqual(len(l), 10)
        self.assertEqual(l[3], 'd')
        self.assertEqual(l[-1], 'j')
        self.assertRaises(IndexError, l.__getitem__, 20)
        self.assertRaises(IndexError, l.__getitem__, -20)

        self.assertEqual(list(l),
                         ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'])

        l[3] = TObjString('z')
        self.assertEqual(list(l),
                         ['a', 'b', 'c', 'z', 'e', 'f', 'g', 'h', 'i', 'j'])

        del l[2]
        self.assertEqual(list(l),
                         ['a', 'b', 'z', 'e', 'f', 'g', 'h', 'i', 'j'])

        self.assert_(TObjString('b') in l)
        self.assert_(not TObjString('x') in l)

        self.assertEqual(list(l[2:6]), ['z', 'e', 'f', 'g'])
        self.assertEqual(list(l[2:6:2]), ['z', 'f'])
        self.assertEqual(list(l[-5:-2]), ['f', 'g', 'h'])
        self.assertEqual(list(l[7:]), ['i', 'j'])
        self.assertEqual(list(l[:3]), ['a', 'b', 'z'])

        del l[2:4]
        self.assertEqual(list(l), ['a', 'b', 'f', 'g', 'h', 'i', 'j'])

        l[2:5] = [TObjString('1'), TObjString('2')]
        self.assertEqual(list(l), ['a', 'b', '1', '2', 'i', 'j'])

        l[6:6] = [TObjString('3')]
        self.assertEqual(list(l), ['a', 'b', '1', '2', 'i', 'j', '3'])

        l.append(TObjString('4'))
        self.assertEqual(list(l), ['a', 'b', '1', '2', 'i', 'j', '3', '4'])

        l.extend([TObjString('5'), TObjString('j')])
        self.assertEqual(list(l),
                         ['a', 'b', '1', '2', 'i', 'j', '3', '4', '5', 'j'])
        self.assertEqual(l.count('b'), 1)
        self.assertEqual(l.count('j'), 2)
        self.assertEqual(l.count('x'), 0)

        self.assertEqual(l.index(TObjString('i')), 4)
        self.assertRaises(ValueError, l.index, TObjString('x'))

        l.insert(3, TObjString('6'))
        l.insert(20, TObjString('7'))
        l.insert(-1, TObjString('8'))
        self.assertEqual(
            list(l),
            ['8', 'a', 'b', '1', '6', '2', 'i', 'j', '3', '4', '5', 'j', '7'])
        self.assertEqual(l.pop(), '7')
        self.assertEqual(l.pop(3), '1')
        self.assertEqual(
            list(l), ['8', 'a', 'b', '6', '2', 'i', 'j', '3', '4', '5', 'j'])

        l.remove(TObjString('j'))
        l.remove(TObjString('3'))

        self.assertRaises(ValueError, l.remove, TObjString('x'))
        self.assertEqual(list(l),
                         ['8', 'a', 'b', '6', '2', 'i', '4', '5', 'j'])

        l.reverse()
        self.assertEqual(list(l),
                         ['j', '5', '4', 'i', '2', '6', 'b', 'a', '8'])

        l.sort()
        self.assertEqual(list(l),
                         ['2', '4', '5', '6', '8', 'a', 'b', 'i', 'j'])

        if sys.hexversion >= 0x3000000:
            l.sort(key=TObjString.GetName)
            l.reverse()
        else:
            l.sort(lambda a, b: cmp(b.GetName(), a.GetName()))
        self.assertEqual(list(l),
                         ['j', 'i', 'b', 'a', '8', '6', '5', '4', '2'])

        l2 = l[:3]
        self.assertEqual(list(l2 * 3),
                         ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'])
        self.assertEqual(list(3 * l2),
                         ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'])

        l2 *= 3
        self.assertEqual(list(l2),
                         ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'])

        l2 = l[:3]
        l3 = l[6:8]
        self.assertEqual(list(l2 + l3), ['j', 'i', 'b', '5', '4'])

        if sys.hexversion >= 0x3000000:
            next = '__next__'
        else:
            next = 'next'
        i = iter(l2)
        self.assertEqual(getattr(i, next)(), 'j')
        self.assertEqual(getattr(i, next)(), 'i')
        self.assertEqual(getattr(i, next)(), 'b')
        self.assertRaises(StopIteration, getattr(i, next))
Ejemplo n.º 27
0
    fit.bkg_n = nbkg
    prefix = "Bg"
    if len(fit_params) > 1:
        prefix = fit.name + "_" + prefix
    (bkg_pdf, bkg_components, bkg_params) = init_pdfs(fit, fpars["background"], prefix, False)
    all_params.update(bkg_params)

    nextra = None
    if 'extra' in fpars:
        prefix = "Ext"
        if len(fit_params) > 1:
            prefix = fit.name + "_" + prefix
        (extra_pdf, extra_components, extra_params) = init_pdfs(fit, fpars["extra"], prefix, False)
        extra_yield = fpars["extra"]["yield"]
        if 'params' in extra_yield:
            _extra_params = TList()
            _extra_constrs = TList()
            for _par in extra_yield["params"]:
                _name = _par["name"]
                if _name in all_params:
                    _extra_params.append(all_params[_name])
                else:
                    _var = RooRealVar(
                            _name, _par["title"], _par["value"],
                            _par["min"], _par["max"], _par["units"]
                            )
                    fit.parameters.push_back(_var)
                    _extra_params.Add(_var)
                    all_params[_name] = _var
                    if "constraint" in _par:
                        _constr = RooGaussian(
Ejemplo n.º 28
0
ds = genpdf['pdf'].generate(RooArgSet(*genpdf['obs']), 150000, RooFit.Verbose())
#saveEta(ds);

# HACK (2/2): restore correct eta range after generation
ds.get().find('eta').setRange(0.0, 0.5)

ds.Print('v')
for o in genpdf['obs']:
    if not o.InheritsFrom('RooAbsCategory'): continue
    ds.table(o).Print('v')

from ROOT import RooDataSet, RooArgSet

from ROOT import TList

mistagresultList = TList();
etaAvgList = TList();
print "**** ADDING CATEGORIES ****"
ds = ds.reduce("qt!=qt::Untagged");
xRegions = createCategoryHistogram(ds,ds.get().find('eta'),NUMCAT);
ds.addColumn(xRegions)
ds.table(xRegions).Print("v")

keepvars = [ds.get().find(name) for name in ['qt', 'qf', 'tageffRegion', 'time']]

for i in xrange(NUMCAT):
    etaAvgList.AddLast(ds.reduce("tageffRegion == tageffRegion::Cat%u" % (i + 1,)).meanVar(ds.get().find('eta')));

etaAvg = ds.meanVar(ds.get().find('eta'));

dspercat = [ ds.reduce(RooArgSet(*keepvars),"tageffRegion == tageffRegion::Cat%u" % (i + 1,)) for i in xrange(NUMCAT) ]
Ejemplo n.º 29
0
reader.BookMVA(
    "testmva",
    "/home/arun/VHbbNtuples_7_6_x/CMSSW_7_6_3_for8X/src/tthAnalysis/HiggsToTauTau/test/training_2lss_1tau_TTV/weights_2lss_1tau_TTV_wCuts_10Var_15Mar2017/mvaTTHvsTTV2lss1tau_BDTG.weights.xml"
)

file = TFile(
    "/home/arun/ttHAnalysis/2016/2017Feb21_fastsim_dR03mvaLoose/histograms/2lss_1tau/forBDTtraining_SS_OS/ntuple_2lss_1tau_SS_OS_all.root"
)
if process == "signal":
    tree = file.Get("2lss_1tau_lepSS_sumOS_Loose/sel/evtntuple/signal/evtTree")
elif process == "TTV":
    tree_ttw = file.Get(
        "2lss_1tau_lepSS_sumOS_Loose/sel/evtntuple/TTW/evtTree")
    tree_ttz = file.Get(
        "2lss_1tau_lepSS_sumOS_Loose/sel/evtntuple/TTZ/evtTree")
    list = TList()
    list.Add(tree_ttw)
    list.Add(tree_ttz)
    tree = TTree.MergeTrees(list)
    tree.SetName("bkgTree")

c1 = TCanvas()
c1.SetFillColor(10)
c1.SetBorderSize(2)
c1.SetLeftMargin(0.12)
c1.SetBottomMargin(0.12)
c1.SetRightMargin(0.05)
c1.SetLogy()

histogram_base = TH1F("histogram_base", "", 100, -1., 1.)
histogram_base.SetTitle("")
Ejemplo n.º 30
0
    def test2Lists(self):
        """Test list/TList behavior and compatibility"""

        l = TList()
        l.Add(TObjString('a'))
        l.Add(TObjString('b'))
        l.Add(TObjString('c'))
        l.Add(TObjString('d'))
        l.Add(TObjString('e'))
        l.Add(TObjString('f'))
        l.Add(TObjString('g'))
        l.Add(TObjString('h'))
        l.Add(TObjString('i'))
        l.Add(TObjString('j'))

        self.assertEqual(len(l), 10)
        self.assertEqual(l[3], 'd')
        self.assertEqual(l[-1], 'j')
        self.assertRaises(IndexError, l.__getitem__, 20)
        self.assertRaises(IndexError, l.__getitem__, -20)

        self.assertEqual(list(l),
                         ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'])

        l[3] = TObjString('z')
        self.assertEqual(list(l),
                         ['a', 'b', 'c', 'z', 'e', 'f', 'g', 'h', 'i', 'j'])

        del l[2]
        self.assertEqual(list(l),
                         ['a', 'b', 'z', 'e', 'f', 'g', 'h', 'i', 'j'])

        self.assert_(TObjString('b') in l)
        self.assert_(not TObjString('x') in l)

        self.assertEqual(list(l[2:6]), ['z', 'e', 'f', 'g'])
        self.assertEqual(list(l[2:6:2]), ['z', 'f'])
        self.assertEqual(list(l[-5:-2]), ['f', 'g', 'h'])
        self.assertEqual(list(l[7:]), ['i', 'j'])
        self.assertEqual(list(l[:3]), ['a', 'b', 'z'])

        del l[2:4]
        self.assertEqual(list(l), ['a', 'b', 'f', 'g', 'h', 'i', 'j'])

        l[2:5] = [TObjString('1'), TObjString('2')]
        self.assertEqual(list(l), ['a', 'b', '1', '2', 'i', 'j'])

        l[6:6] = [TObjString('3')]
        self.assertEqual(list(l), ['a', 'b', '1', '2', 'i', 'j', '3'])

        l.append(TObjString('4'))
        self.assertEqual(list(l), ['a', 'b', '1', '2', 'i', 'j', '3', '4'])

        l.extend([TObjString('5'), TObjString('j')])
        self.assertEqual(list(l),
                         ['a', 'b', '1', '2', 'i', 'j', '3', '4', '5', 'j'])
        self.assertEqual(l.count('b'), 1)
        self.assertEqual(l.count('j'), 2)
        self.assertEqual(l.count('x'), 0)

        self.assertEqual(l.index(TObjString('i')), 4)
        self.assertRaises(ValueError, l.index, TObjString('x'))

        l.insert(3, TObjString('6'))
        l.insert(20, TObjString('7'))
        l.insert(-1, TObjString('8'))
        if not self.legacy_pyroot:
            # The pythonisation of TSeqCollection in experimental PyROOT mimics the
            # behaviour of the Python list, in this case for insert.
            # The Python list insert always inserts before the specified index, so if
            # -1 is specified, insert will place the new element right before the last
            # element of the list.
            self.assertEqual(list(l), [
                'a', 'b', '1', '6', '2', 'i', 'j', '3', '4', '5', 'j', '8', '7'
            ])
            # Re-synchronize with current PyROOT's list
            l.insert(0, TObjString('8'))
            self.assertEqual(list(l), [
                '8', 'a', 'b', '1', '6', '2', 'i', 'j', '3', '4', '5', 'j',
                '8', '7'
            ])
            l.pop(-2)
            self.assertEqual(list(l), [
                '8', 'a', 'b', '1', '6', '2', 'i', 'j', '3', '4', '5', 'j', '7'
            ])
        else:
            self.assertEqual(list(l), [
                '8', 'a', 'b', '1', '6', '2', 'i', 'j', '3', '4', '5', 'j', '7'
            ])
        self.assertEqual(l.pop(), '7')
        self.assertEqual(l.pop(3), '1')
        self.assertEqual(
            list(l), ['8', 'a', 'b', '6', '2', 'i', 'j', '3', '4', '5', 'j'])

        l.remove(TObjString('j'))
        l.remove(TObjString('3'))

        self.assertRaises(ValueError, l.remove, TObjString('x'))
        self.assertEqual(list(l),
                         ['8', 'a', 'b', '6', '2', 'i', '4', '5', 'j'])

        l.reverse()
        self.assertEqual(list(l),
                         ['j', '5', '4', 'i', '2', '6', 'b', 'a', '8'])

        l.sort()
        self.assertEqual(list(l),
                         ['2', '4', '5', '6', '8', 'a', 'b', 'i', 'j'])

        if sys.hexversion >= 0x3000000:
            l.sort(key=TObjString.GetName)
            l.reverse()
        else:
            l.sort(lambda a, b: cmp(b.GetName(), a.GetName()))
        self.assertEqual(list(l),
                         ['j', 'i', 'b', 'a', '8', '6', '5', '4', '2'])

        l2 = l[:3]
        self.assertEqual(list(l2 * 3),
                         ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'])
        self.assertEqual(list(3 * l2),
                         ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'])

        l2 *= 3
        self.assertEqual(list(l2),
                         ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'])

        l2 = l[:3]
        l3 = l[6:8]
        self.assertEqual(list(l2 + l3), ['j', 'i', 'b', '5', '4'])

        if sys.hexversion >= 0x3000000:
            next = '__next__'
        else:
            next = 'next'
        i = iter(l2)
        self.assertEqual(getattr(i, next)(), 'j')
        self.assertEqual(getattr(i, next)(), 'i')
        self.assertEqual(getattr(i, next)(), 'b')
        self.assertRaises(StopIteration, getattr(i, next))
Ejemplo n.º 31
0
    currentRangeTF1List.Delete()
    #for i in range(1,numCat+1):
    #currentRangeTF1List.Last().IsA().Destructor(currentRangeTF1List.Last());

    #s = raw_input("Press Enter to continue...");

    return rangeFit


import os

os.chdir(os.environ['B2DXFITTERSROOT'] + '/tutorial/etaHist')
fileList = os.listdir(os.getcwd())

#TList containing relevant data for each eta set in directory
mainResultList = TList()

#number of tagging categories
numTagCat = 5

#parse etaHist directory for .root files
for i in fileList[:1]:
    if (i[-5:] != '.root'):
        continue
    print i
    inFile = TFile(i)
    inFile.GetListOfKeys().Print()
    inFile.ls()
    #s = raw_input("Press Enter to continue");
    theTH1DHist = TH1F()
    inFile.GetObject("etaHist", theTH1DHist)
Ejemplo n.º 32
0
   def test2Lists( self ):
      """Test list/TList behavior and compatibility"""

      l = TList()
      l.Add( TObjString('a') )
      l.Add( TObjString('b') )
      l.Add( TObjString('c') )
      l.Add( TObjString('d') )
      l.Add( TObjString('e') )
      l.Add( TObjString('f') )
      l.Add( TObjString('g') )
      l.Add( TObjString('h') )
      l.Add( TObjString('i') )
      l.Add( TObjString('j') )

      self.assertEqual( len(l), 10 )
      self.assertEqual( l[3], 'd' )
      self.assertEqual( l[-1], 'j' )
      self.assertRaises( IndexError, l.__getitem__,  20 )
      self.assertRaises( IndexError, l.__getitem__, -20 )

      self.assertEqual( list(l), ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'] )

      l[3] = TObjString('z')
      self.assertEqual( list(l), ['a', 'b', 'c', 'z', 'e', 'f', 'g', 'h', 'i', 'j'] )

      del l[2]
      self.assertEqual( list(l), ['a', 'b', 'z', 'e', 'f', 'g', 'h', 'i', 'j'] )

      self.assert_( TObjString('b') in l )
      self.assert_( not TObjString('x') in l )

      self.assertEqual( list(l[2:6]),   ['z', 'e', 'f', 'g'] )
      self.assertEqual( list(l[2:6:2]), ['z', 'f'] )
      self.assertEqual( list(l[-5:-2]), ['f', 'g', 'h'] )
      self.assertEqual( list(l[7:]),    ['i', 'j'] )
      self.assertEqual( list(l[:3]),    ['a', 'b', 'z'] )

      del l[2:4]
      self.assertEqual( list(l), ['a', 'b', 'f', 'g', 'h', 'i', 'j'] )

      l[2:5] = [ TObjString('1'), TObjString('2') ]
      self.assertEqual( list(l), ['a', 'b', '1', '2', 'i', 'j'] )

      l[6:6] = [ TObjString('3') ]
      self.assertEqual( list(l), ['a', 'b', '1', '2', 'i', 'j', '3'] )

      l.append( TObjString('4') )
      self.assertEqual( list(l), ['a', 'b', '1', '2', 'i', 'j', '3', '4'] )

      l.extend( [ TObjString('5'), TObjString('j') ] )
      self.assertEqual( list(l), ['a', 'b', '1', '2', 'i', 'j', '3', '4', '5', 'j'] )
      self.assertEqual( l.count( 'b' ), 1 )
      self.assertEqual( l.count( 'j' ), 2 )
      self.assertEqual( l.count( 'x' ), 0 )

      self.assertEqual( l.index( TObjString( 'i' ) ), 4 )
      self.assertRaises( ValueError, l.index, TObjString( 'x' ) )

      l.insert(  3, TObjString('6') )
      l.insert( 20, TObjString('7') )
      l.insert( -1, TObjString('8') )
      self.assertEqual( list(l), ['8', 'a', 'b', '1', '6', '2', 'i', 'j', '3', '4', '5', 'j', '7'] )
      self.assertEqual( l.pop(), '7' )
      self.assertEqual( l.pop(3), '1' )
      self.assertEqual( list(l), ['8', 'a', 'b', '6', '2', 'i', 'j', '3', '4', '5', 'j'] )

      l.remove( TObjString( 'j' ) )
      l.remove( TObjString( '3' ) )

      self.assertRaises( ValueError, l.remove, TObjString( 'x' ) )
      self.assertEqual( list(l), ['8', 'a', 'b', '6', '2', 'i', '4', '5', 'j'] )

      l.reverse()
      self.assertEqual( list(l), ['j', '5', '4', 'i', '2', '6', 'b', 'a', '8'] )

      l.sort()
      self.assertEqual( list(l), ['2', '4', '5', '6', '8', 'a', 'b', 'i', 'j'] )

      if sys.hexversion >= 0x3000000:
         l.sort( key=TObjString.GetName )
         l.reverse()
      else:
         l.sort( lambda a, b: cmp(b.GetName(),a.GetName()) )
      self.assertEqual( list(l), ['j', 'i', 'b', 'a', '8', '6', '5', '4', '2'] )

      l2 = l[:3]
      self.assertEqual( list(l2 * 3), ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'] )
      self.assertEqual( list(3 * l2), ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'] )

      l2 *= 3
      self.assertEqual( list(l2), ['j', 'i', 'b', 'j', 'i', 'b', 'j', 'i', 'b'] )

      l2 = l[:3]
      l3 = l[6:8]
      self.assertEqual( list(l2+l3), ['j', 'i', 'b', '5', '4'] )

      i = iter(l2)
      self.assertEqual( i.next(), 'j' )
      self.assertEqual( i.next(), 'i' )
      self.assertEqual( i.next(), 'b' )
      self.assertRaises( StopIteration, i.next )
Ejemplo n.º 33
0
def divideEtaTH1Ds(etaHist,
                   numCat,
                   sourceFileName=None,
                   categoryList=None,
                   drawHist=False):
    # etaHist - the TH1D to split into categories
    # numCat - number of tagging categories
    # categoryList - list of tagging category names
    # drawHist - whether or not to display each category-split histogram
    # sourceFileName - name of etaHist root file

    #normalize etaHist
    etaHist.Scale(1.0 / etaHist.Integral())

    #create default category names if none are given
    if categoryList == None:
        categoryList = []
        for i in range(numCat):
            categoryList += ["Cat" + str(i)]

    histSum = etaHist.Integral()

    #get list of limits
    limList = [-1]

    for i in range(1, numCat):
        start = limList[i - 1] + 1
        targetVal = histSum / numCat
        left = start
        right = etaHist.GetNbinsX()
        #print "\n\n\n",etaHist.Integral(start,start+1),"\n\n\n\n";
        while abs(left - right) > 1:
            if etaHist.Integral(start, int((left + right) / 2)) > targetVal:
                right = (left + right) / 2
            else:
                left = (left + right) / 2
            #print start, left, right, etaHist.Integral(start,int((left+right)/2)),targetVal,",",

        if abs(etaHist.Integral(start, left) -
               targetVal) > abs(etaHist.Integral(start, right) - targetVal):
            limList += [right]
        else:
            limList += [left]

    limList[0] = 0
    limList += [etaHist.GetNbinsX()]

    print limList

    #rangeFit is a 4-element TList containing:
    # - the sourceFileName as element 0
    # - a RooThresholdCategory as element 2, and its RooRealVar as element 1
    # - a tList of tFitResult as element 3
    rangeFit = TList()
    rangeFit.AddLast(TObjString(sourceFileName))
    rangeFit.AddLast(RooRealVar('x', 'x', 0.0, 1.0))
    rangeFit.AddLast(
        RooThresholdCategory("tageffRegion", "region of tageff",
                             rangeFit.At(1), "Cat" + str(numCat)))
    rangeFit.AddLast(TList())

    #plot each category-split histogram, if necessary
    if drawHist == True:

        ROOT.gStyle.SetPalette(ROOT.kOcean)

        #create stack to contain the category TH1Ds
        etaHistStack = THStack("etaHistStack", "Stack of TH1Ds")

        #create category-masking function for TH1D clones
        histCutterFunc = TF1("histCutterFunc",
                             "((x>=[0])?((x<[1])?1.0:0.0):0.0)", 0.0, 1.0)

        for i in range(len(limList) - 1):

            etaHistClone = etaHist.Clone()
            histCutterFunc.SetParameter(
                0,
                etaHist.GetXaxis().GetBinCenter(limList[i]))
            histCutterFunc.SetParameter(
                1,
                etaHist.GetXaxis().GetBinCenter(limList[i + 1]))
            #histCutterFunc.Draw();
            #raw_input("Press Enter to continue to next hisCutterFunc");
            etaHistClone.Multiply(histCutterFunc)
            etaHistClone.SetFillColor(38 + i)
            etaHistStack.Add(etaHistClone)

        etaHistClone = etaHist.Clone()
        etaHistClone.SetFillColor(38 + len(limList))
        #etaHistStack.Add(etaHistClone);

        import time

        os.chdir(os.environ['B2DXFITTERSROOT'] + '/tutorial')

        if (not (os.path.isdir('fits'))):
            os.mkdir('fits')

        os.chdir('fits')
        histCanvas = TCanvas()
        etaHistStack.Draw("hist PFC")
        #etaHistClone.Draw("hist PFC");
        #histCanvas.SaveAs('tagRegionFitList_%f.pdf' % time.time());

    #create thresholds and fitting functions (all linear, but with different ranges corresponding to the categories)

    currentRangeTF1List = TList()

    for i in range(1, numCat + 1):
        if (i <= numCat):
            rangeFit.At(2).addThreshold(
                etaHist.GetXaxis().GetBinCenter(limList[i]),
                categoryList[i - 1])
        currentRangeTF1List.AddLast(
            TF1(
                "fitFuncEtaset", "[0]+[1]*(x-" + str(
                    etaHist.Integral(limList[i - 1], limList[i]) /
                    (limList[i] - limList[i - 1])) + ")",
                etaHist.GetXaxis().GetBinCenter(limList[i - 1]),
                etaHist.GetXaxis().GetBinCenter(limList[i])))
        #currentRangeTF1 = TF1("fitFuncEtaset","[0]+[1]*x",etaHist.GetXaxis().GetBinCenter(limList[i-1]),etaHist.GetXaxis().GetBinCenter(limList[i]));
        rangeFit.Last().AddLast(
            etaHist.Fit(currentRangeTF1List.Last(), "R0S").Get().Clone())
        if (drawHist == True):
            currentRangeTF1List.Last().DrawCopy('same')
        #raw_input('Press Enter to continue to next fit function');
        #currentRangeTF1.IsA().Destructor(currentRangeTF1);
        #print "P0, P1 = ",rangeFit.Last().Last().Parameter(0), rangeFit.Last().Last().Parameter(1);

    histCanvas.SaveAs('tagRegionFitList_%f.pdf' % time.time())
    currentRangeTF1List.Delete()
    #for i in range(1,numCat+1):
    #currentRangeTF1List.Last().IsA().Destructor(currentRangeTF1List.Last());

    #s = raw_input("Press Enter to continue...");

    return rangeFit