def getSelection(l0, l1, jets, met):
    nClJets = len(jets)
    sel = None
    l0pt    = l0.p4.Pt()
    l1pt    = l1.p4.Pt()
    j0      = jets[0]
    mll     = (l0.p4 + l1.p4).M()
    ht      = kin.computeHt(met.p4, [l0.p4, l1.p4]+[j.p4 for j in jets])
    metrel  = kin.computeMetRel(met.p4, [l0.p4, l1.p4]+[j.p4 for j in jets])
    mtl0    = kin.computeMt(l0.p4, met.p4)
    mtl1    = kin.computeMt(l1.p4, met.p4)
    mtmin   = min([mtl0, mtl1])
    mtmax   = max([mtl0, mtl1])
    mlj     = kin.computeMlj(l0.p4, l1.p4, j0.p4)
    detall  = fabs(l0.p4.Eta() - l1.p4.Eta())
    mljj = kin.computeMljj(l0.p4, l1.p4, jets[0].p4, jets[1].p4) if nClJets>1 else None
    ll = kin.getDilepType(l0, l1)
    nj = 'eq1j' if nClJets==1 else 'ge2j'
    if (ll=='mm' and nj=='eq1j'
        and l0pt    >  30.0
        and l1pt    >  20.0
        and detall  <   1.5
        and mtmax   > 100.0
        and ht      > 200.0):
        sel = 'sr_mm_eq1j' if mlj<90.0 else 'cr_mm_eq1j'
    elif (ll=='mm' and nj=='ge2j'
          and l0pt    >  30.0
          and l1pt    >  30.0
          and detall  <   1.5
          and ht      > 220.0):
        sel ='sr_mm_ge2j' if mljj<120.0 else 'cr_mm_ge2j'
    elif (ll=='em' and nj=='eq1j'
          and l0pt    >  30.0
          and l1pt    >  30.0
          and detall  <   1.5
          and ht      > 200.0
          and mtmax   > 110.0):
        sel = 'sr_ee_eq1j' if mlj<90.0 else 'cr_ee_eq1j'
    elif (ll=='em' and nj=='ge2j'
          and l0pt    >  30.0
          and l1pt    >  30.0
          and detall  <   1.5
          and ht      > 200.0
          and mtmax   > 110.0):
        sel = 'sr_em_ge2j' if mljj<120.0 else 'cr_em_ge2j'
    elif (ll=='ee' and nj=='eq1j'
          and l0pt    >  30.0
          and l1pt    >  20.0
          and fabs(mll-91.2) > 10.0
          and metrel  >  55.0
          and ht      > 200.0):
        sel = 'sr_ee_eq1j' if mlj<90.0 else 'cr_ee_eq1j'
    elif (ll=='ee' and nj=='ge2j'
          and l0pt    >  30.0
          and l1pt    >  20.0
          and fabs(mll-91.2) > 10.0
          and metrel  >  30.0
          and mtmax   > 100.0):
        sel = 'sr_ee_ge2j' if mljj<120.0 else 'cr_ee_ge2j'
    return sel
def fillHistosAndCount(histos, files, lls, njs, testRun=False) :
    "Fill the histograms, and provide a dict of event counters[sample][sel] for the summary"
    treename = 'SusySel'
    counts = dict()
    for sample, filename in files.iteritems() :
        countsSample = collections.defaultdict(float)
        histosSample = histos[sample]
        file = r.TFile.Open(filename)
        tree = file.Get(treename)
        nEvents = tree.GetEntries()
        nEventsToProcess = nEvents if not testRun else nEvents/10
        print "processing %s (%d entries %s) %s"%(sample, nEventsToProcess, ", 10% test" if testRun else "", datetime.datetime.now())
        for iEvent, event in enumerate(tree) :
            if iEvent > nEventsToProcess : break
            l0, l1, met, pars = addTlv(event.l0), addTlv(event.l1), addTlv(event.met), event.pars
            jets, lepts = [addTlv(j) for j in event.jets], [addTlv(l) for l in event.lepts]
            ll = getDilepType(l0, l1)
            nJets = len(jets)
            nj = 'eq1j' if nJets==1 else 'ge2j'
            assert nJets>0,"messed something up in the selection upstream"
            if ll not in lls or nj not in njs : continue
            pt0 = l0.p4.Pt()
            pt1 = l1.p4.Pt()
            j0  = jets[0]
            mll  = (l0.p4 + l1.p4).M()
            mtllmet = computeMt(l0.p4 + l1.p4, met.p4)
            ht      = computeHt(met.p4, [l0.p4, l1.p4]+[j.p4 for j in jets])
            metrel  = computeMetRel(met.p4, [l0.p4, l1.p4]+[j.p4 for j in jets])
            mtl0    = computeMt(l0.p4, met.p4)
            mtl1    = computeMt(l1.p4, met.p4)
            mtmin   = min([mtl0, mtl1])
            mtmax   = max([mtl0, mtl1])
            mlj     = computeMlj(l0.p4, l1.p4, j0.p4)
            dphill  = abs(phi_mpi_pi(l0.p4.DeltaPhi(l1.p4)))
            detall  = fabs(l0.p4.Eta() - l1.p4.Eta())
            l3Veto  =  not thirdLepZcandidateIsInWindow(l0, l1, lepts)
            mljj = None
            if nJets >1 :
                j0, j1 = jets[0], jets[1]
                mt2j   = computeMt2j(l0.p4, l1.p4, j0.p4, j1.p4, met.p4)
                mljj   = computeMljj(l0.p4, l1.p4, j0.p4, j1.p4)
                dphijj = fabs(phi_mpi_pi(j0.p4.DeltaPhi(j1.p4)))
                detajj = fabs(j0.p4.Eta() - j1.p4.Eta())
            if passSelection(pt0, pt1, mll, mtllmet, ht, metrel, l3Veto,
                             detall, mtmax, mlj, mljj,
                             ll, nj) :
                llnj = llnjKey(ll, nj)
                weight = pars.weight
                varHistos = histosSample[llnj]
                varValues = dict([(v, eval(v)) for v in variablesToPlot()])
                fillVarHistos(varHistos, varValues, weight, nj)
                countsSample[llnj] += weight
        file.Close()
        file.Delete()
        counts[sample] = countsSample
    return counts
def createOutTree(filenames, dilepChan, nJetChan, tag='', overwrite=False) :
    assert dilepChan in ['ee','mm','em']
    assert nJetChan in ['eq1j', 'ge2j']
    outFilenames = dict()
    for sample, filename in filenames.iteritems() :
        outFilename = '/tmp/'+sample+'_'+dilepChan+'_'+nJetChan+'.root'
        if os.path.exists(outFilename) and not overwrite :
            outFilenames[sample] = outFilename
            continue
        outFile = r.TFile.Open(outFilename, 'recreate')
        outTree = r.TTree("training","Training tree")
        outTree.Branch('vars', vars, '/F:'.join(leafNames))
        outTree.SetDirectory(outFile)
        file = r.TFile.Open(filename)
        tree = file.Get(treename)
        print "processing %s %s %s (%d entries)"%(sample, dilepChan, nJetChan, tree.GetEntries())
        for iEvent, event in enumerate(tree) :
            resetVars(vars)
            l0 = addTlv(event.l0)
            l1 = addTlv(event.l1)
            met = addTlv(event.met)
            jets = [addTlv(j) for j in event.jets]
            lepts = [addTlv(l) for l in event.lepts]
            pars = event.pars
            dilepType = getDilepType(event.l0, event.l1)
            nJets = len(jets)
            if dilepType != dilepChan : continue
            if nJets<1 or (nJets>1 and nJetChan=='eq1j') : continue
            if thirdLepZcandidateIsInWindow(l0, l1, lepts, 20.0) : continue
            mt0, mt1 = computeMt(l0.p4, met.p4), computeMt(l1.p4, met.p4)
            vars.pt0 = l0.p4.Pt()
            vars.pt1 = l1.p4.Pt()
            vars.mll = (l0.p4+l1.p4).M()
            vars.mtmin = min([mt0, mt1])
            vars.mtmax = max([mt0, mt1])
            vars.mtllmet = computeMt(l0.p4 + l1.p4, met.p4)
            vars.ht = computeHt(met.p4, [l0.p4, l1.p4]+[j.p4 for j in jets])
            vars.metrel = computeMetRel(met.p4, [l0.p4, l1.p4]+[j.p4 for j in jets])
            vars.dphill = fabs(phi_mpi_pi(l0.p4.DeltaPhi(l1.p4)))
            vars.detall = fabs(l0.p4.Eta() - l1.p4.Eta())
            if nJets >1 :
                j0, j1 = jets[0], jets[1]
                vars.mt2j = computeMt2j(l0.p4, l1.p4, j0.p4, j1.p4, met.p4)
                vars.mljj = computeMljj(l0.p4, l1.p4, j0.p4, j1.p4)
                vars.dphijj = fabs(phi_mpi_pi(j0.p4.DeltaPhi(j1.p4)))
                vars.detajj = fabs(j0.p4.Eta() - j1.p4.Eta())
            outTree.Fill()
        print "filled ",outTree.GetEntries()," entries"
        outFile.Write()
        outFile.Close()
        outFilenames[sample] = outFile.GetName()
    return outFilenames
def getSelectionChannelOnly(l0, l1, jets, met):
    nClJets = len(jets) #note to self: these are already central-light jets b/c of passCommonCriteria
    sel = None
    ll = kin.getDilepType(l0, l1)
    nj = 'eq1j' if nClJets==1 else 'ge2j'
    mlj  = kin.computeMlj(l0.p4, l1.p4, jets[0].p4)
    mljj = kin.computeMljj(l0.p4, l1.p4, jets[0].p4, jets[1].p4) if nClJets>1 else None
    if (ll=='mm' and nj=='eq1j'):
        sel = 'sr_mm_eq1j' if mlj<90.0 else 'cr_mm_eq1j'
    elif (ll=='mm' and nj=='ge2j'):
        sel ='sr_mm_ge2j' if mljj<120.0 else 'cr_mm_ge2j'
    elif (ll=='em' and nj=='eq1j'):
        sel = 'sr_ee_eq1j' if mlj<90.0 else 'cr_ee_eq1j'
    elif (ll=='em' and nj=='ge2j'):
        sel = 'sr_em_ge2j' if mljj<120.0 else 'cr_em_ge2j'
    elif (ll=='ee' and nj=='eq1j'):
        sel = 'sr_ee_eq1j' if mlj<90.0 else 'cr_ee_eq1j'
    elif (ll=='ee' and nj=='ge2j'):
        sel = 'sr_ee_ge2j' if mljj<120.0 else 'cr_ee_ge2j'
    return sel