Beispiel #1
0
def specialWrite():
    """ ./job-panda.py -swrite
    Write TCuts from waveSkim files into splitSkim files.
    """
    from ROOT import TFile, TNamed, TObject
    cal = dsi.CalInfo()
    runList = cal.GetSpecialRuns("longCal", 5)

    for run in runList:
        dsNum = dsi.GetDSNum(run)
        wavePath = "%s/waves/waveSkimDS%d_run%d.root" % (dsi.specialDir, dsNum,
                                                         run)
        waveFile = TFile(wavePath)
        theCut = waveFile.Get("theCut").GetTitle()
        print(wavePath)

        splitFiles = glob.glob("%s/split/splitSkimDS%d_run%d*.root" %
                               (dsi.specialDir, dsNum, run))
        for idx in range(len(splitFiles)):
            if idx == 0:
                splitPath = "%s/split/splitSkimDS%d_run%d.root" % (
                    dsi.specialDir, dsNum, run)
            else:
                splitPath = "%s/split/splitSkimDS%d_run%d_%d.root" % (
                    dsi.specialDir, dsNum, run, idx)
            if not os.path.isfile(splitPath):
                print("File doesn't exist:", splitPath)
                return

            splitFile = TFile(splitPath, "UPDATE")
            thisCut = TNamed("theCut", theCut)
            thisCut.Write("", TObject.kOverwrite)
        print(splitFiles[-1])
Beispiel #2
0
def splitTree():

    dsNum, subNum = 0, 25

    inPath = "/global/homes/w/wisecg/project/waveskim/waveSkimDS%d_%d.root" % (dsNum, subNum)
    outPath = "./cutSkimDS%d_%d.root" % (dsNum, subNum)

    inFile = TFile(inPath)
    bigTree = inFile.Get("skimTree")

    theCut = inFile.Get("theCut").GetTitle()

    bigTree.Draw(">>elist",theCut,"entrylist")
    elist = gDirectory.Get("elist")
    bigTree.SetEntryList(elist)
    nList = elist.GetN()

    outFile = TFile(outPath,"RECREATE")

    lilTree = TTree()
    lilTree.SetMaxTreeSize(150000000)
    lilTree = bigTree.CopyTree("")

    lilTree.Write("",TObject.kOverwrite)

    thisCut = TNamed("theCut",theCut)
    thisCut.Write()
Beispiel #3
0
def splitTree(dsNum, subNum=None, runNum=None):
    """ ./job-panda.py -split (-sub dsNum subNum) (-run dsNum runNum)

        Split a SINGLE waveSkim file into small (~50MB) files to speed up LAT parallel processing.
        Can call 'batchSplit' instead to submit each run in the list as a job, splitting the files in parallel.
        NOTE: The cut written into the first file is NOT copied into the additional files
              (I couldn't get it to work within this function -- kept getting "file not closed" errors.)
              To clean up, do that with the 'writeCut' function below, potentially AFTER a big parallel job.
    """
    from ROOT import TFile, TTree, gDirectory, TEntryList, TNamed, TObject, gROOT

    print("Splitting tree.  dsNum:", dsNum, "subNum:", subNum, "runNum:",
          runNum)

    # Set input and output paths.  Clear out any files from a previous
    # try before you attempt a copy (avoid the double underscore)
    inPath, outPath = "", ""
    if runNum == None:
        # bg mode
        inPath = "%s/waveSkimDS%d_%d.root" % (dsi.waveDir, dsNum, subNum)
        outPath = "%s/splitSkimDS%d_%d.root" % (dsi.splitDir, dsNum, subNum)
        fileList = sorted(
            glob.glob("%s/splitSkimDS%d_%d*.root" %
                      (dsi.splitDir, dsNum, subNum)))
        for f in fileList:
            os.remove(f)
    elif subNum == None:
        # cal mode
        inPath = "%s/waveSkimDS%d_run%d.root" % (dsi.calWaveDir, dsNum, runNum)
        outPath = "%s/splitSkimDS%d_run%d.root" % (dsi.calSplitDir, dsNum,
                                                   runNum)
        fileList = sorted(
            glob.glob("%s/splitSkimDS%d_run%d*.root" %
                      (dsi.calSplitDir, dsNum, runNum)))
        for f in fileList:
            os.remove(f)

    inFile = TFile(inPath)
    bigTree = inFile.Get("skimTree")
    theCut = inFile.Get("theCut").GetTitle()
    bigTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    bigTree.SetEntryList(elist)
    nList = elist.GetN()

    outFile = TFile(outPath, "RECREATE")
    lilTree = TTree()
    lilTree.SetMaxTreeSize(50000000)  # 50 MB
    thisCut = TNamed("theCut", theCut)
    thisCut.Write("", TObject.kOverwrite)
    lilTree = bigTree.CopyTree(
        "")  # this does NOT write the cut into the extra files
    lilTree.Write("", TObject.kOverwrite)
Beispiel #4
0
def complicatedCopy():

    gFile = TFile("~/project/mjddatadir/gatified/mjd_run13071.root")
    bFile = TFile("~/project/mjddatadir/built/OR_run13071.root")
    gatTree = gFile.Get("mjdTree")
    bltTree = bFile.Get("MGTree")
    gatTree.AddFriend(bltTree)

    theCut = "mH > 1 && !EventDC1Bits && Entry$ < 320"

    gatTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    gatTree.SetEntryList(elist)
    nList = elist.GetN()
    print "Using cut:\n", theCut
    print "Found", gatTree.GetEntries(), "input events."
    print "Found", nList, "events passing cuts."

    print "Starting event loop ..."
    for iList in range(nList):
        entry = gatTree.GetEntryNumber(iList)
        gatTree.LoadTree(entry)
        gatTree.GetEntry(entry)
        nChans = gatTree.channel.size()

        # Loop over hits passing cuts
        numPass = gatTree.Draw("channel", theCut, "GOFF", 1, iList)
        chans = gatTree.GetV1()
        chanList = list(set(int(chans[n]) for n in xrange(numPass)))
        hitList = (iH for iH in xrange(nChans)
                   if gatTree.channel.at(iH) in chanList
                   )  # a 'generator expression'
        for iH in hitList:

            # Load data
            run = gatTree.run
            chan = gatTree.channel.at(iH)
            dataENF = gatTree.trapENFCal.at(iH)
            print "run %d  iList %d  chan %d  enf %.2f" % (run, iList, chan,
                                                           dataENF)

    print "Copying tree with same cut ..."
    outFile = TFile("~/project/lat/copyTree_test.root", "RECREATE")
    outTree = gatTree.CopyTree("")
    outTree.Write()
    print "Wrote", outTree.GetEntries(), "events."
    cutUsed = TNamed("theCut", theCut)
    cutUsed.Write()
Beispiel #5
0
def writeCut(dsNum, subNum=None, runNum=None, calList=[]):
    """ ./job-panda.py -writeCut (-ds dsNum) (-sub dsNum subNum) (-run dsNum subNum) [-cal]
        Assumes the cut used in the FIRST file (even in the whole DS) should be applied
        to ALL files.  This should be a relatively safe assumption.
    """
    from ROOT import TFile, TNamed, TObject
    fileList = []
    bkg = dsi.BkgInfo()

    # bg
    if not calList:
        dsMap = bkg.dsMap()
        # -ds
        if subNum == None and runNum == None:
            for i in range(dsMap[dsNum] + 1):
                inPath = "%s/splitSkimDS%d_%d*" % (dsi.splitDir, dsNum, i)
                fileList.extend(sorted(glob.glob(inPath)))
        # -sub
        elif runNum == None:
            inPath = "%s/splitSkimDS%d_%d*" % (dsi.splitDir, dsNum, subNum)
            fileList.extend(sorted(glob.glob(inPath)))
        # -run
        elif subNum == None:
            inPath = "%s/splitSkimDS%d_run%d*" % (dsi.splitDir, dsNum, runNum)
            fileList.extend(sorted(glob.glob(inPath)))
    # cal
    else:
        dsRanges = bkg.dsRanges()
        for run in calList:
            for key in dsRanges:
                if dsRanges[key][0] <= run <= dsRanges[key][1]:
                    dsNum = key
            inPath = "%s/splitSkimDS%d_run%d*" % (dsi.calSplitDir, dsNum, run)
            fileList.extend(sorted(glob.glob(inPath)))

    # Pull the cut off the FIRST file and add it to the sub-files
    if len(fileList) <= 1:
        print("No files found!  Exiting...")
        exit(1)

    firstFile = TFile(fileList[0])
    theCut = firstFile.Get("theCut").GetTitle()
    print("Applying this cut:\n", theCut)
    for f in fileList:
        print(f)
        subRangeFile = TFile(f, "UPDATE")
        thisCut = TNamed("theCut", theCut)
        thisCut.Write("", TObject.kOverwrite)
Beispiel #6
0
def save_command_line_in_tfile(output_tfile_, command_line_):
    from ROOT import TNamed
    from ROOT import gDirectory

    current_directory = gDirectory
    mkdir_cd_rootfile(output_tfile_, "")
    TNamed("command_line_TNamed", command_line_).Write()
    current_directory.cd()
Beispiel #7
0
 def __init__(self):
     TNamed.__init__(self, 'Kernel', 'Kernel')
     self.FlowList = []
     self.wb = None
     self.ao = None
     self.dataset = None
     self.treename = None
     self.samplename = None
     self.inputpath = None
     self.outputpath = None
     self.outputdump = None
     self.corrections = []
     self.systematics = []
     self.totalflow = None
     self.inputtree = None
     self.outputtree = None
     self._tempVec = []
     self._tempTree = {}
Beispiel #8
0
 def __init__(self):
     TNamed.__init__(self, 'Kernel', 'Kernel')
     self.FlowList = []
     self.wb = None
     self.ao = None
     self.dataset = None
     self.treename = None
     self.samplename = None
     self.inputpath = None
     self.outputpath = None
     self.outputdump = None
     self.corrections = []
     self.systematics = []
     self.totalflow = None
     self.inputtree = None
     self.outputtree = None
     self._tempVec = []
     self._tempTree = {}
    def test11WriteTObjArray(self):
        """Test writing of a TObjArray"""

        f = TFile(self.fname, 'RECREATE')
        t = TTree(self.tname, self.ttitle)
        o = TObjArray()
        t.Branch('mydata', o)

        nameds = [TNamed(name, name) for name in self.testnames]
        for name in nameds:
            o.Add(name)
        self.assertEqual(len(o), len(self.testnames))

        t.Fill()

        f.Write()
        f.Close()
Beispiel #10
0
def AnalyzeDataSet():
    CSVMWP = 0.8484
    DCSVMWP = 0.6324
    NEntries = skimmedTree.GetEntries()
    #    NEntries = 1000
    h_total = TH1F('h_total', 'h_total', 2, 0, 2)
    h_total_mcweight = TH1F('h_total_mcweight', 'h_total_mcweight', 2, 0, 2)

    triglist = [
        'HLT_IsoMu20', 'HLT_Ele27_WPLoose_Gsf',
        'HLT_PFMETNoMu90_PFMHTNoMu90_IDTight_v',
        'HLT_PFMETNoMu120_PFMHTNoMu120_IDTight_v'
    ]

    outfile = TFile(outfilename, 'RECREATE')

    outTree = TTree('outTree', 'tree branches')
    samplepath = TNamed('samplepath', str(sys.argv[1]))

    st_runId = numpy_.zeros(1, dtype=int)
    st_lumiSection = array('L', [0])
    st_eventId = array('L', [0])
    st_pfMetCorrPt = array('f', [0.])
    st_pfMetCorrPhi = array('f', [0.])
    st_isData = array('b', [0])
    for trigs in triglist:
        exec("st_" + trigs + "  = array( 'b', [ 0 ] )")
#    st_HLT_IsoMu20      = array( 'b', [ 0 ] )
#    st_HLT_Ele27_WPLoose_Gsf = array( 'b', [ 0 ] )

    maxn = 10

    st_THINnJet = array('L', [0])  #ROOT.std.vector('int')()
    st_THINjetP4 = ROOT.std.vector('TLorentzVector')()
    st_THINjetCISVV2 = ROOT.std.vector('float')()
    st_THINjetHadronFlavor = ROOT.std.vector('int')()
    st_THINjetNHadEF = ROOT.std.vector('float')()
    st_THINjetCHadEF = ROOT.std.vector('float')()

    st_AK4deepCSVnJet = array('L', [0])  #ROOT.std.vector('int')()
    st_AK4deepCSVjetP4 = ROOT.std.vector('TLorentzVector')()
    st_AK4deepCSVjetDeepCSV_b = ROOT.std.vector('float')()

    st_nEle = array('L', [0])  #ROOT.std.vector('int')()
    st_eleP4 = ROOT.std.vector('TLorentzVector')()
    st_eleIsPassLoose = ROOT.std.vector('bool')()
    st_eleIsPassMedium = ROOT.std.vector('bool')()
    st_eleIsPassTight = ROOT.std.vector('bool')()

    st_nMu = array('L', [0])  #ROOT.std.vector('int')()
    st_muP4 = ROOT.std.vector('TLorentzVector')()
    st_isLooseMuon = ROOT.std.vector('bool')()
    st_isMediumMuon = ROOT.std.vector('bool')()
    st_isTightMuon = ROOT.std.vector('bool')()
    st_muChHadIso = ROOT.std.vector('float')()
    st_muNeHadIso = ROOT.std.vector('float')()
    st_muGamIso = ROOT.std.vector('float')()
    st_muPUPt = ROOT.std.vector('float')()
    st_muCharge = ROOT.std.vector('int')()

    st_trigResult = ROOT.std.vector('bool')()
    st_trigName = ROOT.std.vector('string')()

    st_HPSTau_n = array('L', [0])  #ROOT.std.vector('int')()
    st_HPSTau_4Momentum = ROOT.std.vector('TLorentzVector')()

    mcweight = array('f', [0])
    st_pu_nTrueInt = array('f', [0])  #ROOT.std.vector('std::vector<float>')()

    st_nGenPar = array('L', [0])
    st_genParId = ROOT.std.vector('int')()
    st_genMomParId = ROOT.std.vector('int')()
    st_genParSt = ROOT.std.vector('int')()
    st_genParP4 = ROOT.std.vector('TLorentzVector')()

    WenuRecoil = array('f', [0.])
    Wenumass = array('f', [0.])
    WenuPhi = array('f', [0.])

    WmunuRecoil = array('f', [0.])
    Wmunumass = array('f', [0.])
    WmunuPhi = array('f', [0.])

    ZeeRecoil = array('f', [0.])
    ZeeMass = array('f', [0.])
    ZeePhi = array('f', [0.])

    ZmumuRecoil = array('f', [0.])
    ZmumuMass = array('f', [0.])
    ZmumuPhi = array('f', [0.])

    TOPRecoil = array('f', [0.])
    TOPPhi = array('f', [0.])

    outTree.Branch('st_runId', st_runId, 'st_runId/L')
    outTree.Branch('st_lumiSection', st_lumiSection, 'st_lumiSection/L')
    outTree.Branch('st_eventId', st_eventId, 'st_eventId/L')
    outTree.Branch('st_pfMetCorrPt', st_pfMetCorrPt, 'st_pfMetCorrPt/F')
    outTree.Branch('st_pfMetCorrPhi', st_pfMetCorrPhi, 'st_pfMetCorrPhi/F')
    outTree.Branch('st_isData', st_isData, 'st_isData/O')

    for trigs in triglist:
        exec("outTree.Branch( 'st_" + trigs + "', st_" + trigs + " , 'st_" +
             trigs + "/O')")
#    outTree.Branch( 'st_HLT_IsoMu20', st_HLT_IsoMu20 , 'st_HLT_IsoMu20/O')
#    outTree.Branch( 'st_HLT_Ele27_WPLoose_Gsf', st_HLT_Ele27_WPLoose_Gsf , 'st_HLT_Ele27_WPLoose_Gsf/O')

    outTree.Branch('st_THINnJet', st_THINnJet, 'st_THINnJet/L')
    #outTree.Branch( 'st_THINnJet',st_AK4deepCSVnJet, 'st_AK4deepCSVnJet/L' )
    outTree.Branch('st_THINjetP4', st_THINjetP4)
    #outTree.Branch( 'st_THINjetP4',st_AK4deepCSVjetP4 )
    outTree.Branch('st_THINjetCISVV2', st_THINjetCISVV2)
    #outTree.Branch( 'st_THINjetCISVV2',st_AK4deepCSVjetDeepCSV_b )
    outTree.Branch('st_THINjetHadronFlavor', st_THINjetHadronFlavor)
    outTree.Branch('st_THINjetNHadEF', st_THINjetNHadEF)
    outTree.Branch('st_THINjetCHadEF', st_THINjetCHadEF)

    outTree.Branch('st_AK4deepCSVnJet', st_AK4deepCSVnJet,
                   'st_AK4deepCSVnJet/L')
    outTree.Branch('st_AK4deepCSVjetP4', st_AK4deepCSVjetP4)
    outTree.Branch('st_AK4deepCSVjetDeepCSV_b', st_AK4deepCSVjetDeepCSV_b)

    outTree.Branch('st_nEle', st_nEle, 'st_nEle/L')
    outTree.Branch('st_eleP4', st_eleP4)
    outTree.Branch('st_eleIsPassLoose',
                   st_eleIsPassLoose)  #, 'st_eleIsPassLoose/O' )
    outTree.Branch('st_eleIsPassMedium',
                   st_eleIsPassMedium)  #, 'st_eleIsPassMedium/O' )
    outTree.Branch('st_eleIsPassTight',
                   st_eleIsPassTight)  #, 'st_eleIsPassTight/O' )

    outTree.Branch('st_nMu', st_nMu, 'st_nMu/L')
    outTree.Branch('st_muP4', st_muP4)
    outTree.Branch('st_isLooseMuon', st_isLooseMuon)  #, 'st_isLooseMuon/O' )
    outTree.Branch('st_isMediumMuon',
                   st_isMediumMuon)  #, 'st_isMediumMuon/O' )
    outTree.Branch('st_isTightMuon', st_isTightMuon)  #, 'st_isTightMuon/O' )
    outTree.Branch('st_muChHadIso', st_muChHadIso)  #, 'st_muChHadIso/F')
    outTree.Branch('st_muNeHadIso', st_muNeHadIso)  #, 'st_muNeHadIso/F')
    outTree.Branch('st_muGamIso', st_muGamIso)  #, 'st_muGamIso/F')
    outTree.Branch('st_muPUPt', st_muPUPt)  #, 'st_muPUPt/F')

    outTree.Branch('st_trigName', st_trigName)
    outTree.Branch('st_trigResult', st_trigResult)

    outTree.Branch('st_HPSTau_n', st_HPSTau_n, 'st_HPSTau_n/L')
    outTree.Branch('st_HPSTau_4Momentum', st_HPSTau_4Momentum)

    outTree.Branch('st_pu_nTrueInt', st_pu_nTrueInt, 'st_pu_nTrueInt/F')
    outTree.Branch('mcweight', mcweight, 'mcweight/F')
    outTree.Branch('st_nGenPar', st_nGenPar, 'st_nGenPar/L')  #nGenPar/I
    outTree.Branch('st_genParId', st_genParId)  #vector<int>
    outTree.Branch('st_genMomParId', st_genMomParId)
    outTree.Branch('st_genParSt', st_genParSt)
    outTree.Branch('st_genParP4', st_genParP4)

    outTree.Branch('WenuRecoil', WenuRecoil, 'WenuRecoil/F')
    outTree.Branch('Wenumass', Wenumass, 'Wenumass/F')
    outTree.Branch('WenuPhi', WenuPhi, 'WenuPhi/F')

    outTree.Branch('WmunuRecoil', WmunuRecoil, 'WmunuRecoil/F')
    outTree.Branch('Wmunumass', Wmunumass, 'Wmunumass/F')
    outTree.Branch('WmunuPhi', WmunuPhi, 'WmunuPhi/F')

    outTree.Branch('ZeeRecoil', ZeeRecoil, 'ZeeRecoil/F')
    outTree.Branch('ZeeMass', ZeeMass, 'ZeeMass/F')
    outTree.Branch('ZeePhi', ZeePhi, 'ZeePhi/F')

    outTree.Branch('ZmumuRecoil', ZmumuRecoil, 'ZmumuRecoil/F')
    outTree.Branch('ZmumuMass', ZmumuMass, 'ZmumuMass/F')
    outTree.Branch('ZmumuPhi', ZmumuPhi, 'ZmumuPhi/F')

    outTree.Branch('TOPRecoil', TOPRecoil, 'TOPRecoil/F')
    outTree.Branch('TOPPhi', TOPPhi, 'TOPPhi/F')

    for ievent in range(NEntries):

        #    print "\n*****\nWARNING: *Test run* Processing 200 events only.\n*****\n"
        #    for ievent in range(200):
        if ievent % 100 == 0:
            print "Processed " + str(ievent) + " of " + str(
                NEntries) + " events."
        skimmedTree.GetEntry(ievent)
        ## Get all relevant branches
        run = skimmedTree.__getattr__('runId')
        lumi = skimmedTree.__getattr__('lumiSection')
        event = skimmedTree.__getattr__('eventId')
        #        print "Run:"+str(run)+"; Lumi:"+str(lumi)+"; Event:"+str(event)
        trigName = skimmedTree.__getattr__('hlt_trigName')
        trigResult = skimmedTree.__getattr__('hlt_trigResult')
        filterName = skimmedTree.__getattr__('hlt_filterName')
        filterResult = skimmedTree.__getattr__('hlt_filterResult')

        pfMet = skimmedTree.__getattr__('pfMetCorrPt')
        pfMetPhi = skimmedTree.__getattr__('pfMetCorrPhi')

        nTHINJets = skimmedTree.__getattr__('THINnJet')
        thinjetP4 = skimmedTree.__getattr__('THINjetP4')
        thinJetCSV = skimmedTree.__getattr__('THINjetCISVV2')
        passThinJetLooseID = skimmedTree.__getattr__('THINjetPassIDLoose')
        THINjetHadronFlavor = skimmedTree.__getattr__('THINjetHadronFlavor')
        thinjetNhadEF = skimmedTree.__getattr__('THINjetNHadEF')
        thinjetChadEF = skimmedTree.__getattr__('THINjetCHadEF')

        try:
            nTHINdeepCSVJets = skimmedTree.__getattr__('AK4deepCSVnJet')
            thindeepCSVjetP4 = skimmedTree.__getattr__('AK4deepCSVjetP4')
            thinJetdeepCSV = skimmedTree.__getattr__('AK4deepCSVjetDeepCSV_b')
        except:
            if ievent == 0:
                print "\n**********WARNING: Looks like the ntuple is from an older version, as DeepCSV jet collection is missing. DeepCSV information will NOT be stored.**********\n"

        nEle = skimmedTree.__getattr__('nEle')
        eleP4 = skimmedTree.__getattr__('eleP4')
        eleIsPassLoose = skimmedTree.__getattr__('eleIsPassLoose')
        eleIsPassMedium = skimmedTree.__getattr__('eleIsPassMedium')
        eleIsPassTight = skimmedTree.__getattr__('eleIsPassTight')
        eleCharge = skimmedTree.__getattr__('eleCharge')

        nMu = skimmedTree.__getattr__('nMu')
        muP4 = skimmedTree.__getattr__('muP4')
        isLooseMuon = skimmedTree.__getattr__('isLooseMuon')
        isMediumMuon = skimmedTree.__getattr__('isMediumMuon')
        isTightMuon = skimmedTree.__getattr__('isTightMuon')
        muChHadIso = skimmedTree.__getattr__('muChHadIso')
        muNeHadIso = skimmedTree.__getattr__('muNeHadIso')
        muGamIso = skimmedTree.__getattr__('muGamIso')
        muPUPt = skimmedTree.__getattr__('muPUPt')
        muCharge = skimmedTree.__getattr__('muCharge')

        nTau = skimmedTree.__getattr__('HPSTau_n')
        tauP4 = skimmedTree.__getattr__('HPSTau_4Momentum')
        isDecayModeFinding = skimmedTree.__getattr__('disc_decayModeFinding')
        passLooseTauIso = skimmedTree.__getattr__(
            'disc_byLooseIsolationMVA3oldDMwLT')
        isData = skimmedTree.__getattr__('isData')
        mcWeight = skimmedTree.__getattr__('mcWeight')
        pu_nTrueInt = skimmedTree.__getattr__('pu_nTrueInt')  #int()

        #        print skimmedTree.__getattr__('pu_nTrueInt')
        #        print pu_nTrueInt
        #        print

        nGenPar = skimmedTree.__getattr__('nGenPar')
        genParId = skimmedTree.__getattr__('genParId')
        genMomParId = skimmedTree.__getattr__('genMomParId')
        genParSt = skimmedTree.__getattr__('genParSt')
        genParP4 = skimmedTree.__getattr__('genParP4')

        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # MC Weights ----------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        mcweight[0] = 0.0
        if isData == 1: mcweight[0] = 1.0
        if not isData:
            if mcWeight < 0: mcweight[0] = -1.0
            if mcWeight > 0: mcweight[0] = 1.0

        h_total.Fill(1.)
        h_total_mcweight.Fill(1., mcweight[0])

        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        ## Trigger selection
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        itrig_ = 0
        trig1 = False
        trig2 = False
        trig3 = False
        trig4 = False
        trig5 = False
        trig6 = False
        trig7 = False
        trig8 = False
        trig9 = False
        trig10 = False
        trig11 = False
        trig12 = False
        trig1 = CheckFilter(trigName, trigResult,
                            'HLT_PFMET170_')  # added from  monojet
        trig2 = CheckFilter(trigName, trigResult, 'HLT_PFMET170_NoiseCleaned')
        trig3 = CheckFilter(trigName, trigResult,
                            'HLT_PFMET170_JetIdCleaned_v')
        trig4 = CheckFilter(trigName, trigResult, 'HLT_PFMET170_HBHECleaned_v')
        trig5 = CheckFilter(trigName, trigResult,
                            'HLT_PFMETNoMu90_PFMHTNoMu90_IDTight_v')
        trig6 = CheckFilter(trigName, trigResult,
                            'HLT_PFMETNoMu100_PFMHTNoMu100_IDTight_v'
                            )  #added from  tt+DM all hadronic analysis
        trig7 = CheckFilter(trigName, trigResult,
                            'HLT_PFMETNoMu110_PFMHTNoMu110_IDTight_v')
        trig8 = CheckFilter(trigName, trigResult,
                            'HLT_PFMETNoMu120_PFMHTNoMu120_IDTight_v')
        trig9 = CheckFilter(trigName, trigResult, 'HLT_PFMET110_PFMHT110_')
        trig10 = CheckFilter(
            trigName, trigResult,
            'HLT_IsoMu24_v')  #added from tt+DM all hadronic analysis
        trig11 = CheckFilter(
            trigName, trigResult,
            'HLT_IsoTkMu24_v')  #added from tt+DM all hadronic analysis
        trig12 = CheckFilter(
            trigName, trigResult,
            'HLT_Ele27_WPTight_Gsf')  #added from Siew Yan slides
        #        trig13 = CheckFilter(trigName, trigResult, 'HLT_IsoMu20')   #Added from AN CR
        trig13 = CheckFilter(
            trigName, trigResult,
            'HLT_IsoMu24')  #Instead of IsoMu20 which is absent era E onwards
        trig14 = CheckFilter(trigName, trigResult,
                             'HLT_Ele27_WPLoose_Gsf')  #Added from AN CR

        #        print list(trigName)
        #        print [bool(i) for i in list(trigResult)]

        #        print ievent
        #        for itrig in range(len(list(trigResult))):
        #            if bool(list(trigResult)[itrig]): print list(trigName)[itrig]

        #        if 'HLT_IsoMu20' in list(trigName):
        #            print 'HLT_IsoMu20'
        #        if 'HLT_Ele27_WPLoose_Gsf' in list(trigName):
        #            print 'HLT_Ele27_WPLoose_Gsf'

        #        print list(trigName)
        #        for itr in list(trigName):
        #            if itr.find('IsoMu')!=-1: print itr
        ##            if itr.find('HLT_Ele27_WPLoose_Gsf')!=-1: print itr
        ###        print (trig13,trig14)
        #        print

        #        if not isData:
        #            trigstatus  = False # triggers are not required for MC
        #        if isData:
        #            trigstatus =  trig1 | trig2 | trig3 | trig4 | trig5 | trig6 | trig7 | trig8 | trig9 | trig10 | trig11 | trig12  #to include data with above triggers
        #        if not isData:
        #           if trigstatus == True : continue
        trigstatus = trig1 | trig2 | trig3 | trig4 | trig5 | trig6 | trig7 | trig8 | trig9 | trig10 | trig11 | trig12 | trig13 | trig14

        if not trigstatus: continue  #Currently doing this for both MC and data

        for itrig in range(len(list(trigName))):
            st_trigName.push_back(list(trigName)[itrig])
            st_trigResult.push_back(bool(list(trigResult)[itrig]))

#        print (isData,trigstatus)
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
## Filter selection
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        filterstatus = False
        filter1 = False
        filter2 = False
        filter3 = False
        filter4 = False
        filter5 = False
        filter6 = False
        ifilter_ = 0
        filter1 = CheckFilter(filterName, filterResult, 'Flag_HBHENoiseFilter')
        filter2 = CheckFilter(filterName, filterResult,
                              'Flag_globalTightHalo2016Filter')
        filter3 = CheckFilter(filterName, filterResult, 'Flag_eeBadScFilter')
        filter4 = CheckFilter(filterName, filterResult, 'Flag_goodVertices')
        filter5 = CheckFilter(filterName, filterResult,
                              'Flag_EcalDeadCellTriggerPrimitiveFilter')

        filter6 = True  #Flag_HBHENoiseIsoFilter

        if not isData:
            filterstatus = True
        if isData:
            filterstatus = filter1 & filter2 & filter3 & filter4 & filter5 & filter6
        if filterstatus == False: continue

        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        ## PFMET Selection
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------

        #        if samplename=="all":
        pfmetstatus = (pfMet > 200.0)
        #           if pfmetstatus == False : continue

        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------

        thinjetpassindex = []
        nBjets = 0
        for ithinjet in range(nTHINJets):
            j1 = thinjetP4[ithinjet]
            #if (j1.Pt() > 30.0)&(abs(j1.Eta())<2.4)&(bool(passThinJetLooseID[ithinjet])==True)&(bool(passThinJetPUID[ithinjet]) == True):
            if (j1.Pt() > 30.0) & (abs(j1.Eta()) < 2.4) & (bool(
                    passThinJetLooseID[ithinjet]) == True):
                thinjetpassindex.append(ithinjet)
            if thinJetCSV[ithinjet] > CSVMWP: nBjets += 1
#        print ('njet: ',len(thinjetpassindex))
#        if len(thinjetpassindex) < 1 : continue
#        print nBjets
#        if nBjets < 1: continue

        thindCSVjetpassindex = []
        ndBjets = 0

        try:
            for jthinjet in range(nTHINdeepCSVJets):
                j1 = thindeepCSVjetP4[jthinjet]
                #if (j1.Pt() > 30.0)&(abs(j1.Eta())<2.4)&(bool(passThinJetLooseID[ithinjet])==True)&(bool(passThinJetPUID[ithinjet]) == True):
                if (j1.Pt() > 30.0) & (abs(j1.Eta()) < 2.4) & (bool(
                        passThinJetLooseID[jthinjet]) == True):
                    thindCSVjetpassindex.append(jthinjet)
                if thinJetdeepCSV[jthinjet] > DCSVMWP: ndBjets += 1
            if len(thinjetpassindex) < 1 and len(thindCSVjetpassindex) < 1:
                continue

        except:
            if len(thinjetpassindex) < 1: continue
#        print ('njet: ',len(thinjetpassindex))
#        if len(thindCSVjetpassindex) < 1 : continue
#        print nBjets
#        if nBjets < 1: continue

# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
## Electron Veto
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        myEles = []
        for iele in range(nEle):
            if (eleP4[iele].Pt() > 10.) & (abs(eleP4[iele].Eta()) < 2.5) & (
                    bool(eleIsPassLoose[iele]) == True):
                myEles.append(iele)

        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        ## Muon Veto
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        myMuos = []
        for imu in range(nMu):
            if (muP4[imu].Pt() > 10.) & (abs(muP4[imu].Eta()) < 2.4) & (bool(
                    isLooseMuon[imu]) == True):
                relPFIso = (muChHadIso[imu] + max(
                    0., muNeHadIso[imu] + muGamIso[imu] -
                    0.5 * muPUPt[imu])) / muP4[imu].Pt()
                if relPFIso < 0.25:
                    myMuos.append(imu)

        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        ## Tau Veto
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        # ----------------------------------------------------------------------------------------------------------------------------------------------------------------
        myTaus = []
        for itau in range(nTau):
            if (tauP4[itau].Pt() > 18.) & (abs(tauP4[itau].Eta()) < 2.3) & (
                    bool(isDecayModeFinding[itau]) == True) & (bool(
                        passLooseTauIso[itau]) == True):
                myTaus.append(itau)

        st_runId[0] = long(run)
        st_lumiSection[0] = lumi
        st_eventId[0] = event
        #        print '-----------'+str(st_runId)+", "+str(st_lumiSection)+", "+str(st_eventId)

        st_pfMetCorrPt[0] = pfMet
        st_pfMetCorrPhi[0] = pfMetPhi
        st_isData[0] = isData

        st_HLT_PFMETNoMu90_PFMHTNoMu90_IDTight_v[0] = trig5
        st_HLT_PFMETNoMu120_PFMHTNoMu120_IDTight_v[0] = trig8
        st_HLT_IsoMu20[0] = trig13
        st_HLT_Ele27_WPLoose_Gsf[0] = trig14

        st_THINjetP4.clear()
        st_THINjetCISVV2.clear()
        st_THINjetHadronFlavor.clear()
        st_THINjetNHadEF.clear()
        st_THINjetCHadEF.clear()

        st_AK4deepCSVjetP4.clear()
        st_AK4deepCSVjetDeepCSV_b.clear()

        st_eleP4.clear()
        st_muP4.clear()
        st_muChHadIso.clear()
        st_muGamIso.clear()
        st_muNeHadIso.clear()
        st_HPSTau_4Momentum.clear()

        st_genParId.clear()
        st_genMomParId.clear()
        st_genParSt.clear()
        st_genParP4.clear()

        st_THINnJet[0] = len(thinjetpassindex)
        for ithinjet in thinjetpassindex:
            st_THINjetP4.push_back(thinjetP4[ithinjet])
            st_THINjetCISVV2.push_back(thinJetCSV[ithinjet])
            st_THINjetHadronFlavor.push_back(THINjetHadronFlavor[ithinjet])
            st_THINjetNHadEF.push_back(thinjetNhadEF[ithinjet])
            st_THINjetCHadEF.push_back(thinjetChadEF[ithinjet])

        try:
            st_AK4deepCSVnJet[0] = len(thindCSVjetpassindex)
            for ithinjet in thindCSVjetpassindex:
                st_AK4deepCSVjetP4.push_back(thindeepCSVjetP4[ithinjet])
                st_AK4deepCSVjetDeepCSV_b.push_back(thinJetdeepCSV[ithinjet])
        except:
            pass

        st_nEle[0] = len(myEles)
        for iele in myEles:
            st_eleP4.push_back(eleP4[iele])
            st_eleIsPassLoose.push_back(bool(eleIsPassLoose[iele]))
            st_eleIsPassMedium.push_back(bool(eleIsPassMedium[iele]))
            st_eleIsPassTight.push_back(bool(eleIsPassTight[iele]))

        st_nMu[0] = len(myMuos)
        for imu in myMuos:
            st_muP4.push_back(muP4[imu])
            st_isLooseMuon.push_back(bool(isLooseMuon[imu]))
            st_isTightMuon.push_back(bool(isTightMuon[imu]))
            st_isMediumMuon.push_back(bool(isMediumMuon[imu]))
            st_muChHadIso.push_back(muChHadIso[imu])
            st_muNeHadIso.push_back(muNeHadIso[imu])
            st_muGamIso.push_back(muGamIso[imu])
            st_muPUPt.push_back(muPUPt[imu])

        st_HPSTau_n[0] = len(myTaus)
        for itau in myTaus:
            st_HPSTau_4Momentum.push_back(tauP4[itau])

        st_pu_nTrueInt[0] = pu_nTrueInt
        #        print pu_nTrueInt
        #        print st_pu_nTrueInt[0]
        st_nGenPar[0] = nGenPar
        for igp in range(nGenPar):
            st_genParId.push_back(genParId[igp])
            st_genMomParId.push_back(genMomParId[igp])
            st_genParSt.push_back(genParSt[igp])
            st_genParP4.push_back(genParP4[igp])

        ## Fill variables for the CRs.
        WenuRecoil[0] = -1.0
        Wenumass[0] = -1.0
        WenuPhi[0] = -10.

        WmunuRecoil[0] = -1.0
        Wmunumass[0] = -1.0
        WmunuPhi[0] = -1.0

        ZeeMass[0] = -1.0
        ZeeRecoil[0] = -1.0
        ZeePhi[0] = -10.

        ZmumuMass[0] = -1.0
        ZmumuRecoil[0] = -1.0
        ZmumuPhi[0] = -10.

        TOPRecoil[0] = -1.0
        TOPPhi[0] = -10.

        ## for dielectron
        if len(myEles) >= 2:
            #            ele1 = myEles[0]
            #            ele2 = myEles[1]
            #            p4_ele1 = eleP4[ele1]
            #            p4_ele2 = eleP4[ele2]

            for iele1 in myEles:
                p4_ele1 = eleP4[iele1]
                for iele2 in myEles:
                    if iele2 > iele1 and eleCharge[iele1] * eleCharge[
                            iele2] < 0:
                        p4_ele2 = eleP4[iele2]
                        ee_mass = (p4_ele1 + p4_ele2).M()
                        zeeRecoilPx = -(pfMet * math.cos(pfMetPhi) -
                                        p4_ele1.Px() - p4_ele2.Px())
                        zeeRecoilPy = -(pfMet * math.sin(pfMetPhi) -
                                        p4_ele1.Py() - p4_ele2.Py())
                        ZeeRecoilPt = math.sqrt(zeeRecoilPx * zeeRecoilPx +
                                                zeeRecoilPy * zeeRecoilPy)
                        if ee_mass > 70.0 and ee_mass < 110.0 and ZeeRecoilPt > 200.:
                            ZeeRecoil[0] = ZeeRecoilPt
                            ZeeMass[0] = ee_mass
                            ZeePhi[0] = arctan(-zeeRecoilPx, -zeeRecoilPy)
                            break

#            ee_mass = ( p4_ele1 + p4_ele2 ).M()
#
#            #if not  ( (ee_mass > 70.0 ) & (ee_mass < 110.0) ): continue
#            if not ( eleCharge[ele1] * eleCharge[ele2] > 0 ) :
#                zeeRecoilPx = -( pfMet*math.cos(pfMetPhi) - p4_ele1.Px() - p4_ele2.Px())
#                zeeRecoilPy = -( pfMet*math.sin(pfMetPhi) - p4_ele1.Py() - p4_ele2.Py())
#                ZeeRecoil[0] =  math.sqrt(zeeRecoilPx * zeeRecoilPx  +  zeeRecoilPy*zeeRecoilPy)
#                ZeeMass[0] = ee_mass

## hardrecoil cut for ZJETS sample
#            if samplename == "ZJETS":
#               ZeeRecoilstatus =(ZeeRecoil > 200)
#               print(samplename,ZeeRecoilstatus)
#               if ZeeRecoilstatus == False : continue

## for dimu
        if len(myMuos) >= 2:
            #            mu1 = myMuos[0]
            #            mu2 = myMuos[1]
            #            p4_mu1 = muP4[mu1]
            #            p4_mu2 = muP4[mu2]
            #
            #            mumu_mass = ( p4_mu1 + p4_mu2 ).M()
            #
            #            #if not  ( (mumu_mass > 70.0 ) & (mumu_mass < 110.0) ): continue
            #            if not ( muCharge[mu1] * muCharge[mu2] > 0 ) :
            #                zmumuRecoilPx = -( pfMet*math.cos(pfMetPhi)  - p4_mu1.Px() - p4_mu2.Px())
            #                zmumuRecoilPy = -( pfMet*math.sin(pfMetPhi)  - p4_mu1.Py() - p4_mu2.Py())
            #                ZmumuRecoil[0] =  math.sqrt(zmumuRecoilPx * zmumuRecoilPx  +  zmumuRecoilPy*zmumuRecoilPy)
            #                ZmumuMass[0] = mumu_mass
            for imu1 in myMuos:
                p4_mu1 = muP4[imu1]
                for imu2 in myMuos:
                    if imu2 > imu1 and muCharge[imu1] * muCharge[imu2] < 0:
                        p4_mu2 = muP4[imu2]
                        mumu_mass = (p4_mu1 + p4_mu2).M()
                        zmumuRecoilPx = -(pfMet * math.cos(pfMetPhi) -
                                          p4_mu1.Px() - p4_mu2.Px())
                        zmumuRecoilPy = -(pfMet * math.sin(pfMetPhi) -
                                          p4_mu1.Py() - p4_mu2.Py())
                        ZmumuRecoilPt = math.sqrt(
                            zmumuRecoilPx * zmumuRecoilPx +
                            zmumuRecoilPy * zmumuRecoilPy)
                        if mumu_mass > 70.0 and mumu_mass < 110.0 and ZmumuRecoilPt > 200.:
                            ZmumuRecoil[0] = ZmumuRecoilPt
                            ZmumuMass[0] = mumu_mass
                            ZeePhi[0] = arctan(-zmumuRecoilPx, -zmumuRecoilPy)
                            break

        ## hardrecoil cut for ZJETS sample
#        if samplename == "ZJETS":
        if len(myEles) >= 2:
            ZRecoilstatus = (ZeeRecoil[0] > 200)
        elif len(myMuos) >= 2:
            ZRecoilstatus = (ZmumuRecoil[0] > 200)
        else:
            ZRecoilstatus = False
#         print(samplename,ZRecoilstatus)
#         if ZRecoilstatus == False : continue

## for Single electron
        if len(myEles) == 1:
            ele1 = myEles[0]
            p4_ele1 = eleP4[ele1]

            e_mass = MT(
                p4_ele1.Pt(), pfMet, DeltaPhi(p4_ele1.Phi(), pfMetPhi)
            )  #transverse mass defined as sqrt{2pT*MET*(1-cos(dphi)}

            #if not  ( (e_mass > 50.0 ) & (e_mass < 160.0) ): continue

            WenuRecoilPx = -(pfMet * math.cos(pfMetPhi) - p4_ele1.Px())
            WenuRecoilPy = -(pfMet * math.sin(pfMetPhi) - p4_ele1.Py())
            WenuRecoilPt = math.sqrt(WenuRecoilPx * WenuRecoilPx +
                                     WenuRecoilPy * WenuRecoilPy)
            if WenuRecoilPt > 200.:
                WenuRecoil[0] = WenuRecoilPt
                Wenumass[0] = e_mass
                WenuPhi[0] = arctan(-WenuRecoilPx, -WenuRecoilPy)

        ## hardrecoil cut for WJETS sample
#        if samplename == "WJETS":
#           WenuRecoilstatus =(WenuRecoil > 200)
#           print(samplename,WenuRecoilstatus)
#           if WenuRecoilstatus == False : continue

## for Single muon
        if len(myMuos) == 1:
            mu1 = myMuos[0]
            p4_mu1 = muP4[mu1]

            mu_mass = MT(
                p4_mu1.Pt(), pfMet, DeltaPhi(p4_mu1.Phi(), pfMetPhi)
            )  #transverse mass defined as sqrt{2pT*MET*(1-cos(dphi)}

            #if not  ( (mu_mass > 50.0 ) & (mu_mass < 160.0) ): continue

            WmunuRecoilPx = -(pfMet * math.cos(pfMetPhi) - p4_mu1.Px())
            WmunuRecoilPy = -(pfMet * math.sin(pfMetPhi) - p4_mu1.Py())
            WmunuRecoilPt = math.sqrt(WmunuRecoilPx * WmunuRecoilPx +
                                      WmunuRecoilPy * WmunuRecoilPy)
            if WmunuRecoilPt > 200.:
                WmunuRecoil[0] = WmunuRecoilPt
                Wmunumass[0] = mu_mass
                WmunuPhi[0] = arctan(-WmunuRecoilPx, -WmunuRecoilPy)

        ## hardrecoil cut for WJETS sample
#        if samplename == "WJETS":
        if len(myEles) == 1:
            WRecoilstatus = (WenuRecoil[0] > 200)
        elif len(myMuos) == 1:
            WRecoilstatus = (WmunuRecoil[0] > 200)
        else:
            WRecoilstatus = False
#            print(samplename,WRecoilstatus)
#            if WRecoilstatus == False : continue

## for Single electron && Single Muon
        if len(myEles) >= 1 and len(myMuos) >= 1:
            #           ele1 = myEles[0]
            #           p4_ele1 = eleP4[ele1]
            #           mu1 = myMuos[0]
            #           p4_mu1 = muP4[mu1]
            #
            #           #e_mass = MT(p4_ele1.Pt(),pfMet, DeltaPhi(p4_ele1.Phi(),pfMetPhi)) #transverse mass defined as sqrt{2pT*MET*(1-cos(dphi)}
            #           #mu_mass = MT(p4_mu1.Pt(),pfMet, DeltaPhi(p4_mu1.Phi(),pfMetPhi))
            #           #if not  ( (e_mass > 50.0 ) & (e_mass < 160.0) ): continue
            #
            #           TOPenumunuRecoilPx = -( pfMet*math.cos(pfMetPhi) - p4_mu1.Px() -p4_ele1.Px())
            #           TOPenumunuRecoilPy = -( pfMet*math.sin(pfMetPhi) - p4_mu1.Py() -p4_ele1.Py())
            #           TOPRecoil[0] =  math.sqrt(TOPenumunuRecoilPx * TOPenumunuRecoilPx  +  TOPenumunuRecoilPy*TOPenumunuRecoilPy)
            for iele in myEles:
                p4_ele1 = eleP4[iele]
                for imu in myMuos:
                    p4_mu1 = muP4[imu]
                    TOPenumunuRecoilPx = -(pfMet * math.cos(pfMetPhi) -
                                           p4_mu1.Px() - p4_ele1.Px())
                    TOPenumunuRecoilPy = -(pfMet * math.sin(pfMetPhi) -
                                           p4_mu1.Py() - p4_ele1.Py())
                    TOPenumunuRecoilPt = math.sqrt(
                        TOPenumunuRecoilPx * TOPenumunuRecoilPx +
                        TOPenumunuRecoilPy * TOPenumunuRecoilPy)
                    if TOPenumunuRecoilPt > 200:
                        TOPRecoil[0] = TOPenumunuRecoilPt
                        TOPPhi[0] = arctan(-TOPenumunuRecoilPx,
                                           -TOPenumunuRecoilPy)
                        break

        TOPRecoilstatus = (TOPRecoil[0] > 200)

        if pfmetstatus == False and ZRecoilstatus == False and WRecoilstatus == False and TOPRecoilstatus == False:
            continue


#        if ZRecoilstatus:
#            print ('Z: ',nEle, nMu, ZeeMass[0], ZmumuMass[0])
#        if WRecoilstatus:
#            print ('W: ', Wenumass[0], Wmunumass[0])
#        if TOPRecoilstatus:
#            print ('T: ',nEle, nMu, TOPenumunuRecoilPt)

        outTree.Fill()

    h_total_mcweight.Write()
    h_total.Write()
    samplepath.Write()
    outfile.Write()
Beispiel #11
0
 def tag(self, name, content):
     named = TNamed(name, content)
     oldDir = gDirectory
     self.file.cd()
     named.Write()
     oldDir.cd()
Beispiel #12
0
 def __init__(self, key, value):
     """Initialize a string value."""
     TNamed.__init__(self, key, value)
                # Declare the file
                file_Output_S = TFile(
                    "%s_VFAT%s_ID_%s_ScurveOutput.root" %
                    (TestName, pos, port), "RECREATE", "", 1)
                dir_S_ChipID = file_Output_S.mkdir("ChipID")
                dir_S_Thresholds = file_Output_S.mkdir("Thresholds")
                dir_S_TrimDACValues = file_Output_S.mkdir("TrimDACValues")
                dir_S_SCurveMeanByChan = file_Output_S.mkdir(
                    "SCurveMeanByChannel")
                dir_S_SCurveSigma = file_Output_S.mkdir("SCurvesSigma")
                dir_S_SCurveByChan = file_Output_S.mkdir("SCurveByChannel")
                dir_S_SCurveExample = file_Output_S.mkdir("SCurveExample")
                dir_S_SCurveSeparation = file_Output_S.mkdir(
                    "SCurveSeparation")
                chipID = TNamed("VFAT%s" % (pos), str(port))
                dir_A_ChipID.cd()
                chipID.Write()
                dir_S_ChipID.cd()
                chipID.Write()
                Canvas = TCanvas()

                #Number of the example channel for which the SCURVE and its fit are printed.
                SCUVRE = 15
                # build a rectangle in axes coords
                VCALmean14 = [
                ]  #to plot VCal means for one channel (SCUVRE) for all chips
                VCALcov14 = [
                ]  #to plot VCal covs for one channel (SCUVRE) for all chips

                Threshold_2 = False
Beispiel #14
0
def main(argv):

    print("=======================================")
    print("LAT started:",time.strftime('%X %x %Z'))
    startT = time.clock()
    # gROOT.ProcessLine("gErrorIgnoreLevel = 3001;") # suppress ROOT error messages
    global batMode
    intMode, batMode, rangeMode, fileMode, gatMode, singleMode, pathMode, cutMode = False, False, False, False, False, False, False, False
    dontUseTCuts = False
    dsNum, subNum, runNum, plotNum = -1, -1, -1, 1
    pathToInput, pathToOutput, manualInput, manualOutput, customPar = ".", ".", "", "", ""

    if len(argv)==0: return
    for i,opt in enumerate(argv):
        if opt == "-r":
            rangeMode, dsNum, subNum = True, int(argv[i+1]), int(argv[i+2])
            print("Scanning DS-%d sub-range %d" % (dsNum, subNum))
        if opt == "-p":
            pathMode, manualInput, manualOutput = True, argv[i+1], argv[i+2]
            print("Manually set input/output files:\nInput: %s\nOutput: %s" % (manualInput, manualOutput))
        if opt == "-d":
            pathToInput, pathToOutput = argv[i+1], argv[i+2]
            print("Custom paths: Input %s,  Output %s" % (pathToInput,pathToOutput))
        if opt == "-f":
            fileMode, dsNum, runNum = True, int(argv[i+1]), int(argv[i+2])
            print("Scanning DS-%d, run %d" % (dsNum, runNum))
        if opt == "-g":
            gatMode, runNum = True, int(argv[i+1])
            print("GATDataSet mode.  Scanning run %d" % (runNum))
        if opt == "-s":
            singleMode, pathToInput = True, argv[i+1]
            print("Single file mode.  Scanning {}".format(pathToInput))
        if opt == "-i":
            intMode, plotNum = True, int(argv[i+1])
            print("Interactive mode selected. Use \"p\" for previous and \"q\" to exit.")
        if opt == "-x":
            dontUseTCuts = True
            print("DC TCuts deactivated.  Retaining all events ...")
        if opt == "-c":
            cutMode, customPar = True, str(argv[i+1])
            print("Using custom cut parameter: {}".format(customPar))
        if opt == "-b":
            batMode = True
            import matplotlib
            # if os.environ.get('DISPLAY','') == '':
                # print('No display found. Using non-interactive Agg backend')
            matplotlib.use('Agg')
            print("Batch mode selected.  A new file will be created.")
    import matplotlib.pyplot as plt
    from matplotlib import gridspec
    import matplotlib.ticker as mtick
    plt.style.use('pltTalks.mplstyle')
    from matplotlib.colors import LogNorm, Normalize

    # File I/O
    inFile, outFile, bltFile = TFile(), TFile(), TFile()
    gatTree, bltTree, out = TTree(), TTree(), TTree()
    theCut, inPath, outPath = "", "", ""

    # Set input and output files
    if rangeMode:
        inPath = "%s/waveSkimDS%d_%d.root" % (pathToInput, dsNum, subNum)
        outPath = "%s/latSkimDS%d_%d.root" % (pathToOutput, dsNum, subNum)
    if fileMode:
        inPath = "%s/waveSkimDS%d_run%d.root" % (pathToInput, dsNum, runNum)
        outPath = "%s/latSkimDS%d_run%d.root" % (pathToOutput, dsNum, runNum)
    if pathMode:
        inPath, outPath = manualInput, manualOutput
    if gatMode:
        ds = GATDataSet()
        gatPath = ds.GetPathToRun(runNum,GATDataSet.kGatified)
        bltPath = ds.GetPathToRun(runNum,GATDataSet.kBuilt)
        outPath = "%s/lat_run%d.root" % (pathToOutput, runNum)
    if pathMode and gatMode:
        outPath = manualOutput


    # Initialize trees
    if rangeMode or fileMode or pathMode:
        inFile = TFile(inPath)
        gatTree = inFile.Get("skimTree")
        print(gatTree.GetEntries(),"entries in input tree.")
    elif gatMode:
        inFile = TFile(gatPath)
        bltFile = TFile(bltPath)
        gatTree = inFile.Get("mjdTree")
        bltTree = bltFile.Get("MGTree")
        gatTree.AddFriend(bltTree)
    if singleMode:
        inFile = TFile(pathToInput)
        gatTree = inFile.Get("skimTree")

    # apply cut to tree
    if (rangeMode or fileMode or pathMode) and not dontUseTCuts:
        try:
            theCut = inFile.Get("theCut").GetTitle()
        except ReferenceError:
            theCut = ""
    if cutMode:
        # theCut += customPar
        # theCut = "(channel==672 || channel==674) && mH==2" # sync chan: 672, extp chan: 674
        # theCut += " && fitSlo < 10"
        # theCut = "trapENFCal > 1 && trapENFCal < 10 && riseNoise > 2"
        theCut = "trapENFCal > 20 && trapENFCal < 100 && riseNoise > 2"
        print("WARNING: Custom cut in use! : ",theCut)

    gatTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    gatTree.SetEntryList(elist)
    nList = elist.GetN()
    print("Using cut:\n",theCut)
    print("Found",gatTree.GetEntries(),"input entries.")
    print("Found",nList,"entries passing cuts.")

    # Output: In batch mode (-b) only, create an output file+tree & append new branches.
    if batMode and not intMode:
        outFile = TFile(outPath, "RECREATE")
        print("Attempting tree copy to",outPath)
        out = gatTree.CopyTree("")
        out.Write()
        print("Wrote",out.GetEntries(),"entries.")
        cutUsed = TNamed("theCut",theCut)
        cutUsed.Write()

    waveS1, waveS2 = std.vector("double")(), std.vector("double")()
    waveS3, waveS4, waveS5 = std.vector("double")(), std.vector("double")(), std.vector("double")()
    bcMax, bcMin = std.vector("double")(), std.vector("double")()
    bandMax, bandTime = std.vector("double")(), std.vector("double")()
    den10, den50, den90 = std.vector("double")(), std.vector("double")(), std.vector("double")()
    oppie = std.vector("double")()
    fitMu, fitAmp, fitSlo = std.vector("double")(), std.vector("double")(), std.vector("double")()
    fitTau, fitBL = std.vector("double")(), std.vector("double")()
    matchMax, matchWidth, matchTime = std.vector("double")(), std.vector("double")(), std.vector("double")()
    pol0, pol1, pol2, pol3 = std.vector("double")(), std.vector("double")(), std.vector("double")(), std.vector("double")()
    fails, fitChi2, fitLL = std.vector("int")(), std.vector("double")(), std.vector("double")()
    riseNoise = std.vector("double")()
    t0_SLE, t0_ALE, lat, latF = std.vector("double")(), std.vector("double")(), std.vector("double")(), std.vector("double")()
    latAF, latFC, latAFC = std.vector("double")(), std.vector("double")(), std.vector("double")()
    nMS = std.vector("int")()
    tE50, latE50, wfStd = std.vector("double")(), std.vector("double")(), std.vector("double")()
    wfAvgBL, wfRMSBL = std.vector("double")(), std.vector("double")()
    fitErr = std.vector("int")()

    # It's not possible to put the "out.Branch" call into a class initializer (waveLibs::latBranch). You suck, ROOT.
    b1, b2 = out.Branch("waveS1",waveS1), out.Branch("waveS2",waveS2)
    b3, b4, b5 = out.Branch("waveS3",waveS3), out.Branch("waveS4",waveS4), out.Branch("waveS5",waveS5)
    b7, b8 = out.Branch("bcMax",bcMax), out.Branch("bcMin",bcMin)
    b9, b10 = out.Branch("bandMax",bandMax), out.Branch("bandTime",bandTime)
    b11, b12, b13 = out.Branch("den10",den10), out.Branch("den50",den50), out.Branch("den90",den90)
    b14 = out.Branch("oppie",oppie)
    b15, b16, b17 = out.Branch("fitMu", fitMu), out.Branch("fitAmp", fitAmp), out.Branch("fitSlo", fitSlo)
    b18, b19 = out.Branch("fitTau",fitTau), out.Branch("fitBL",fitBL)
    b20, b21, b22 = out.Branch("matchMax", matchMax), out.Branch("matchWidth", matchWidth), out.Branch("matchTime", matchTime)
    b23, b24, b25, b26 = out.Branch("pol0", pol0), out.Branch("pol1", pol1), out.Branch("pol2", pol2), out.Branch("pol3", pol3)
    b27, b28, b29 = out.Branch("fails",fails), out.Branch("fitChi2",fitChi2), out.Branch("fitLL",fitLL)
    b30 = out.Branch("riseNoise",riseNoise)
    b31, b32, b33, b34 = out.Branch("t0_SLE",t0_SLE), out.Branch("t0_ALE",t0_ALE), out.Branch("lat",lat), out.Branch("latF",latF)
    b35, b36, b37 = out.Branch("latAF",latAF), out.Branch("latFC",latFC), out.Branch("latAFC",latAFC)
    b38 = out.Branch("nMS",nMS)
    b39, b40, b41 = out.Branch("tE50", tE50), out.Branch("latE50", latE50), out.Branch("wfStd", wfStd)
    b42, b43 = out.Branch("wfAvgBL", wfAvgBL), out.Branch("wfRMSBL", wfRMSBL)
    b44 = out.Branch("fitErr",fitErr)

    # make a dictionary that can be iterated over (avoids code repetition in the loop)
    brDict = {
        "waveS1":[waveS1, b1], "waveS2":[waveS2, b2],
        "waveS3":[waveS3, b3], "waveS4":[waveS4, b4], "waveS5":[waveS5, b5],
        "bcMax":[bcMax, b7], "bcMin":[bcMin, b8],
        "bandMax":[bandMax, b9], "bandTime":[bandTime, b10],
        "den10":[den10, b11], "den50":[den50, b12], "den90":[den90, b13],
        "oppie":[oppie, b14],
        "fitMu":[fitMu, b15], "fitAmp":[fitAmp, b16], "fitSlo":[fitSlo, b17],
        "fitTau":[fitTau, b18], "fitBL":[fitBL,b19],
        "matchMax":[matchMax, b20], "matchWidth":[matchWidth, b21], "matchTime":[matchTime, b22],
        "pol0":[pol0, b23], "pol1":[pol1, b24], "pol2":[pol2, b25], "pol3":[pol3, b26],
        "fails":[fails,b27], "fitChi2":[fitChi2,b28], "fitLL":[fitLL,b29],
        "riseNoise":[riseNoise,b30],
        "t0_SLE":[t0_SLE,b31], "t0_ALE":[t0_ALE,b32], "lat":[lat,b33], "latF":[latF,b34],
        "latAF":[latAF,b35], "latFC":[latFC,b36], "latAFC":[latAFC,b37],
        "nMS":[nMS,b38], "tE50":[tE50,b39], "latE50":[latE50,b40], "wfStd":[wfStd,b41],
        "wfAvgBL":[wfAvgBL,b42], "wfRMSBL":[wfRMSBL,b43],
        "fitErr":[fitErr,b44]
    }

    # Make a figure (-i option: select different plots)
    # fig = plt.figure(figsize=(12,9), facecolor='w')
    fig = plt.figure()
    if plotNum==0 or plotNum==7 or plotNum==8:
        p0 = plt.subplot(111)  # 0-raw waveform, 7-new trap filters
    elif plotNum==1 or plotNum==2:
        p0 = plt.subplot(211)  # 1-wavelet, 2-time points, bandpass filters, tail slope
        p1 = plt.subplot(212)
    elif plotNum==3:
        p0 = plt.subplot2grid((2,5), (0,0), colspan=3)  # oppie / freq-domain matched filter
        p1 = plt.subplot2grid((2,5), (0,3), colspan=2)
        p2 = plt.subplot2grid((2,5), (1,0), colspan=3)
    elif plotNum==4:
        p0 = plt.subplot(111)  # time-domain matched filter
    elif plotNum==5:
        p0 = plt.subplot(111)  # bandpass / bandTime
    elif plotNum==6:
        p0 = plt.subplot2grid((6,10), (0,0), colspan=10, rowspan=3) # waveform fit
        p1 = plt.subplot2grid((6,10), (3,0), colspan=10, rowspan=1) # residual
        p2 = plt.subplot2grid((6,10), (4,0), colspan=2, rowspan=2) # traces
        p3 = plt.subplot2grid((6,10), (4,2), colspan=2, rowspan=2)
        p4 = plt.subplot2grid((6,10), (4,4), colspan=2, rowspan=2)
        p5 = plt.subplot2grid((6,10), (4,6), colspan=2, rowspan=2)
        p6 = plt.subplot2grid((6,10), (4,8), colspan=2, rowspan=2)
    elif plotNum==9:
        p0 = plt.subplot2grid((5,1), (0,0)) # 9- wpt on wf fit residual
        p1 = plt.subplot2grid((5,1), (1,0), rowspan=2)
        p2 = plt.subplot2grid((5,1), (3,0), rowspan=2)
    if not batMode: plt.show(block=False)


    # Load a fast signal template - used w/ the freq-domain matched filter
    # print("Generating signal template ...")
    tSamp, tR, tZ, tAmp, tST, tSlo = 5000, 0, 15, 100, 2500, 10
    # tOrig, tOrigTS = wl.MakeSiggenWaveform(tSamp,tR,tZ,tAmp,tST,tSlo) # Damn you to hell, PDSF
    templateFile = np.load("%s/data/lat_template.npz" % os.environ['LATDIR'])
    if dsNum==2 or dsNum==6: templateFile = np.load("%s/data/lat_ds2template.npz" % os.environ['LATDIR'])
    tOrig, tOrigTS = templateFile['arr_0'], templateFile['arr_1']


    # Load stuff from DS1 forced acq. runs
    npzfile = np.load("%s/data/fft_forcedAcqDS1.npz" % os.environ['LATDIR'])
    noise_asd, noise_xFreq, avgPwrSpec, xPwrSpec, data_forceAcq, data_fft = npzfile['arr_0'],npzfile['arr_1'],npzfile['arr_2'],npzfile['arr_3'],npzfile['arr_4'],npzfile['arr_5']


    # Loop over events
    print("Starting event loop ...")
    iList = -1
    while True:
        iList += 1
        if intMode==True and iList != 0:
            value = input()
            if value=='q': break        # quit
            if value=='p': iList -= 2   # go to previous
            if (value.isdigit()):
                iList = int(value)      # go to entry number
        elif intMode==False and batMode==False:
            plt.pause(0.00001)          # rapid-draw mode
        if iList >= nList: break        # bail out, goose!

        entry = gatTree.GetEntryNumber(iList);
        gatTree.LoadTree(entry)
        gatTree.GetEntry(entry)
        nChans = gatTree.channel.size()
        event = MGTEvent()
        if gatMode: event = bltTree.event

        # Reset all branch vectors
        # NOTE: The events sometimes contain 'straggler' hits that do not pass the
        # given TCut.  This line sets ALL the new parameters to -88888 by default.
        # If you see this value in a plot, then you must be including hits that
        # passed the cut in wave-skim but did not pass the (different?) cut in LAT.
        for key in brDict: brDict[key][0].assign(nChans,-88888)
        brDict["fails"][0].assign(nChans,0) # set error code to 'true' by default
        errorCode = [0,0,0,0]


        # Loop over hits passing cuts
        numPass = gatTree.Draw("channel",theCut,"GOFF",1,iList)
        chans = gatTree.GetV1()
        chanList = list(set(int(chans[n]) for n in range(numPass)))
        hitList = (iH for iH in range(nChans) if gatTree.channel.at(iH) in chanList)  # a 'generator expression'
        for iH in hitList:

            # ------------------------------------------------------------------------
            # Waveform processing

            # load data
            run = gatTree.run
            chan = gatTree.channel.at(iH)
            dataENFCal = gatTree.trapENFCal.at(iH)
            dataENM = gatTree.trapENM.at(iH)
            dataTSMax = gatTree.trapENMSample.at(iH)*10. - 4000
            wf = MGTWaveform()
            iEvent = 0
            if gatMode:
                wf = event.GetWaveform(iH)
                iEvent = entry
            else:
                wf = gatTree.MGTWaveforms.at(iH)
                iEvent = gatTree.iEvent

            # print("%d:  run %d  chan %d  trapENFCal %.2f" % (iList, run, chan, dataENFCal))

            # be absolutely sure you're matching the right waveform to this hit
            if wf.GetID() != chan:
                print("ERROR -- Vector matching failed.  iList %d  run %d  iEvent %d" % (iList,run,iEvent))
                return

            # Let's start the show - grab a waveform.
            # Remove first 4 samples when we have multisampling
            # Remove last 2 samples to get rid of the ADC spike at the end of all wf's.
            truncLo, truncHi = 0, 2
            if dsNum==6 or dsNum==2: truncLo = 4
            signal = wl.processWaveform(wf,truncLo,truncHi)
            data = signal.GetWaveRaw()
            data_blSub = signal.GetWaveBLSub()
            dataTS = signal.GetTS()
            dataBL,dataNoise = signal.GetBaseNoise()

            # wavelet packet transform
            wp = pywt.WaveletPacket(data_blSub, 'db2', 'symmetric', maxlevel=4)
            nodes = wp.get_level(4, order='freq')
            wpCoeff = np.array([n.data for n in nodes],'d')
            wpCoeff = abs(wpCoeff)

            # wavelet parameters
            # First get length of wavelet on the time axis, the scale axis will always be the same
            # due to the number of levels in the wavelet
            wpLength = len(wpCoeff[1,:])
            waveS1[iH] = np.sum(wpCoeff[0:1,1:wpLength//4+1]) # python3 : floor division (//) returns an int
            waveS2[iH] = np.sum(wpCoeff[0:1,wpLength//4+1:wpLength//2+1])
            waveS3[iH] = np.sum(wpCoeff[0:1,wpLength//2+1:3*wpLength//4+1])
            waveS4[iH] = np.sum(wpCoeff[0:1,3*wpLength//4+1:-1])
            waveS5[iH] = np.sum(wpCoeff[2:-1,1:-1])
            S6 = np.sum(wpCoeff[2:9,1:wpLength//4+1])
            S7 = np.sum(wpCoeff[2:9,wpLength//4+1:wpLength//2+1])
            S8 = np.sum(wpCoeff[2:9,wpLength//2+1:3*wpLength//4+1])
            S9 = np.sum(wpCoeff[2:9,3*wpLength//4+1:-1])
            S10 = np.sum(wpCoeff[9:,1:wpLength//4+1])
            S11 = np.sum(wpCoeff[9:,wpLength//4+1:wpLength//2+1])
            S12 = np.sum(wpCoeff[9:,wpLength//2+1:3*wpLength//4+1])
            S13 = np.sum(wpCoeff[9:,3*wpLength//4+1:-1])
            sumList = [S6, S7, S8, S9, S10, S11, S12, S13]
            bcMax[iH] = np.max(sumList)
            bcMin[iH] = 1. if np.min(sumList) < 1 else np.min(sumList)

            # reconstruct waveform w/ only lowest frequency.
            new_wp = pywt.WaveletPacket(data=None, wavelet='db2', mode='symmetric')
            new_wp['aaa'] = wp['aaa'].data
            data_wlDenoised = new_wp.reconstruct(update=False)
            # resize in a smart way
            diff = len(data_wlDenoised) - len(data_blSub)
            if diff > 0: data_wlDenoised = data_wlDenoised[diff:]

            # waveform high/lowpass filters - parameters are a little arbitrary

            B1,A1 = butter(2, [1e5/(1e8/2),1e6/(1e8/2)], btype='bandpass')
            data_bPass = lfilter(B1, A1, data_blSub)

            # used in the multisite tagger
            B2, A2 = butter(1, 0.08)
            data_filt = filtfilt(B2, A2, data_blSub)
            data_filtDeriv = wl.wfDerivative(data_filt)
            filtAmp = np.amax(data_filtDeriv) # scale the max to match the amplitude
            data_filtDeriv = data_filtDeriv * (dataENM / filtAmp)

            B3, A3 = butter(2,1e6/(1e8/2), btype='lowpass')
            data_lPass = lfilter(B3, A3, data_blSub)

            idx = np.where((dataTS > dataTS[0]+100) & (dataTS < dataTS[-1]-100))
            windowingOffset = dataTS[idx][0] - dataTS[0]

            bandMax[iH] = np.amax(data_bPass[idx])
            bandTime[iH] = dataTS[ np.argmax(data_bPass[idx])] - windowingOffset


            # timepoints of low-pass waveforms
            tpc = MGWFTimePointCalculator();
            tpc.AddPoint(.2)
            tpc.AddPoint(.5)
            tpc.AddPoint(.9)
            mgtLowPass = wl.MGTWFFromNpArray(data_lPass)
            tpc.FindTimePoints(mgtLowPass)
            den10[iH] = tpc.GetFromStartRiseTime(0)*10
            den50[iH] = tpc.GetFromStartRiseTime(1)*10
            den90[iH] = tpc.GetFromStartRiseTime(2)*10


            # ================ xgauss waveform fitting ================

            amp, mu, sig, tau, bl = dataENM, dataTSMax, 600., -72000., dataBL
            floats = np.asarray([amp, mu, sig, tau, bl])
            temp = xgModelWF(dataTS, floats)
            if not batMode: MakeTracesGlobal()

            # get the noise of the denoised wf
            denoisedNoise,_,_ = wl.baselineParameters(data_wlDenoised)

            # NOTE: fit is to wavelet-denoised data, BECAUSE there are no HF components in the model,
            # AND we'll still calculate fitChi2 w/r/t the data, not the denoised data.
            # datas = [dataTS, data, dataNoise] # fit data
            datas = [dataTS, data_wlDenoised + dataBL, denoisedNoise] # fit wavelet-denoised data w/ Bl added back in

            # Set bounds - A,mu,sig,tau,bl.
            # bnd = ((None,None),(None,None),(None,None),(None,None),(None,None))   # often gets caught at sig=0
            bnd = ((None,None),(None,None),(2.,None),(-72001.,-71999.),(None,None)) # gets caught much less often.

            # L-BGFS-B with numerical gradient.
            start = time.clock()
            result = op.minimize(lnLike, floats, args=datas, method="L-BFGS-B", options=None, bounds=bnd)
            fitSpeed = time.clock() - start

            fitErr[iH] = 0
            if not result["success"]:
                # print("fit fail: ", result["message"])
                fitErr[iH] = 1
                errorCode[0] = 1

            amp, mu, sig, tau, bl = result["x"]

            # save parameters

            fitMu[iH], fitAmp[iH], fitSlo[iH], fitTau[iH], fitBL[iH] = mu, amp, sig, tau, bl
            floats = np.asarray([amp, mu, sig, tau, bl])
            fit = xgModelWF(dataTS, floats)

            # print("%d/%d iH %d  e %-10.2f  fs %-8.2f  f %d" % (iList, nList, iH, dataENFCal, fitSlo[iH], fitErr[iH]))

            # log-likelihood of this fit
            fitLL[iH] = result["fun"]

            # chi-square of this fit
            # Textbook is (observed - expected)^2 / expected,
            # but we'll follow MGWFCalculateChiSquare.cc and do (observed - expected)^2 / NDF.
            # NOTE: we're doing the chi2 against the DATA, though the FIT is to the DENOISED DATA.
            fitChi2[iH] = np.sum(np.square(data-fit)) / (len(data)-1)/dataNoise


            # get wavelet coeff's for rising edge only.  normalize to bcMin
            # view this w/ plot 1

            # find the window of rising edge
            fit_blSub = fit - bl
            fitMaxTime = dataTS[np.argmax(fit_blSub)]
            fitStartTime = dataTS[0]
            idx = np.where(fit_blSub < 0.1)
            if len(dataTS[idx] > 0): fitStartTime = dataTS[idx][-1]
            fitRiseTime50 = (fitMaxTime + fitStartTime)/2.

            # bcMin is 32 samples long in the x-direction.
            # if we make the window half as wide, it'll have the same # of coeff's as bcMin.
            # this is still 'cheating' since we're not summing over the same rows.
            numXRows = wpCoeff.shape[1]
            wpCtrRise = int((fitRiseTime50 - dataTS[0]) / (dataTS[-1] - dataTS[0]) * numXRows)
            wpLoRise = wpCtrRise - 8
            if wpLoRise < 0: wpLoRise = 0
            wpHiRise = wpCtrRise + 8
            if wpHiRise > numXRows: wpHiRise = numXRows

            # sum all HF wavelet components for this edge.
            riseNoise[iH] = np.sum(wpCoeff[2:-1,wpLoRise:wpHiRise]) / bcMin[iH]

            # print("%d %d %d %d e %-5.2f  bmax %-6.2f  bmin %-6.2f  mu %-5.2f  a %-5.2f  s %-5.2f  bl %-5.2f  rn %.2f" % (run,iList,iH,chan,dataENFCal,bcMax[iH],bcMin[iH],fitMu[iH],fitAmp[iH],fitSlo[iH],fitBL[iH],riseNoise[iH]))

            # =========================================================

            # optimal matched filter (freq. domain)
            # we use the pysiggen fast template (not the fit result) to keep this independent of the wf fitter.

            # pull in the template, shift it, and make sure it's the same length as the data
            guessTS = tOrigTS - 15000.
            idx = np.where((guessTS > -5) & (guessTS < dataTS[-1]))
            guessTS, guess = guessTS[idx], tOrig[idx]
            if len(guess)!=len(data):
                if len(guess)>len(data):
                    guess, guessTS = guess[0:len(data)], guessTS[0:len(data)]
                else:
                    guess = np.pad(guess, (0,len(data)-len(guess)), 'edge')
                    guessTS = np.pad(guessTS, (0,len(data)-len(guessTS)), 'edge')

            data_fft = np.fft.fft(data_blSub) # can also try taking fft of the low-pass data
            temp_fft = np.fft.fft(guess)

            datafreq = np.fft.fftfreq(data.size) * 1e8
            power_vec = np.interp(datafreq, noise_xFreq, noise_asd) # load power spectra from file

            # Apply the filter
            optimal = data_fft * temp_fft.conjugate() / power_vec
            optimal_time = 2 * np.fft.ifft(optimal)

            # Normalize the output
            df = np.abs(datafreq[1] - datafreq[0]) # freq. bin size
            sigmasq = 2 * (temp_fft * temp_fft.conjugate() / power_vec).sum() * df
            sigma = np.sqrt(np.abs(sigmasq))
            SNR = abs(optimal_time) / (sigma)
            oppie[iH] = np.amax(SNR)


            # time-domain matched filter.  use the baseline-subtracted wf as data, and fit_blSub too.

            # make a longer best-fit waveform s/t it can be shifted L/R.
            matchTS = np.append(dataTS, np.arange(dataTS[-1], dataTS[-1] + 20000, 10)) # add 2000 samples
            match = xgModelWF(matchTS, [amp, mu+10000., sig, tau, bl]) # shift mu accordingly
            match = match[::-1] - bl # time flip and subtract off bl

            # line up the max of the 'match' (flipped wf) with the max of the best-fit wf
            # this kills the 1-1 matching between matchTS and dataTS (each TS has some offset)
            matchMaxTime = matchTS[np.argmax(match)]
            matchTS = matchTS + (fitMaxTime - matchMaxTime)

            # resize match, matchTS to have same # samples as data, dataTS.
            # this is the only case we really care about
            # ("too early" and "too late" also happen, but the shift is larger than the trigger walk, making it unphysical)
            if matchTS[0] <= dataTS[0] and matchTS[-1] >= dataTS[-1]:
                idx = np.where((matchTS >= dataTS[0]) & (matchTS <= dataTS[-1]))
                match, matchTS = match[idx], matchTS[idx]
                sizeDiff = len(dataTS)-len(matchTS)
                if sizeDiff < 0:
                    match, matchTS = match[:sizeDiff], matchTS[:sizeDiff]
                elif sizeDiff > 0:
                    match = np.hstack((match, np.zeros(sizeDiff)))
                    matchTS = np.hstack((matchTS, dataTS[-1*sizeDiff:]))
                if len(match) != len(data):
                    print("FIXME: match filter array manip is still broken.")

            # compute match filter parameters
            matchMax[iH], matchWidth[iH], matchTime[iH] = -888, -888, -888
            if len(match)==len(data):
                smoothMF = gaussian_filter(match * data_blSub, sigma=5.)
                matchMax[iH] = np.amax(smoothMF)
                matchTime[iH] = matchTS[ np.argmax(smoothMF) ]
                idx = np.where(smoothMF > matchMax[iH]/2.)
                if len(matchTS[idx]>1):
                    matchWidth[iH] = matchTS[idx][-1] - matchTS[idx][0]


            # Fit tail slope to polynomial.  Guard against fit fails

            idx = np.where(dataTS >= fitMaxTime)
            tail, tailTS = data[idx], dataTS[idx]
            popt1,popt2 = 0,0
            try:
                popt1,_ = op.curve_fit(wl.tailModelPol, tailTS, tail)
                pol0[iH], pol1[iH], pol2[iH], pol3[iH] = popt1[0], popt1[1], popt1[2], popt1[3]
            except:
                # print("curve_fit tailModelPol failed, run %i  event %i  channel %i" % (run, iList, chan))
                errorCode[2] = 1
                pass

            # =========================================================

            # new trap filters.
            # params: t0_SLE, t0_ALE, lat, latF, latAF, latFC, latAFC

            # calculate trapezoids

            # standard trapezoid - prone to walking, less sensitive to noise.  use to find energy
            eTrap = wl.trapFilter(data_blSub, 400, 250, 7200.)
            eTrapTS = np.arange(0, len(eTrap)*10., 10)
            eTrapInterp = interpolate.interp1d(eTrapTS, eTrap)

            # short trapezoid - triggers more quickly, sensitive to noise.  use to find t0
            sTrap = wl.trapFilter(data_blSub, 100, 150, 7200.)
            sTrapTS = np.arange(0, len(sTrap)*10., 10)

            # asymmetric trapezoid - used to find the t0 only
            aTrap = wl.asymTrapFilter(data_blSub, 4, 10, 200, True) # (0.04us, 0.1us, 2.0us)
            aTrapTS = np.arange(0, len(aTrap)*10., 10)

            # find leading edges (t0 times)

            # limit the range from 0 to 10us, and use an ADC threshold of 1.0 as suggested by DCR
            t0_SLE[iH],_ = wl.walkBackT0(sTrap, eTrapTS[-1]+7000-4000-2000, 1., 0, 1000) # (in ns) finds leading edge from short trap
            t0_ALE[iH],_ = wl.walkBackT0(aTrap, eTrapTS[-1]+7000-4000-2000, 1., 0, 1000) # (in ns) finds leading edge from asymmetric trap

            # standard energy trapezoid w/ a baseline padded waveform
            data_pad = np.pad(data_blSub,(200,0),'symmetric')
            pTrap = wl.trapFilter(data_pad, 400, 250, 7200.)
            pTrapTS = np.linspace(0, len(pTrap)*10, len(pTrap))
            pTrapInterp = interpolate.interp1d(pTrapTS, pTrap)

            # calculate energy parameters
            # standard amplitude.  basically trapEM, but w/o NL correction if the input WF doesn't have it.
            lat[iH] = np.amax(eTrap)

            # Calculate DCR suggested amplitude, using the 50% to the left and right of the maximum point
            t0_F50,t0fail1 = wl.walkBackT0(pTrap, thresh=lat[iH]*0.5, rmin=0, rmax=len(pTrap)-1)
            t0_B50,t0fail2 = wl.walkBackT0(pTrap, thresh=lat[iH]*0.5, rmin=0, rmax=len(pTrap)-1, forward=True)
            t0_E50 = (t0_F50 + t0_B50)/2.0

            #TODO -- if it's necessary due to the trigger walk, we could potentially add a way to recursively increase the threshold until a timepoint is found, however it will still always fail for most noise events
            if not t0fail1 or not t0fail2:
                latE50[iH] = 0 # Set amplitude to 0 if one of the evaluations failed
            else:
                latE50[iH] = pTrapInterp(t0_E50) # Maybe I should call this latDCR50 to confuse people
            tE50[iH] = t0_B50 - t0_F50 # Save the difference between the middle points, can be used as a cut later

            # standard amplitude with t0 from the shorter traps
            # If either fixed pickoff time (t0) is < 0, use the first sample as the amplitude (energy).
            latF[iH] = eTrapInterp( np.amax([t0_SLE[iH]-7000+4000+2000, 0.]) ) # This should be ~trapEF
            latAF[iH] = eTrapInterp( np.amax([t0_ALE[iH]-7000+4000+2000, 0.]) )

            # amplitude from padded trapezoid, with t0 from short traps and a correction function
            # function is under development.  currently: f() = exp(p0 + p1*E), p0 ~ 7.8, p1 ~ -0.45 and -0.66
            # functional walk back distance is *either* the minimum of the function value, or 5500 (standard value)

            # t0_corr = -7000+6000+2000 # no correction
            t0_corr = -7000+6000+2000 - np.amin([np.exp(7.8 - 0.45*lat[iH]),1000.])
            t0A_corr = -7000+6000+2000 - np.amin([np.exp(7.8 - 0.66*lat[iH]),1000.])

            latFC[iH] = pTrapInterp( np.amax([t0_SLE[iH] + t0_corr, 0.]) )
            latAFC[iH] = pTrapInterp( np.amax([t0_ALE[iH] + t0A_corr, 0.]) )


            # =========================================================

            # the genius multisite event tagger - plot 8

            # decide a threshold
            dIdx = np.argmax(data_filtDeriv)
            dMax = data_filtDeriv[dIdx]
            dRMS,_,_ = wl.baselineParameters(data_filtDeriv)
            # msThresh = np.amax([dMax * .2, dRMS * 5.])
            # msThresh = dMax * .15
            msThresh = 50.  # I don't know.  this seems like a good value

            # run peak detect algorithm
            maxtab,_ = wl.peakdet(data_filtDeriv, msThresh)

            # profit
            msList = []
            for iMax in range(len(maxtab)):
                idx = int(maxtab[iMax][0])
                val = maxtab[iMax][1]
                msList.append(dataTS[idx])
                # print("%d  idx %d  TS %d  val %.2f  thresh %.2f" % (iList, idx, dataTS[idx], val, msThresh))
            nMS[iH] = len(maxtab)

            # =========================================================
            # wfStd analysis
            wfAvgBL[iH] = dataBL
            wfRMSBL[iH] = dataNoise
            wfStd[iH] = np.std(data[5:-5])

            # ------------------------------------------------------------------------
            # End waveform processing.

            # Calculate error code
            fails[iH] = 0
            for i,j in enumerate(errorCode):
                if j==1: fails[iH] += int(j)<<i
            # print("fails:",fails[iH])

            # Make plots!
            if batMode: continue
            if plotNum==0: # raw data
                p0.cla()
                p0.plot(dataTS,data,'b')
                p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f" % (run,iList,chan,dataENFCal))
                p0.set_xlabel("Time (ns)", ha='right', x=1.)
                p0.set_ylabel("Voltage (ADC)", ha='right', y=1.)

            if plotNum==1: # wavelet plot
                p0.cla()
                p0.margins(x=0)
                p0.plot(dataTS,data_blSub,color='blue',label='data (%.2f keV)' % dataENFCal)
                p0.plot(dataTS,data_wlDenoised,color='cyan',label='denoised',alpha=0.7)
                p0.axvline(fitRiseTime50,color='green',label='fit 50%',linewidth=2)
                p0.plot(dataTS,fit_blSub,color='red',label='bestfit',linewidth=2)
                # p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f  flo %.0f  fhi %.0f  fhi-flo %.0f" % (run,iList,chan,dataENFCal,fitStartTime,fitMaxTime,fitMaxTime-fitStartTime))
                p0.legend(loc='best')
                p0.set_xlabel("Time (ns)", ha='right', x=1.)
                p0.set_ylabel("Voltage (ADC)", ha='right', y=1.)

                p1.cla()
                p1.imshow(wpCoeff, interpolation='nearest', aspect="auto", origin="lower",extent=[0, 1, 0, len(wpCoeff)],cmap='viridis')
                p1.axvline(float(wpLoRise)/numXRows,color='orange',linewidth=2)
                p1.axvline(float(wpHiRise)/numXRows,color='orange',linewidth=2)
                # p1.set_title("waveS5 %.2f  bcMax %.2f  bcMin %.2f  riseNoise %.2f" % (waveS5[iH], bcMax[iH], bcMin[iH], riseNoise[iH]))
                # p1.set_xlabel("Time (%wf)", ha='right', x=1.)
                p1.set_ylabel("WPT Coefficients", ha='right', y=1.)

            if plotNum==2: # time points, bandpass filters, tail slope
                p0.cla()
                p0.plot(dataTS,data,color='blue',label='data')
                p0.axvline(den10[iH],color='black',label='lpTP')
                p0.axvline(den50[iH],color='black')
                p0.axvline(den90[iH],color='black')
                p0.plot(dataTS,fit,color='magenta',label='bestfit')
                if errorCode[2]!=1: p0.plot(tailTS, wl.tailModelPol(tailTS, *popt1), color='orange',linewidth=2, label='tailPol')
                p0.legend(loc='best')
                p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f" % (run,iEvent,chan,dataENFCal))

                p1.cla()
                p1.plot(dataTS,data_lPass,color='blue',label='lowpass')
                p1.plot(dataTS,data_filtDeriv,color='green',label='filtDeriv')
                p1.plot(dataTS,data_filt,color='black',label='filtfilt')
                p1.plot(dataTS,data_bPass,color='red',label='bpass')
                p1.axvline(bandTime[iH],color='orange',label='bandTime')
                p1.legend(loc='best')

            if plotNum==3: # freq-domain matched filter
                p0.cla()
                p0.plot(dataTS,data,'b')
                p0.plot(dataTS,temp,'r')
                p0.plot(dataTS,fit,color='cyan')
                p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f" % (run,iEvent,chan,dataENFCal))

                data_asd, data_xFreq = plt.psd(data, Fs=1e8, NFFT=2048, pad_to=2048, visible=False)
                temp_asd, temp_xFreq = plt.psd(temp, Fs=1e8, NFFT=2048, pad_to=2048, visible=False)

                p1.cla()
                p1.loglog(data_xFreq, np.sqrt(data_asd), 'b')
                p1.loglog(noise_xFreq, np.sqrt(noise_asd), 'g')
                p1.loglog(temp_xFreq, np.sqrt(temp_asd), 'r')
                p1.set_xlabel('Frequency (Hz)')
                p1.set_ylabel('ASD')
                p1.grid('on')

                p2.cla()
                p2.plot(dataTS, SNR)
                p2.set_title('oppie %.1f' % (oppie[iH]))
                p2.set_xlabel('Offset time (s)')
                p2.set_ylabel('SNR')

            if plotNum==4: # time domain match filter plot
                p0.cla()
                p0.plot(dataTS,data_blSub,color='blue',label='data',alpha=0.7)
                p0.plot(dataTS,fit_blSub,color='red',label='bestfit',linewidth=3)
                p0.axvline(matchTime[iH],color='orange',label='matchTime',linewidth=2)
                p0.plot(matchTS,smoothMF,color='magenta',label='smoothMF',linewidth=3)
                p0.plot(matchTS,match,color='cyan',label='match',linewidth=3)
                p0.set_xlabel('Time (s)')
                p0.set_ylabel('Voltage (arb)')
                p0.legend(loc='best')
                p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f  matchMax %.2f  matchTime %.2f  matchWidth %.2f" % (run,iEvent,chan,dataENFCal,matchMax[iH],matchTime[iH],matchWidth[iH]))

            if plotNum==5: # bandTime plot
                p0.cla()
                p0.plot(dataTS,data_blSub,color='blue',label='data',alpha=0.7)
                p0.plot(dataTS,data_lPass,color='magenta',label='lowpass',linewidth=4)
                p0.plot(dataTS,data_bPass,color='red',label='bpass',linewidth=4)
                p0.axvline(bandTime[iH],color='orange',label='bandTime',linewidth=4)
                p0.legend(loc='best')
                p0.set_xlabel('Time (ns)')
                p0.set_ylabel('ADC (arb)')
                p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f" % (run,iEvent,chan,dataENFCal))

            if plotNum==6: # waveform fit plot
                p0.cla()
                p0.plot(dataTS,data,color='blue',label='data')
                # p0.plot(dataTS,data_wlDenoised,color='cyan',label='wlDenoised',alpha=0.5)
                p0.plot(dataTS,temp,color='orange',label='xgauss guess')
                p0.plot(dataTS,fit,color='red',label='xgauss fit')
                p0.set_title("Run %d  evt %d  chan %d  trapENFCal %.1f  trapENM %.1f  deltaBL %.1f\n  amp %.2f  mu %.2f  sig %.2f  tau %.2f  chi2 %.2f  spd %.3f" % (run,iList,chan,dataENFCal,dataENM,dataBL-bl,amp,mu,sig,tau,fitChi2[iH],fitSpeed))
                p0.legend(loc='best')
                p1.cla()
                p1.plot(dataTS,data-fit,color='blue',label='residual')
                p1.legend(loc='best')
                p2.cla()
                p2.plot(ampTr[1:],label='amp',color='red')
                p2.legend(loc='best')
                p3.cla()
                p3.plot(muTr[1:],label='mu',color='green')
                p3.legend(loc='best')
                p4.cla()
                p4.plot(sigTr[1:],label='sig',color='blue')
                p4.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
                p4.legend(loc='best')
                p5.cla()
                p5.plot(tauTr[1:],label='tau',color='black')
                p5.legend(loc='best')
                p6.cla()
                p6.plot(blTr[1:],label='bl',color='magenta')
                p6.legend(loc='best')

                print(gatTree.fitSlo.at(iH), sig)

            if plotNum==7: # new traps plot
                p0.cla()
                p0.plot(dataTS, data_blSub, color='blue', label='data')
                p0.plot(sTrapTS, sTrap, color='red', label='sTrap')
                p0.axvline(t0_SLE[iH], color='red')
                p0.plot(aTrapTS, aTrap, color='orange', label='aTrap')
                p0.axvline(t0_ALE[iH], color='orange')
                p0.plot(eTrapTS, eTrap, color='green', label='eTrap')
                p0.axhline(lat[iH],color='green')
                p0.plot(pTrapTS, pTrap, color='magenta', label='pTrap')
                p0.axhline(latAFC[iH], color='magenta')
                p0.axhline(latE50[iH], color='cyan')
                p0.set_title("trapENFCal %.2f  trapENM %.2f || latEM %.2f  latEF %.2f  latEAF %.2f  latEFC %.2f  latEAFC %.2f  latE50 %.2f" % (dataENFCal,dataENM,lat[iH],latF[iH],latAF[iH],latFC[iH],latAFC[iH], latE50[iH]))
                p0.legend(loc='best')

            if plotNum==8: # multisite tag plot
                p0.cla()
                p0.plot(dataTS, data_blSub, color='blue', label='data')
                p0.plot(dataTS, data_filtDeriv, color='red', label='filtDeriv')
                for mse in msList: p0.axvline(mse, color='green')
                p0.axhline(msThresh,color='red')
                p0.legend()

            if plotNum==9: # wavelet vs wf fit residual plot

                # wavelet packet transform on wf fit residual
                fitResid = data-fit
                wpRes = pywt.WaveletPacket(fitResid, 'db2', 'symmetric', maxlevel=4)
                nodesRes = wpRes.get_level(4, order='freq')
                wpCoeffRes = np.array([n.data for n in nodesRes], 'd')
                wpCoeffRes = abs(wpCoeffRes)
                R6 = np.sum(wpCoeffRes[2:9,1:wpLength//4+1])
                R7 = np.sum(wpCoeffRes[2:9,wpLength//4+1:wpLength//2+1])
                R8 = np.sum(wpCoeffRes[2:9,wpLength//2+1:3*wpLength//4+1])
                R9 = np.sum(wpCoeffRes[2:9,3*wpLength//4+1:-1])
                R10 = np.sum(wpCoeffRes[9:,1:wpLength//4+1])
                R11 = np.sum(wpCoeffRes[9:,wpLength//4+1:wpLength//2+1])
                R12 = np.sum(wpCoeffRes[9:,wpLength//2+1:3*wpLength//4+1])
                R13 = np.sum(wpCoeffRes[9:,3*wpLength//4+1:-1])
                RsumList = [R6, R7, R8, R9, R10, R11, R12, R13]
                bcMinRes = 1. if np.min(RsumList) < 1 else np.min(RsumList)
                riseNoiseRes = np.sum(wpCoeffRes[2:-1,wpLoRise:wpHiRise]) / bcMinRes
                rnCut = 1.1762 + 0.00116 * np.log(1 + np.exp((dataENFCal-7.312)/0.341))

                p0.cla()
                p0.margins(x=0)
                p0.plot(dataTS,data_blSub,color='blue',label='data')
                # p0.plot(dataTS,data_wlDenoised,color='cyan',label='denoised',alpha=0.7)
                # p0.axvline(fitRiseTime50,color='green',label='fit 50%',linewidth=2)
                p0.plot(dataTS,fit_blSub,color='red',label='bestfit',linewidth=2)
                # p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f  flo %.0f  fhi %.0f  fhi-flo %.0f" % (run,iList,chan,dataENFCal,fitStartTime,fitMaxTime,fitMaxTime-fitStartTime))
                # p0.legend(loc='best')
                p0.set_title("Run %d  Entry %d  Channel %d  ENFCal %.2f  flo %.0f  fhi %.0f  fhi-flo %.0f  approxFitE %.2f" % (run,iList,chan,dataENFCal,fitStartTime,fitMaxTime,fitMaxTime-fitStartTime,fitAmp[iH]*0.4))

                p1.cla()
                p1.plot(dataTS,fitResid,color='blue')

                p2.cla()
                p2.set_title("riseNoise %.2f  rnCut %.2f  riseNoiseRes %.2f  bcMinRes %.2f  bcMin %.2f  max %.2f" % (riseNoise[iH],rnCut,riseNoiseRes,bcMinRes,bcMin[iH],wpCoeffRes.max()))
                p2.imshow(wpCoeffRes, interpolation='nearest', aspect="auto", origin="lower",extent=[0, 1, 0, len(wpCoeff)],cmap='viridis')

            plt.tight_layout()
            plt.pause(0.000001)
            # ------------------------------------------------------------------------

        # End loop over hits, fill branches
        if batMode:
            for key in brDict:
                brDict[key][1].Fill()
            if iList%5000 == 0 and iList!=0:
                out.Write("",TObject.kOverwrite)
                print("%d / %d entries saved (%.2f %% done), time: %s" % (iList,nList,100*(float(iList)/nList),time.strftime('%X %x %Z')))

    # End loop over events
    if batMode and not intMode:
        out.Write("",TObject.kOverwrite)
        print("Wrote",out.GetBranch("channel").GetEntries(),"entries in the copied tree,")
        print("and wrote",b1.GetEntries(),"entries in the new branches.")

    stopT = time.clock()
    print("Stopped:",time.strftime('%X %x %Z'),"\nProcess time (min):",(stopT - startT)/60)
    print(float(nList)/((stopT-startT)/60.),"entries per minute.")
Beispiel #15
0
def main(argv):
    # gROOT.ProcessLine("gErrorIgnoreLevel = 3001;") # suppress ROOT error messages
    startT = time.clock()
    print "Started:", time.strftime('%X %x %Z')
    intMode, batMode, rangeMode, fileMode, dsNum, subNum, runNum = False, False, False, False, -1, -1, -1
    for i, opt in enumerate(argv):
        if opt == "-i":
            intMode = True
            print "Interactive mode selected. Use \"p\" for previous and \"q\" to exit."
        if opt == "-r":
            rangeMode = True
            dsNum, subNum = int(argv[i + 1]), int(argv[i + 2])
            print "Scanning DS-%d range %d" % (dsNum, subNum)
        if opt == "-f":
            fileMode = True
            dsNum, runNum = int(argv[i + 1]), int(argv[i + 2])
        if opt == "-b":
            batMode = True
            import matplotlib
            if os.environ.get('DISPLAY', '') == '':
                print('No display found. Using non-interactive Agg backend')
                matplotlib.use('Agg')
            print "Batch mode selected.  A new file will be created."
    import matplotlib.pyplot as plt

    # Set file I/O and cuts
    inFile = "./data/waveSkimDS4_test.root"
    outFile = "./data/waveletSkimDS4_test.root"
    if (rangeMode):
        inFile = "~/project/v2-waveskim/waveSkimDS%d_%d.root" % (dsNum, subNum)
        outFile = "~/project/v2-processwfs/waveletSkimDS%d_%d.root" % (dsNum,
                                                                       subNum)
    if (fileMode):
        # inFile = "~/project/cal-wave-skim/waveSkimDS%d_run%d.root" % (dsNum,runNum)
        # outFile = "~/project/cal-wavelet-skim/waveletSkimDS%d_run%d.root" % (dsNum,runNum)
        inFile = "./waveSkimDS%d_run%d.root" % (dsNum, runNum)
        outFile = "./waveletSkimDS%d_run%d.root" % (dsNum, runNum)

    inputFile = TFile(inFile)
    waveTree = inputFile.Get("skimTree")
    print "Found", waveTree.GetEntries(), "input entries."

    theCut = inputFile.Get("theCut").GetTitle()
    # theCut = "trapENFCal > 0.8 && gain==0 && mHClean==1 && isGood && !muVeto && !isLNFill1 && !isLNFill2 && P!=0 && D!=0 && trapETailMin<0"
    # theCut += " && Entry$ < 2000"
    # theCut += " && trapENFCal > 20"
    print "Using cut:\n", theCut, "\n"

    print "Attempting entrylist draw ..."
    waveTree.Draw(">>elist", theCut, "GOFF entrylist")
    elist = gDirectory.Get("elist")
    waveTree.SetEntryList(elist)
    nList = elist.GetN()
    print "Found", nList, "entries passing cuts."

    # In batch mode ONLY, create an output file+tree & append new branches
    outputFile = TFile()
    outTree = TTree()
    if (batMode == True):
        outputFile = TFile(outFile, "RECREATE")
        print "Attempting tree copy to", outFile
        outTree = waveTree.CopyTree("")
        outTree.Write()
        print "Wrote", outTree.GetEntries(), "entries."
        cutUsed = TNamed("cutUsedHere", theCut)
        cutUsed.Write()

    waveS0 = std.vector("double")()
    waveS1 = std.vector("double")()
    waveS2 = std.vector("double")()
    waveS3 = std.vector("double")()
    waveS4 = std.vector("double")()
    waveS5 = std.vector("double")()
    wpar4 = std.vector("double")()
    waveEnt = std.vector("double")()
    butterMax = std.vector("double")()
    butterTime = std.vector("double")()
    tOffset = std.vector("double")()
    lastZeroTime = std.vector("double")()
    pol0 = std.vector("double")()
    pol1 = std.vector("double")()
    pol2 = std.vector("double")()
    pol3 = std.vector("double")()
    exp0 = std.vector("double")()
    exp1 = std.vector("double")()
    rt10 = std.vector("double")()
    rt20 = std.vector("double")()
    rt50 = std.vector("double")()
    rt90 = std.vector("double")()

    bWaveS0 = outTree.Branch("waveS0", waveS0)
    bWaveS1 = outTree.Branch("waveS1", waveS1)
    bWaveS2 = outTree.Branch("waveS2", waveS2)
    bWaveS3 = outTree.Branch("waveS3", waveS3)
    bWaveS4 = outTree.Branch("waveS4", waveS4)
    bWaveS5 = outTree.Branch("waveS5", waveS5)
    bWPar4 = outTree.Branch("wpar4", wpar4)
    bWaveEnt = outTree.Branch("waveEnt", waveEnt)
    bButterMax = outTree.Branch("butterMax", butterMax)
    bButterTime = outTree.Branch("butterTime", butterTime)
    bTOffset = outTree.Branch("tOffset", tOffset)
    bLastZeroTime = outTree.Branch("lastZeroTime", lastZeroTime)
    bPol0 = outTree.Branch("pol0", pol0)
    bPol1 = outTree.Branch("pol1", pol1)
    bPol2 = outTree.Branch("pol2", pol2)
    bPol3 = outTree.Branch("pol3", pol3)
    bExp0 = outTree.Branch("exp0", exp0)
    bExp1 = outTree.Branch("exp1", exp1)
    bRt10 = outTree.Branch("rt10", rt10)
    bRt20 = outTree.Branch("rt20", rt20)
    bRt50 = outTree.Branch("rt50", rt50)
    bRt90 = outTree.Branch("rt90", rt90)

    # Make a figure
    fig = plt.figure(figsize=(7, 7), facecolor='w')
    p1 = plt.subplot(211)
    p2 = plt.subplot(212)
    if not batMode: plt.show(block=False)

    # Begin loop over events
    iList = -1
    while True:
        iList += 1
        if intMode == True and iList != 0:
            value = raw_input()
            if value == 'q': break  # quit
            if value == 'p': iList -= 2  # go to previous
            if (value.isdigit()):
                iList = int(value)  # go to entry number
        elif intMode == False and batMode == False:
            plt.pause(0.001)  # rapid-draw mode
        if iList >= nList: break  # bail out, goose!

        entry = waveTree.GetEntryNumber(iList)
        waveTree.LoadTree(entry)
        waveTree.GetEntry(entry)
        nChans = waveTree.channel.size()
        waveS1.assign(nChans, -999)
        waveS0.assign(nChans, -999)
        waveS1.assign(nChans, -999)
        waveS2.assign(nChans, -999)
        waveS3.assign(nChans, -999)
        waveS4.assign(nChans, -999)
        waveS5.assign(nChans, -999)
        wpar4.assign(nChans, -999)
        waveEnt.assign(nChans, -999)
        butterMax.assign(nChans, -999)
        butterTime.assign(nChans, -999)
        tOffset.assign(nChans, -999)
        lastZeroTime.assign(nChans, -999)
        pol0.assign(nChans, -999)
        pol1.assign(nChans, -999)
        pol2.assign(nChans, -999)
        pol3.assign(nChans, -999)
        exp0.assign(nChans, -999)
        exp1.assign(nChans, -999)
        rt10.assign(nChans, -999)
        rt20.assign(nChans, -999)
        rt50.assign(nChans, -999)
        rt90.assign(nChans, -999)

        # Loop over hits passing cuts
        numPass = waveTree.Draw("channel", theCut, "GOFF", 1, iList)
        chans = waveTree.GetV1()
        chanList = list(set(int(chans[n]) for n in xrange(numPass)))
        hitList = (iH for iH in xrange(nChans)
                   if waveTree.channel.at(iH) in chanList
                   )  # a 'generator expression'
        for iH in hitList:
            run = waveTree.run
            iEvent = waveTree.iEvent
            chan = waveTree.channel.at(iH)
            energy = waveTree.trapENFCal.at(iH)
            wf = waveTree.MGTWaveforms.at(iH)
            startTime = waveTree.triggerTrapt0.at(iH)
            blrwfFMR50 = waveTree.blrwfFMR50.at(iH)

            if wf.GetID() != chan:
                print "Warning!!  Vector matching failed!  iList %d  run %d  iEvent %d" % (
                    iList, run, iEvent)
                break

            signal = wl.processWaveform(wf, opt='full')
            waveBLSub = signal.GetWaveBLSub()
            waveFilt = signal.GetWaveFilt()
            waveTS = signal.GetTS()
            waveletYOrig, waveletYTrans = signal.GetWavelet()

            # waveFTX, waveFTY, waveFTPow = signal.GetFourier()
            # waveTrap = signal.GetTrapezoid()
            # waveTrapF = signal.GetFiltTrapezoid()
            # waveTrapX = np.linspace(0, len(waveTrap), num=len(waveTrap))
            # _,waveletFilt = wl.waveletTransform(waveFilt,level=3)  # testing other levels on the filtered WF

            # Wavelet cut parameters
            waveS0[iH] = np.sum(waveletYTrans[0:1, 1:-1])
            waveS1[iH] = np.sum(waveletYTrans[0:1, 1:33])
            waveS2[iH] = np.sum(waveletYTrans[0:1, 33:65])
            waveS3[iH] = np.sum(waveletYTrans[0:1, 65:97])
            waveS4[iH] = np.sum(waveletYTrans[0:1, 97:-1])
            waveS5[iH] = np.sum(waveletYTrans[2:-1, 1:-1])
            wpar4[iH] = np.amax(waveletYTrans[0:1, 1:-1])

            # Waveform entropy parameters
            d1 = 2. * np.multiply(
                waveletYTrans[0:1, 1:65],
                np.log(waveletYTrans[0:1, 1:65] / waveS0[iH] / 2.0))
            d2 = np.multiply(waveletYTrans[0:1, 65:-1],
                             np.log(waveletYTrans[0:1, 65:-1] / waveS0[iH]))
            waveEnt[iH] = np.abs(np.sum(d1)) + np.abs(np.sum(d2))

            # Smoothed derivative params
            waveFiltDeriv = wl.wfDerivative(waveFilt)
            butterMax[iH] = np.amax(waveFiltDeriv)
            butterTime[iH] = np.argmax(
                waveFiltDeriv[100:]) * 10 + signal.GetOffset() + 1000
            tOffset[iH] = signal.GetOffset()
            # print "max %.3f  ts %.0f ns  offset %.0f ns" % (butterMax[iH], butterTime[iH], tOffset[iH])

            # Make a super denoised waveform
            wp = pywt.WaveletPacket(data=waveBLSub,
                                    wavelet='haar',
                                    mode='symmetric',
                                    maxlevel=3)
            new_wp = pywt.WaveletPacket(data=None,
                                        wavelet='haar',
                                        mode='symmetric')
            new_wp['aaa'] = wp['aaa'].data
            superDenoisedWF = new_wp.reconstruct(update=False)
            mgtSDWF = wl.MGTWFFromNpArray(superDenoisedWF)

            # Try to get start time by finding the super denoised last zero crossing
            lastZero = 0
            zeros = np.asarray(np.where(superDenoisedWF < 0.1))
            if (zeros.size > 0): lastZero = zeros[0, -1]
            lastZeroTime[iH] = waveTS[lastZero]

            # Time point calculator
            timePointCalc = MGWFTimePointCalculator()
            timePointCalc.AddPoint(.1)
            timePointCalc.AddPoint(.2)
            timePointCalc.AddPoint(.5)
            timePointCalc.AddPoint(.9)
            timePointCalc.FindTimePoints(mgtSDWF)
            rt10[iH] = timePointCalc.GetFromStartRiseTime(0) * 10
            rt20[iH] = timePointCalc.GetFromStartRiseTime(1) * 10
            rt50[iH] = timePointCalc.GetFromStartRiseTime(2) * 10
            rt90[iH] = timePointCalc.GetFromStartRiseTime(3) * 10
            # print "lastZero %.2f  startTime %.2f  10 %.2f  20 %.2f  50 %.2f  90 %.2f" % (lastZeroTime[iH], startTime,rt10, rt20, rt50, rt90)

            # Fit tail slope (2 methods).  Guard against fit fails
            idxMax = np.where(waveTS > rt90[iH])  # returns an array/tuple
            idxMax = idxMax[0][0]  # "cast" to int
            tail, tailTS = waveBLSub[idxMax:], waveTS[idxMax:]
            try:
                popt1, _ = curve_fit(wl.tailModelPol, tailTS, tail)
                pol0[iH], pol1[iH], pol2[iH], pol3[iH] = popt1[0], popt1[
                    1], popt1[2], popt1[3]
            except:
                pass
            try:
                popt2, _ = curve_fit(wl.tailModelExp,
                                     tailTS,
                                     tail,
                                     p0=[energy, 72000])
                exp0[iH], exp1[iH] = popt2[0], popt2[1]
            except:
                # print "curve_fit tailModelExp failed, run %i  event %i  channel %i" % (run, iList, chan)
                pass

            # Make a plot
            if not batMode:
                print "i %d  run %d  iEvent(i) %d  iH(j+1) %d/%d  chan %d  energy %.2f  bt %.2f" % (
                    iList, run, iEvent, iH + 1, nChans, chan, energy,
                    butterTime[iH])

                # Waveform plot
                p1.cla()
                p1.plot(waveTS, waveBLSub, color='blue')
                p1.plot(waveTS, superDenoisedWF, color='red')
                p1.plot(tailTS, wl.tailModelPol(tailTS, *popt1), color='cyan')
                p1.plot(tailTS,
                        wl.tailModelExp(tailTS, *popt2),
                        color='orange')
                p1.set_xlabel("Time (ns)")
                p1.set_ylabel("ADC")
                p1.set_title(
                    "Run %d  Ch %d  Energy %.2f  \n S5 %.2f  (S3-S2)/E %.2f  (S3-S4)/E %.2f"
                    % (run, chan, energy, waveS5[iH],
                       ((waveS3[iH] - waveS2[iH]) / energy),
                       ((waveS3[iH] - waveS4[iH]) / energy)))

                # Wavelet plot
                p2.cla()
                p2.imshow(waveletYTrans,
                          interpolation='nearest',
                          aspect="auto",
                          origin="lower",
                          extent=[0, 1, 0, len(waveletYTrans)],
                          cmap='jet')

                plt.tight_layout()
                plt.subplots_adjust(hspace=.2, top=0.92)
                plt.pause(0.00001)

        # End loop over hits
        if (batMode == True):
            bWaveS0.Fill()
            bWaveS1.Fill()
            bWaveS2.Fill()
            bWaveS3.Fill()
            bWaveS4.Fill()
            bWaveS5.Fill()
            bWPar4.Fill()
            bWaveEnt.Fill()
            bButterMax.Fill()
            bButterTime.Fill()
            bTOffset.Fill()
            bLastZeroTime.Fill()
            bPol0.Fill()
            bPol1.Fill()
            bPol2.Fill()
            bPol3.Fill()
            bExp0.Fill()
            bExp1.Fill()
            bRt10.Fill()
            bRt20.Fill()
            bRt50.Fill()
            bRt90.Fill()
            if iList % 5000 == 0:
                outTree.Write("", TObject.kOverwrite)
                print "%d / %d entries saved (%.2f %% done)." % (
                    iList, nList, 100 * (float(iList) / nList))

    # End loop over events
    if (batMode == True):
        outTree.Write("", TObject.kOverwrite)
        print "Wrote", outTree.GetBranch(
            "channel").GetEntries(), "entries in the copied tree,"
        print "and wrote", bWaveS0.GetEntries(), "entries in the new branches."

    stopT = time.clock()
    print "Process time (min): ", (stopT - startT) / 60
Beispiel #16
0
def main():


    EffName = "Eff"
    WMMName = "MM"

    # Bin Definitions
    binHeader = "bindefs"
    # binname = fluxmodel
    binname = "bin0"

    # ROOT Output
    OutFile = "response_matrix.root"
    # fout = ROOT.TFile(OutFile, "RECREATE")
    fout = ROOT.TFile(OutFile, "UPDATE")
    # Check if bin directory exists, quit if so, otherwise create it!
    if ( not fout.GetDirectory(binname) ):
        pdir = fout.mkdir(binname,"Bin number 0")
    else:
        fout.Close()
        print("\n=========================\n")
        errormessage = "Directory %s already exists!\nEither try another bin number or delete %s and start again. Exiting...\n"%(binname,OutFile)
        raise ValueError(errormessage)

    # Go to home of ROOT file
    fout.cd(binname)

    # formatted_df = pd.read_csv('../formatted-dataframe.csv')
    formatted_df_outfile = os.path.join('/data/user/jbourbeau/composition/unfolding',
                                        'unfolding-dataframe-PyUnfold-formatted.csv')
    df_flux = pd.read_csv(formatted_df_outfile, index_col='log_energy_bin_idx')
    counts = df_flux['counts'].values

    res_mat_file = os.path.join('/data/user/jbourbeau/composition/unfolding',
                                'response.txt')
    # res_mat_file = os.path.join('/data/user/jbourbeau/composition/unfolding',
    #                             'block_response.txt')
    block_response = np.loadtxt(res_mat_file)
    res_mat_err_file = os.path.join('/data/user/jbourbeau/composition/unfolding',
                                    'response_err.txt')
    # res_mat_err_file = os.path.join('/data/user/jbourbeau/composition/unfolding',
    #                                 'block_response_err.txt')
    block_response_err = np.loadtxt(res_mat_err_file)

    cbins = len(counts)+1
    carray = np.arange(cbins, dtype=float)
    print('carray = {}'.format(carray))
    print('cbins = {}'.format(cbins))

    ebins = len(counts)+1
    earray = np.arange(ebins, dtype=float)
    print('earray = {}'.format(earray))
    cbins -= 1
    ebins -= 1

    # Prepare Combined Weighted Histograms - To be Normalized by Model After Filling
    # Isotropic Weights of Causes - For Calculating Combined Species Efficiency
    Eff = TH1F("%s"%(EffName), 'Non-Normed Combined Efficiency', cbins, carray)
    Eff.GetXaxis().SetTitle('Causes')
    Eff.GetYaxis().SetTitle("Efficiency")
    Eff.SetStats(0)
    Eff.Sumw2()
    # Isotropic Weighted Mixing Matrix - For Calculating Combined Species MM
    WMM = TH2F("%s"%(WMMName), 'Weighted Combined Mixing Matrix',
               cbins, carray, ebins, earray)
    WMM.GetXaxis().SetTitle('Causes')
    WMM.GetYaxis().SetTitle('Effects')
    WMM.SetStats(0)
    WMM.Sumw2()
    # # Isotropic Weighted Mixing Matrix - For Calculating Combined Species MM
    # ModelMM = TH2F("Model_%s"%(WMMName),'Model Weighted Combined Matrix',cbins,carray,ebins,earray)
    # ModelMM.GetXaxis().SetTitle('Causes')
    # ModelMM.GetYaxis().SetTitle('Effects')
    # ModelMM.SetStats(0)
    # ModelMM.Sumw2()
    # Raw Number of Events in Each Bin
    NMM = TH2F("NMM", 'Number of Events Matrix', cbins, carray, ebins, earray)
    NMM.GetXaxis().SetTitle('Causes')
    NMM.GetYaxis().SetTitle('Effects')
    NMM.SetStats(0)
    NMM.Sumw2()


    throwArea = 1.
    for ci in xrange(0,cbins):
        # Calculate Flux-Weighted Efficiency
        # Eval = Eff.GetBinContent(ci+1)/binwidth[ci]/throwArea
        # dEval = Eff.GetBinError(ci+1)/binwidth[ci]/throwArea
        dEval, Eval = 1, 1
        Eff.SetBinContent(ci+1, Eval)
        # Eff.SetBinError(ci+1, dEval)

        # Normalize Species Probability Matrix
        sum2 = 0
        for ek in xrange(0,ebins):
            # Calculate Flux-Weighted Mixing Matrix
            # wmm_val = WMM.GetBinContent(ci+1,ek+1)/binwidth[ci]/throwArea
            # dwmm_val = WMM.GetBinError(ci+1,ek+1)/binwidth[ci]/throwArea

            # wmm_val, dwmm_val =  block_response[ebins-ek-1][ci], block_response_err[ebins-ek-1][ci]
            wmm_val, dwmm_val =  block_response[ek][ci], block_response_err[ek][ci]

            sum2 += dwmm_val**2
            WMM.SetBinContent(ci+1, ek+1, wmm_val)
            WMM.SetBinError(ci+1, ek+1, dwmm_val)

        Eff.SetBinError(ci+1, np.sqrt(sum2))

    # Write model name to file
    modelName = 'whatever'
    MODELNAME = TNamed("ModelName",modelName)
    MODELNAME.Write()
    # Write Cuts to file
    cuts = ''
    cuts_theta = cuts.replace("rec.zenithAngle","theta")
    CUTS = TNamed("cuts",cuts_theta)
    CUTS.Write()
    # Write the cause and effect arrays to file
    CARRAY = TH1F("CARRAY","Cause Array", cbins, carray)
    CARRAY.GetXaxis().SetTitle('Causes')
    EARRAY = TH1F("EARRAY","Effect Array", ebins, earray)
    EARRAY.GetXaxis().SetTitle('Effects')
    CARRAY.Write()
    EARRAY.Write()

    # Write the weighted histograms to file
    addtitle = r'Test'
    WMM.SetTitle(WMM.GetTitle()+addtitle)
    # ModelMM.SetTitle(ModelMM.GetTitle()+addtitle)
    Eff.SetTitle(Eff.GetTitle()+addtitle)
    WMM.Write()
    Eff.Write()
    # ModelMM.Write()
    NMM.Write()

    fout.Write()
    fout.Close()
    print("Saving output file %s\n"%OutFile)

    print("\n=========================\n")
    print("Finished here! Exiting...")
Beispiel #17
0
    def plotDiscriminant(self, discriminant, signal_idx, bkg_idx, weights = None, save_disc = True, rejection_power=True):
        '''
        Plot the discriminants and the resulting ROC curve derived from them.

        Keyword args:
        discriminant --- The score of the BDT (set in the setProbas method)
        signal_idx --- The true indices of all signal events
        bkg_idx ---The true indices of all background events
        save_disc --- Flag indicating if the discriminant plots should be saved.
        rejection_power --- Whether or not to calculate bkg power: 1/eff in addtion to 1-eff
        '''
        import ROOT as root
        from ROOT import TH2D, TCanvas, TFile, TNamed, TH1F, TLegend
        import numpy as np
        from root_numpy import fill_hist
        import functions as fn
        import os
        

        # stop showing plots to screen
        root.gROOT.SetBatch(True)

        if not os.path.exists(self.output_path):
            os.makedirs(self.output_path)
        fo = TFile.Open(self.output_path+"/"+self.output_prefix+str(self.job_id)+'.root','RECREATE')

        bins = 100
        # when creating the plots do it over the range of all probas (scores)
        discriminant_bins = np.linspace(np.min(discriminant), np.max(discriminant), bins)

        hist_bkg = TH1F("Background Discriminant","Discriminant",bins, np.min(discriminant), np.max(discriminant))
        hist_sig = TH1F("Signal Discriminant","Discriminant",bins, np.min(discriminant), np.max(discriminant))

        # fill the signal and background histograms
        if weights is not None:
            fill_hist(hist_bkg,discriminant[bkg_idx], weights[bkg_idx])
            fill_hist(hist_sig,discriminant[signal_idx], weights[signal_idx])
        else:
            fill_hist(hist_bkg,discriminant[bkg_idx])
            fill_hist(hist_sig,discriminant[signal_idx])
        if hist_bkg.Integral() != 0:
            hist_bkg.Scale(1/hist_bkg.Integral())
        if hist_sig.Integral() != 0:
            hist_sig.Scale(1/hist_sig.Integral())

        hist_sig.SetLineColor(4)
        hist_bkg.SetLineColor(2)
        #hist_sig.SetFillColorAlpha(4, 0.5);
        hist_sig.SetFillStyle(3004)
        #hist_bkg.SetFillColorAlpha(2, 0.5);
        hist_bkg.SetFillStyle(3005)
        hist_sig.Write()
        hist_bkg.Write()
        c = TCanvas()
        leg = TLegend(0.8,0.55,0.9,0.65);leg.SetFillColor(root.kWhite)
        leg.AddEntry(hist_sig, "Signal","l")
        leg.AddEntry(hist_bkg, "Background", "l")
        max_y = max(hist_sig.GetMaximum(), hist_bkg.GetMaximum())
        hist_sig.SetMaximum(max_y*1.2)
        hist_bkg.SetMaximum(max_y*1.2)

        hist_sig.GetXaxis().SetTitle("Signal Probability")
        hist_sig.GetYaxis().SetTitle("Normalised Entries")

        hist_bkg.GetXaxis().SetTitle("Signal Probability")
        hist_bkg.GetYaxis().SetTitle("Normalised Entries")

        hist_sig.Draw('hist')
        hist_bkg.Draw('histsame')
        c.Write()
        if save_disc == True:
            if not os.path.exists('disc_plots'):
                os.makedirs('disc_plots')
            c.SaveAs('disc_plots/discriminants_'+str(self.job_id)+'.png')

        # before deciding whether to do a left or right cut for the roc curve we have to find the median.
        sig_median = np.median(discriminant[signal_idx])
        bkg_median = np.median(discriminant[bkg_idx])
        if sig_median > bkg_median:
            roc_cut = 'R'
        else:
            roc_cut = 'L'

        # create the single sided roccurve with the code from Sam
        self.roc_graph = fn.RocCurve_SingleSided(hist_sig, hist_bkg, self.sig_eff,self.bkg_eff, roc_cut)
        self.roc_graph.SetName('BackgroundRejection')
        self.roc_graph.SetTitle('BackgroundRejection')
        self.roc_graph.Write()
        
        # get teh background rejection power at 50% signal efficiency
        # store the efficiencies first
        self.ROC_sig_efficiency, self.ROC_bkg_rejection = fn.getEfficiencies(self.roc_graph)
        self.bkgRejectionPower()
        # write the roc score as a string to the output file
        rej_string = 'rejection_power_'+str(self.ROC_rej_power_05)
        rej_n = TNamed(rej_string,rej_string)
        rej_n.Write()


        if rejection_power:
            c.SetLogy()
            self.roc_graph_power = fn.RocCurve_SingleSided(hist_sig, hist_bkg, self.sig_eff,self.bkg_eff, roc_cut, rejection=False)
            c.cd()
            self.roc_graph_power.SetName('BackgroundPower')
            self.roc_graph_power.SetTitle('BackgroundPower')
            self.roc_graph_power.Write()

        # write the decision function to the root file as well, if it is defined.
        if len(self.decision_function) > 0:
            self.decisionFunctionCanvas()
            # add the legends
            leg2 = TLegend(0.8,0.55,0.9,0.65);leg2.SetFillColor(root.kWhite)
            leg2.AddEntry(self.df_sig, "Signal","l")
            leg2.AddEntry(self.df_bkg, "Background", "l")
            # canvas to draw them on
            c2 = TCanvas('Decision Functions')
            self.df_sig.Draw('hist')
            self.df_bkg.Draw('histsame')
            leg2.Draw('same')
            c2.Write()
            # now write the df histograms as well
            self.df_sig.Write()
            self.df_bkg.Write()
            
        self.hist_sig = hist_sig.Clone(); self.hist_sig.SetDirectory(0)
        self.hist_bkg = hist_bkg.Clone(); self.hist_bkg.SetDirectory(0)
        
        fo.Close()
Beispiel #18
0
def main(argv):
    """ Interactive-fit or 'rapid'-fit waveforms that pass a given TCut.
        BUG: Doesn't always work with a TChain.  Add input files together
        with hadd and use a single TFile.
    """
    scanSpeed = 0.2
    iList = -1
    opt1 = ""
    intMode, batMode = False, False
    if (len(argv) >= 1): opt1 = argv[0]
    if "-i" in (opt1):
        intMode = True
        print "Interactive mode selected."
    if "-b" in (opt1):
        batMode = True
        print "Batch mode selected.  A new file will be created."

    # Load template waveform (created by gen-template.py)
    npzfile = np.load("./data/genTemplateWF.npz")
    temp, tempTS, tempE, tempST = npzfile['arr_0'] + 1, npzfile[
        'arr_1'], npzfile['arr_2'], npzfile['arr_3'] * 10

    # Set cuts
    # theCut = inputFile.Get("cutUsedHere").GetTitle()
    # DS3 "big DC" cut - PRELIMINARY
    theCut = "trapENFCal > 0.8 && gain==0 && mHClean==1 && isGood && !muVeto && !wfDCBits && !isLNFill1 && !isLNFill2 && trapETailMin < 0 && channel!=596 && channel!=676 && channel!=676 && channel!=612 && channel!=1104 && channel!=1200 && channel!=1334 && channel!=1336 && tOffset < 10 && waveS5/TMath::Power(trapENFCal,1/4) < 1200 && (waveS3-waveS2)/trapENFCal > 100 && (waveS3-waveS2)/trapENFCal < 300 && !(channel==692 && (run==16974 || run==16975 || run==16976 || run==16977 || run==16978 || run==16979)) && butterTime < 11000"
    # theCut += " && trapENFCal > 1.5 && trapENFCal < 2.1"
    # theCut += " && trapENFCal < 20 && trapENFCal > 2 && run > 17480"
    # theCut += " && kvorrT/trapENFCal > 2.2 && trapENFCal > 2 && trapENFCal < 10"

    # Set file I/O and create entry lists
    startT = time.clock()
    inFile = "~/project/wavelet-skim/waveletSkimDS3_1.root"
    # inFile = "~/project/wavelet-skim/hadd/waveletSkimDS3.root"
    outFile = "~/project/fit-skim/fitSkimDS3_1.root"
    inputFile = TFile(inFile)
    waveTree = inputFile.Get("skimTree")
    print "Found", waveTree.GetEntries(
    ), "input entries.  Using cut:\n", theCut, "\n"
    waveTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    waveTree.SetEntryList(elist)
    nList = elist.GetN()
    print "Found", nList, "entries passing cuts."
    stopT = time.clock()
    print "Data loading time (s): ", (stopT - startT)

    # In batch mode ONLY, create an output file+tree & append new branches
    outputFile = TFile()
    outTree = TTree()
    if batMode:
        outputFile = TFile(outFile, "RECREATE")
        print "Attempting tree copy to", outFile
        outTree = waveTree.CopyTree("")
        outTree.Write()
        print "Wrote", outTree.GetEntries(), "entries."
        cutUsed = TNamed("cutUsedHere", theCut)
        cutUsed.Write()
    fitStart = std.vector("double")()
    fitE = std.vector("double")()
    fitSlo = std.vector("double")()
    fitStartSD = std.vector("double")()
    fitESD = std.vector("double")()
    fitSloSD = std.vector("double")()
    fitChi2NDF = std.vector("double")()
    fitLnLike = std.vector("double")()
    tExp1 = std.vector("double")()
    tExp2 = std.vector("double")()
    tPol0 = std.vector("double")()
    tPol1 = std.vector("double")()
    tPol2 = std.vector("double")()
    tPol3 = std.vector("double")()
    baseAvg = std.vector("double")()
    baseNoise = std.vector("double")()
    bFitStart = outTree.Branch("fitStart", fitStart)
    bFitE = outTree.Branch("fitE", fitE)
    bFitSlo = outTree.Branch("fitSlo", fitSlo)
    bFitStart_sd = outTree.Branch("fitStartSD", fitStartSD)
    bFitE_sd = outTree.Branch("fitESD", fitESD)
    bFitSlo_sd = outTree.Branch("fitSloSD", fitSloSD)
    bFitChi2NDF = outTree.Branch("fitChi2NDF", fitChi2NDF)
    bFitLnLike = outTree.Branch("fitLnLike", fitLnLike)
    bTExp1 = outTree.Branch("tExp1", tExp1)
    bTExp2 = outTree.Branch("tExp2", tExp2)
    bTPol0 = outTree.Branch("tPol0", tPol0)
    bTPol1 = outTree.Branch("tPol1", tPol1)
    bTPol2 = outTree.Branch("tPol2", tPol2)
    bTPol3 = outTree.Branch("tPol3", tPol3)
    bBaseAvg = outTree.Branch("baseAvg", baseAvg)
    bBaseNoise = outTree.Branch("baseNoise", baseNoise)

    # Make a figure
    # with PdfPages('multipage_pdf.pdf') as pdf:
    fig = plt.figure(figsize=(11, 7), facecolor='w')
    p1 = plt.subplot2grid((6, 7), (0, 0), colspan=4, rowspan=2)  # original
    p2 = plt.subplot2grid((6, 7), (2, 0), colspan=4, rowspan=3)  # rising edge
    p3 = plt.subplot2grid((6, 7), (5, 0), colspan=4, rowspan=1)  # residual
    p4 = plt.subplot2grid((6, 7), (0, 4), colspan=3, rowspan=2)  # trace 1
    p5 = plt.subplot2grid((6, 7), (2, 4), colspan=3, rowspan=2,
                          sharex=p4)  # trace 2
    p6 = plt.subplot2grid((6, 7), (4, 4), colspan=3, rowspan=2,
                          sharex=p4)  # trace 3
    if not batMode: plt.show(block=False)

    # Setup multiprocessing
    manager = Manager()
    returnDict = manager.dict()

    # Loop over events
    while (True):
        saveMe = False
        iList += 1
        if intMode == True and iList != 0:
            value = raw_input()
            if value == 'q': break
            if value == 's': saveMe = True
            if value == 'p': iList -= 2  # previous
            if (value.isdigit()): iList = int(value)  # go to entry
        if iList >= elist.GetN(): break

        entry = waveTree.GetEntryNumber(iList)
        waveTree.LoadTree(entry)
        waveTree.GetEntry(entry)
        nChans = waveTree.channel.size()
        numPass = waveTree.Draw("channel", theCut, "GOFF", 1, iList)
        chans = waveTree.GetV1()
        chanList = list(set(int(chans[n]) for n in xrange(numPass)))

        fitStart.assign(nChans, -9999)
        fitE.assign(nChans, -9999)
        fitSlo.assign(nChans, -9999)
        fitStartSD.assign(nChans, -9999)
        fitESD.assign(nChans, -9999)
        fitSloSD.assign(nChans, -9999)
        fitChi2NDF.assign(nChans, -9999)
        fitLnLike.assign(nChans, -9999)
        tExp1.assign(nChans, -9999)
        tExp2.assign(nChans, -9999)
        tPol0.assign(nChans, -9999)
        tPol1.assign(nChans, -9999)
        tPol2.assign(nChans, -9999)
        tPol3.assign(nChans, -9999)
        baseAvg.assign(nChans, -9999)
        baseNoise.assign(nChans, -9999)

        # Loop over hits passing cuts
        hitList = (iH for iH in xrange(nChans)
                   if waveTree.channel.at(iH) in chanList
                   )  # a 'generator expression'
        for iH in hitList:

            # Load waveform for this hit
            run = waveTree.run
            chan = waveTree.channel.at(iH)
            dataE = waveTree.trapENFCal.at(iH)
            dataST = waveTree.butterTime.at(iH)  # replace with blrwfFMR50?
            toe = waveTree.kvorrT.at(iH) / dataE
            print "%d / %d  Run %d  nCh %d  chan %d  trapENF %.1f  t/e %.1f" % (
                iList, nList, run, nChans, chan, dataE, toe)
            signal = wl.processWaveform(waveTree.MGTWaveforms.at(iH),
                                        opt='full')
            waveBLSub = signal.GetWaveBLSub()
            waveFilt = signal.GetWaveFilt()
            waveTS = signal.GetTS()
            dataBaseline, dataNoise = signal.GetBaseNoise()

            # Denoise the data waveform (take only lowest-frequency components)
            wp = pywt.WaveletPacket(data=waveBLSub,
                                    wavelet='haar',
                                    mode='symmetric',
                                    maxlevel=3)
            new_wp = pywt.WaveletPacket(data=None,
                                        wavelet='haar',
                                        mode='symmetric')
            new_wp['aaa'] = wp['aaa'].data
            waveDenoised = new_wp.reconstruct(update=False)

            # Window the fit around rising edge - start time calculator method
            loWin, hiWin = dataST - 1000, dataST + 4000  # ns
            if loWin < waveTS[0] or hiWin > waveTS[-1]:
                print "Window out of range!  dataST: %.1f  loWin %.1f  hiWin %.1f" % (
                    dataST, loWin, hiWin)
            idx = np.where((waveTS >= loWin) & (waveTS <= hiWin))
            data = waveBLSub[idx]
            # data = waveDenoised[idx]
            dataTS = waveTS[idx]

            # Pack into lists
            rawList = [waveBLSub, waveTS, dataE, dataST]
            dataList = [data, dataTS, dataE, dataST, loWin, hiWin]
            tempList = [temp, tempTS, tempE, tempST]

            # Optionally save something to a file
            if saveMe:
                print "Saved entry", iList, iH
                np.savez("./data/tailSlopeInputs.npz", rawList, tempList)

            # Recreate the guess and the guess's rising edge
            guessFull, guessFullTS = wm.MakeModel(dataList,
                                                  tempList,
                                                  [dataST, dataE, 1.],
                                                  opt="full")
            guess, guessTS = wm.MakeModel(dataList,
                                          tempList, [dataST, dataE, 1.],
                                          opt="!fancy")

            # Make an "almost complete" guess - no MCMC
            # st, en, slo = dataST-100, dataE, 5
            # InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
            # model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)

            # Fit with MCMC and get best-fit parameters
            numSteps, burnIn = 3000, 1800  # default: 10000, 5000.  fast: 3000, 1800  long test: 20000,10000
            wfModel = wm.TemplateModel(dataList, dataNoise, tempList)
            p = Process(target=RunMCMC,
                        args=(wfModel, numSteps, burnIn, returnDict))
            p.start()
            p.join()
            startTimeTr = returnDict["startTimeTr"]
            energyTr = returnDict["energyTr"]
            slownessTr = returnDict["slownessTr"]
            st = np.median(startTimeTr[burnIn:])
            en = np.median(energyTr[burnIn:])
            slo = np.median(slownessTr[burnIn:])
            InterpFn = interpolate.interp1d(tempTS,
                                            temp,
                                            kind="linear",
                                            copy="False",
                                            assume_sorted="True")
            model, modelTS = wm.MakeModel(dataList,
                                          tempList, [st, en, slo],
                                          fn=InterpFn)

            # Save some extra parameters for the ROOT output
            # Calculate residual, Chi2/NDF, likelihood, etc.
            st_std = np.std(startTimeTr[burnIn:])
            en_std = np.std(energyTr[burnIn:])
            slo_std = np.std(slownessTr[burnIn:])
            residual = model - data
            frac = (np.power(data - model, 2)) / np.abs(model)
            chi2NDF = np.sum(frac) / len(model)
            inv_sigma2 = 1.0 / (dataNoise**2)
            lnLike = -0.5 * (np.sum((data - model)**2 * inv_sigma2 -
                                    np.log(inv_sigma2)))

            # ** Do a separate & simple fit of the tail slope **
            # TODO: Add this to process-waveforms.py
            idxMax = np.where(
                guessFull == guessFull.max())  # returns an array/tuple
            idxMax = idxMax[0][0]  # "cast" to int
            tail, tailTS = waveDenoised[idxMax:], waveTS[idxMax:]
            popt, _ = curve_fit(wl.tailModelPol, tailTS, tail)  # poly fit
            pol0, pol1, pol2, pol3 = popt[0], popt[1], popt[2], popt[3]
            a, b = dataE, 72000
            popt2, _ = curve_fit(wl.tailModelExp, tailTS, tail,
                                 p0=[a, b])  # expo fit
            e1, e2 = popt2[0], popt2[1]

            # Assign values to output vectors and fill branches
            fitStart[iH], fitStartSD[iH] = st, st_std
            fitE[iH], fitESD[iH] = en, en_std
            fitSlo[iH], fitSloSD[iH] = slo, slo_std
            fitChi2NDF[iH] = chi2NDF
            fitLnLike[iH] = lnLike
            tExp1[iH], tExp2[iH] = e1, e2
            tPol0[iH], tPol1[iH], tPol2[iH], tPol3[iH] = pol0, pol1, pol2, pol3
            baseAvg[iH] = dataBaseline
            baseNoise[iH] = dataNoise
            if batMode:
                bFitStart.Fill()
                bFitE.Fill()
                bFitSlo.Fill()
                bFitStart_sd.Fill()
                bFitE_sd.Fill()
                bFitSlo_sd.Fill()
                bFitChi2NDF.Fill()
                bFitLnLike.Fill()
                bTExp1.Fill()
                bTExp2.Fill()
                bTPol0.Fill()
                bTPol1.Fill()
                bTPol2.Fill()
                bTPol3.Fill()
                bBaseAvg.Fill()
                bBaseNoise.Fill()
                if iList % 5000 == 0:
                    outTree.Write("", TObject.kOverwrite)
                    print "%d / %d entries saved (%.2f %% done)." % (
                        iList, nList, 100 * (float(iList) / nList))

            # If not in batch mode, fill the figure
            if batMode: continue
            p1.cla()
            p1.set_ylabel("ADC")
            p1.set_title(
                "Run %d  Channel %d  Entry %d\ntrapENFCal %.1f  T/E %.1f  ST %.1f"
                % (run, chan, iList, dataE, toe, dataST))
            p1.plot(waveTS, waveBLSub, color='blue')
            p1.plot(waveTS, waveDenoised, color='red', alpha=0.8)
            p1.plot(guessFullTS, guessFull, color='orange', linewidth=2)
            p1.axvline(x=dataST, color='green', linewidth=2)
            p1.axvline(x=loWin, color='black')
            p1.axvline(x=hiWin, color='black')
            p1.plot(tailTS,
                    wl.tailModelPol(tailTS, *popt),
                    color='cyan',
                    linewidth=1)  # tail poly fit
            p1.plot(tailTS,
                    wl.tailModelExp(tailTS, *popt2),
                    color='cyan',
                    linewidth=1)  # tail expo fit

            p2.cla()
            p2.plot(dataTS, data, color='blue', label='Data')
            p2.plot(guessTS, guess, color='orange', label='Guess')
            # special: plot the values of the trace after burn-in
            # to see how the model is covering the "money-zone"/rising edge after it's converged.
            # for i in range(burnIn,numSteps):
            # st_tr, en_tr, slo_tr = M.trace('startTime')[i], M.trace('energy')[i], M.trace('slowness')[i]
            # trace, traceTS = wm.MakeModel(dataList, tempList, [st_tr,en_tr,slo_tr], fn=InterpFn)
            # p2.plot(traceTS, trace, color='red',alpha=0.1,linewidth=2)
            p2.plot(modelTS,
                    model,
                    color='red',
                    linewidth=3,
                    alpha=0.7,
                    label='Best Fit')
            p2.legend(loc=4)

            p3.cla()
            p3.set_xlabel("Time [ns]", x=0.95, ha='right')
            p3.set_ylabel("Residual [ADC]")
            p3.plot(modelTS, residual, color='red')
            p3.axhline(y=0, color='blue', alpha=0.3)
            p3.axhline(y=dataNoise, color='blue', alpha=0.3)
            p3.axhline(y=-1.0 * dataNoise, color='blue', alpha=0.3)

            p4.cla()
            minST = tempST - tempTS[-1] + hiWin
            maxST = tempST - tempTS[0] + loWin
            p4.set_title(
                "startTime %.1f  Energy %.2f\nSlow %.1f  chi2/ndf %.1f  Min %d  Max %d"
                % (st, en, slo, chi2NDF, minST, maxST))
            p4.plot(startTimeTr[:])
            p4.set_ylabel('startTime')
            p4.axvline(x=burnIn, color='red', alpha=0.5)

            p5.cla()
            p5.plot(energyTr[:])
            p5.set_ylabel('energy')
            p5.axvline(x=burnIn, color='red', alpha=0.5)

            p6.cla()
            p6.plot(slownessTr[:])
            p6.set_ylabel('slowness')
            p6.axvline(x=burnIn, color='red', alpha=0.5)

            plt.tight_layout()
            plt.subplots_adjust(hspace=0.35)
            plt.pause(scanSpeed)
            # pdf.savefig()

    # End loop over events
    if batMode:
        outTree.Write("", TObject.kOverwrite)
        print "Wrote", outTree.GetBranch(
            "channel").GetEntries(), "entries in the copied tree,"
        print "and wrote", bFitStart.GetEntries(
        ), "entries in the new branches."
Beispiel #19
0
def ApplyChannelCuts(dsNum, cutType, dType):
    """ ./lat3.py -cut [dsNum] [cutType]
    This runs over whole datasets.
    cutTypes:
        fs, rn, wf, fs+rn, fs+wf, rn+wf, fs+rn+wf

    dTypes:
        bkg, cal
    """
    # load the database
    calDB = db.TinyDB('../calDB.json')
    pars = db.Query()

    # setup a loop over modules and dataset ranges
    gROOT.ProcessLine("gErrorIgnoreLevel = 3001;")
    cInfo = ds.CalInfo()
    nMods = [1]
    if dsNum == 4: nMods = [2]
    if dsNum == 5: nMods = [1,2]
    # Dummy string for file writing -- adds nothing to the directories if background
    dString = ""
    if dType == "cal": dString = "cal"

    for modNum in nMods:
        # Changed so range of idx are set here to take advantage of module number
        nRanges = []
        if dType == "bkg":
            nRanges = [0, ds.dsMap[dsNum]]
            # if dsNum==5: nRanges[0] = 80 # exclude DS-5A
        elif dType == "cal":
            nRanges = [0, len(cInfo.master['ds%d_m%d'%(dsNum, modNum)])-1]
        else:
            print "cal or bkg not set, returning"
            return 0

        # Loop over bkgIdx, even though for calibration runs this will represent calIdx
        for bkgIdx in range(nRanges[0], nRanges[1]+1):
            # load the chains and find the right calIdx's.
            skimTree = TChain("skimTree")

            # build the file list
            fRegex = ""
            fList = []
            if dType == "bkg":
                fRegex = "/global/homes/w/wisecg/project/bg-lat/latSkimDS%d_%d_*.root" % (dsNum, bkgIdx)
                fList = glob.glob(fRegex)
                skimTree.Add(fRegex)
            elif dType == "cal":
                calList = cInfo.GetCalList("ds%d_m%d" % (dsNum, modNum), bkgIdx, runLimit=10)
                for i in calList:
                    fList += glob.glob("/global/homes/w/wisecg/project/cal-lat/latSkimDS%d_run%d_*.root"%(dsNum,i))
                    skimTree.Add("/global/homes/w/wisecg/project/cal-lat/latSkimDS%d_run%d_*.root" % (dsNum, i))
            file0 = fList[0]
            print "DS-%d subset %d, Mod-%d.  N_files: %d" % (dsNum, bkgIdx, modNum, len(fList))

            # Print some basic info about files
            f = TFile(file0)
            firstRun, lastRun, calIdxLo, calIdxHi = 0,0,0,0
            theCut = f.Get("theCut").GetTitle()
            if dType == "bkg":
                skimTree.GetEntry(0)
                firstRun = skimTree.run
                skimTree.GetEntry(skimTree.GetEntries()-1)
                lastRun = skimTree.run
                calIdxLo = cInfo.GetCalIdx("ds%d_m%d" % (dsNum, modNum), firstRun)
                calIdxHi = cInfo.GetCalIdx("ds%d_m%d" % (dsNum, modNum), lastRun)
            elif dType == "cal":
                # All the idx are the same for calibration!
                calIdxLo = calIdxHi = bkgIdx
                firstRun, lastRun = calList[0], calList[-1]

            print "    Entries %d  firstRun %d  lastRun %d  calIdxLo %d  calIdxHi %d" % (skimTree.GetEntries(),firstRun,lastRun,calIdxLo,calIdxHi)

            # build the channel list  (remove 692 and 1232 from DS5 for now.)
            chList = ds.GetGoodChanList(dsNum)
            if dsNum==5 and modNum==1:
                chList = [ch for ch in chList if ch < 1000 and ch!=692]
            if dsNum==5 and modNum==2:
                chList = [ch for ch in chList if ch > 1000 and ch!=1232]

            # -- create a dict of cuts for each channel, covering each calIdx. --
            cutDict = {}
            for calIdx in range(calIdxLo, calIdxHi+1):

                runCovMin = cInfo.master["ds%d_m%d" % (dsNum, modNum)][calIdx][1]
                runCovMax = cInfo.master["ds%d_m%d" % (dsNum, modNum)][calIdx][2]
                runCut = "run>=%d && run<=%d" % (runCovMin, runCovMax)

                fsD = ds.getDBRecord("fitSlo_ds%d_idx%d_m%d_Peak" % (dsNum, calIdx, modNum), False, calDB, pars)
                rnSD = ds.getDBRecord("riseNoise_ds%d_idx%d_m%d_SoftPlus" % (dsNum, calIdx, modNum), False, calDB, pars)
                rnCD = ds.getDBRecord("riseNoise_ds%d_idx%d_m%d_Continuum" % (dsNum, calIdx, modNum), False, calDB, pars)
                wfD = ds.getDBRecord("wfstd_ds%d_idx%d_mod%d" % (dsNum, calIdx, modNum), False, calDB, pars)

                for ch in chList:

                    fsCut, rnCut, wfCut, chanCut = None, None, None, None

                    # print ch,":",fsD[ch][2]

                    # fitSlo: check the 90% value is positive
                    if fsD[ch][2] > 0:
                        fsCut = "fitSlo<%.2f" % fsD[ch][2]

                    # riseNoise: check the softplus curvature is positive
                    if rnSD[ch][3] > 0:
                        rnCut = "riseNoise<(%.3f+%.5f*TMath::Log(1+TMath::Exp((trapENFCal-(%.3f))/%.3f)))" % (max(rnSD[ch][0],rnCD[ch][4]), rnSD[ch][1], rnSD[ch][2], rnSD[ch][3])

                    # wfStd: check if ralph says this is ok to use
                    if wfD!=0 and ch in wfD.keys() and wfD[ch][0]==u'y':
                        wfCut = "abs(wfstd - sqrt((%.4e + %.4e*trapENFCal + %.4e*trapENFCal**2 + %.2e*pow(trapENFCal,3) + %.2e*pow(trapENFCal,4))**2 + %.4f)) < (%.2f+%.3f*trapENFCal)" % (wfD[ch][3], wfD[ch][4], wfD[ch][5], wfD[ch][6], wfD[ch][7], wfD[ch][8], wfD[ch][9], wfD[ch][10])

                    # set the combination channel cut
                    if cutType == "fs" and fsCut!=None:
                        chanCut = "(%s && %s)" % (runCut, fsCut)

                    if cutType == "rn" and rnCut!=None:
                        chanCut = "(%s && %s)" % (runCut, rnCut)

                    if cutType == "wf" and wfCut!=None:
                        chanCut = "(%s && %s)" % (runCut, wfCut)

                    if cutType == "fs+rn" and fsCut!=None and rnCut!=None:
                        chanCut = "(%s && %s && %s)" % (runCut, fsCut, rnCut)

                    if cutType == "fs+wf" and fsCut!=None and wfCut!=None:
                        chanCut = "(%s && %s && %s)" % (runCut, fsCut, wfCut)

                    if cutType == "rn+wf" and rnCut!=None and wfCut!=None:
                        chanCut = "(%s && %s && %s)" % (runCut, rnCut, wfCut)

                    if cutType == "fs+rn+wf" and fsCut!=None and rnCut!=None and wfCut!=None:
                        chanCut = "(%s && %s && %s && %s)" % (runCut, fsCut, rnCut, wfCut)

                    # create dict entry for this channel or append to existing, taking care of parentheses and OR's.
                    if ch in cutDict.keys() and chanCut!=None:
                        cutDict[ch] += " || %s" % chanCut
                    elif ch not in cutDict.keys() and chanCut!=None:
                        cutDict[ch] = "(%s" % chanCut

            # close the parens for each channel entry
            for key in cutDict:
                cutDict[key] += ")"

            # -- finally, loop over each channel we have an entry for, get its cut, and create an output file. --
            for ch in cutDict:
                # TODO: threshold cut (or at least save the value for each bkgIdx)

                chanCut = theCut + "&& gain==0 && channel==%d" % ch

                if cutType == "fs":
                    outFile = "~/project/cuts/%sfs/%sfitSlo-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch)
                    chanCut += "&& fitSlo>0 && %s" % cutDict[ch]

                if cutType == "rn":
                    outFile = "~/project/cuts/%srn/%sriseNoise-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch)
                    chanCut += "&& %s" % cutDict[ch]

                if cutType == "wf":
                    outFile = "~/project/cuts/%swf/%swfstd-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch)
                    chanCut += "&& %s" % cutDict[ch]

                if cutType == "fs+rn":
                    outFile = "~/project/cuts/%sfs_rn/%sfs_rn-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch)
                    chanCut += "&& fitSlo>0 && %s" % cutDict[ch]

                if cutType == "fs+wf":
                    outFile = "~/project/cuts/%sfs_wf/%sfs_wf-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch)
                    chanCut += "&& fitSlo>0 && %s" % cutDict[ch]

                if cutType == "rn+wf":
                    outFile = "~/project/cuts/%srn_wf/%srn_wf-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch)
                    chanCut += "&& %s" % cutDict[ch]

                if cutType == "fs+rn+wf":
                    outFile = "~/project/cuts/%sfs_rn_wf/%sfs_rn_wf-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch)
                    chanCut += "&& fitSlo>0 && %s" % cutDict[ch]

                print "    Writing to:",outFile
                print "    Cut used:",chanCut,"\n"
                outFile = TFile(outFile,"RECREATE")
                outTree = TTree()
                outTree = skimTree.CopyTree(chanCut)
                outTree.Write()
                cutUsed = TNamed("chanCut",chanCut)
                cutUsed.Write()
                print "Wrote",outTree.GetEntries(),"entries."
                outFile.Close()
Beispiel #20
0
def createinputs(fname, sampleset, observables, bins, **kwargs):
    """Create histogram inputs in ROOT file for datacards.
       fname:       filename pattern of ROOT file
       sampleset:   SampleSet object
       observables: list of Variables objects
       bins:        list of Selection objects
  """
    #LOG.header("createinputs")
    outdir = kwargs.get('outdir', "")
    tag = kwargs.get('tag', "")  # file tag
    htag = kwargs.get('htag', "")  # hist tag for systematic
    filters = kwargs.get('filter',
                         None)  # only create histograms for these processes
    vetoes = kwargs.get('veto', None)  # veto these processes
    parallel = kwargs.get('parallel', True)  # MultiDraw histograms in parallel
    recreate = kwargs.get('recreate', False)  # recreate ROOT file
    replaceweight = kwargs.get('replaceweight', None)  # replace weight
    extraweight = kwargs.get('weight', "")  # extraweight
    shiftQCD = kwargs.get('shiftQCD', 0)  # e.g 0.30 for 30%
    verbosity = kwargs.get('verb', 0)
    option = 'RECREATE' if recreate else 'UPDATE'
    method = 'QCD_OSSS' if filters == None or 'QCD' in filters else None
    method = kwargs.get('method', method)

    # FILE LOGISTICS: prepare file and directories
    files = {}
    ensuredir(outdir)
    fname = os.path.join(outdir, fname)
    for obs in observables:
        obsname = obs.filename
        ftag = tag + obs.tag
        fname_ = repkey(fname, OBS=obsname, TAG=tag)
        file = TFile.Open(fname_, option)
        if recreate:
            print ">>> created file %s" % (fname_)
        for selection in bins:
            if not obs.plotfor(selection): continue
            obs.changecontext(selection)
            ensureTDirectory(file, selection.filename, cd=True, verb=verbosity)
            if recreate:
                string = joincuts(selection.selection, obs.cut)
                TNamed("selection", string).Write(
                )  # write exact selection string to ROOT file for the record / debugging
                #TNamed("weight",sampleset.weight).Write()
                LOG.verb(
                    "%s selection %r: %r" % (obsname, selection.name, string),
                    verbosity, 1)
        files[obs] = file

    # GET HISTS
    for selection in bins:
        bin = selection.filename  # bin name
        print ">>>\n>>> " + color(
            " %s " % (bin), 'magenta', bold=True, ul=True)
        if htag:
            print ">>> systematic uncertainty: %s" % (color(
                htag.lstrip('_'), 'grey'))
        if recreate or verbosity >= 1:
            print ">>> %r" % (selection.selection)
        hists = sampleset.gethists(observables,
                                   selection,
                                   method=method,
                                   split=True,
                                   parallel=parallel,
                                   filter=filters,
                                   veto=vetoes)

        # SAVE HIST
        ljust = 4 + max(11, len(htag))  # extra space
        TAB = LOG.table("%10.1f %10d  %-18s  %s")
        TAB.printheader('events', 'entries', 'variable',
                        'process'.ljust(ljust))
        for obs, hist in hists.iterhists():
            name = lreplace(hist.GetName(), obs.filename).strip(
                '_')  # histname = $VAR_$NAME (see Sample.gethist)
            if not name.endswith(htag):
                name += htag  # HIST = $PROCESS_$SYSTEMATIC
            name = repkey(name, BIN=bin)
            drawopt = 'E1' if 'data' in name else 'EHIST'
            lcolor = kBlack if any(
                s in name
                for s in ['data', 'ST', 'VV']) else hist.GetFillColor()
            hist.SetOption(drawopt)
            hist.SetLineColor(lcolor)
            hist.SetFillStyle(0)  # no fill in ROOT file
            hist.SetName(name)
            hist.GetXaxis().SetTitle(obs.title)
            for i, yval in enumerate(hist):
                if yval < 0:
                    print ">>> replace bin %d (%.3f<0) of %r" % (
                        i, yval, hist.GetName())
                    hist.SetBinContent(i, 0)
            files[obs].cd(bin)  # $FILE:$BIN/$PROCESS_$SYSTEMATC
            hist.Write(name, TH1.kOverwrite)
            TAB.printrow(hist.GetSumOfWeights(), hist.GetEntries(),
                         obs.printbins(), name)
            deletehist(hist)  # clean memory

    # CLOSE
    for obs, file in files.iteritems():
        file.Close()