예제 #1
0
def test_chain_draw_hist_init_first():
    if sys.version_info[0] >= 3:
        raise SkipTest("Python 3 support not implemented")
    hist = Hist(100, 0, 1)
    chain = TreeChain('tree', FILE_PATHS)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)
예제 #2
0
def test_chain_draw_hist_init_first():
    if sys.version_info[0] >= 3:
        raise SkipTest("Python 3 support not implemented")
    hist = Hist(100, 0, 1)
    chain = TreeChain('tree', FILE_PATHS)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)
예제 #3
0
    def test_chain_draw(self):

        chain = TreeChain('tree', self.file_paths)
        hist = Hist(100, 0, 1)
        chain.draw('a_x', hist=hist)
        assert_equals(hist.Integral() > 0, True)

        # check that Draw can be repeated
        hist2 = Hist(100, 0, 1)
        chain.draw('a_x', hist=hist2)
        assert_equals(hist.Integral(), hist2.Integral())
예제 #4
0
파일: test_tree.py 프로젝트: suvayu/rootpy
def test_chain_draw():

    chain = TreeChain("tree", FILE_PATHS)
    hist = Hist(100, 0, 1)
    chain.draw("a_x", hist=hist)
    assert_equal(hist.Integral() > 0, True)

    # check that Draw can be repeated
    hist2 = Hist(100, 0, 1)
    chain.draw("a_x", hist=hist2)
    assert_equal(hist.Integral(), hist2.Integral())
예제 #5
0
def get_histograms_from_trees(
                              trees = [],
                              branch = 'var',
                              weightBranch = 'EventWeight',
                              selection = '1',
                              files = {},
                              verbose = False,
                              nBins = 40,
                              xMin = 0,
                              xMax = 100,
                              ignoreUnderflow = True,
                              ):
    histograms = {}
    nHistograms = 0

    # Setup selection and weight string for ttree draw
    weightAndSelection = '( %s ) * ( %s )' % ( weightBranch, selection )

    for sample, input_file in files.iteritems():

        histograms[sample] = {}

        for tree in trees:

            tempTree = tree
            if 'data' in sample and ( 'Up' in tempTree or 'Down' in tempTree ) :
                tempTree = tempTree.replace('_'+tempTree.split('_')[-1],'')


            chain = None;
            if isinstance( input_file, list ):
                for f in input_file:
                    chain.Add(f)
            else:
                chain = TreeChain(tempTree, [input_file]);


            weightAndSelection = '( %s ) * ( %s )' % ( weightBranch, selection )

            root_histogram = Hist( nBins, xMin, xMax, type='D')
            chain.Draw(branch, weightAndSelection, hist = root_histogram)
            if not is_valid_histogram( root_histogram, tree, input_file):
                return

            # When a tree is filled with a dummy variable, it will end up in the underflow, so ignore it
            if ignoreUnderflow:
                root_histogram.SetBinContent(0, 0)
                root_histogram.SetBinError(0,0)

            gcd()
            nHistograms += 1
            histograms[sample][tree] = root_histogram.Clone()

    return histograms
예제 #6
0
def test_chain_draw():
    if sys.version_info[0] >= 3:
        raise SkipTest("Python 3 support not implemented")
    chain = TreeChain('tree', FILE_PATHS)
    hist = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)

    # check that Draw can be repeated
    hist2 = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist2)
    assert_equal(hist.Integral(), hist2.Integral())
예제 #7
0
def load_tree(infiles):
    from rootpy.tree import TreeChain
    from rootpy.ROOT import gROOT
    gROOT.SetBatch(True)

    print('[INFO] Opening files: {0}'.format(infiles))
    tree = TreeChain('ntupler/tree', infiles)
    tree.define_collection(name='hits', prefix='vh_', size='vh_size')
    tree.define_collection(name='simhits', prefix='vc_', size='vc_size')
    tree.define_collection(name='tracks', prefix='vt_', size='vt_size')
    tree.define_collection(name='particles', prefix='vp_', size='vp_size')
    return tree
예제 #8
0
def test_chain_iter():
    if sys.version_info[0] >= 3:
        raise SkipTest("Python 3 support not implemented")
    chain = TreeChain('tree', FILE_PATHS)
    assert_equal(len(chain), 3)  # 3 files
    entries = 0
    for entry in chain:
        entries += 1
    assert_equal(entries, 300)
    entries = 0
    for entry in chain:
        entries += 1
    assert_equal(entries, 300)
    assert_equal(chain.GetEntries(), 300)
    assert_equal(chain.GetEntriesFast(), 300)
예제 #9
0
    def loop(self):

        tree = TreeChain(self.treeName, self.inFile)
        self.Nev = 0
        for f in self.inFile:
            tmpf = root_open(f)

            for path, dirs, objects in tmpf.walk():
                if "Events" in objects:
                    evHist = tmpf.Get(path + "/Events")
                    self.Nev += evHist.GetBinContent(1)
                    break
            tmpf.close()

        counter = 0
        for e in tree:
            self.Fill(0, "zp_pt", e.mu1_pt)

            #just profile stuff
            self.Profile(0, "test_profile", e.mu1_pt, e.mu1_pt / 14)
            if (e.mu1_pt > 30):
                #make cuts
                self.Fill(1, "zp_pt", e.mu1_pt / 15)

        self.writeFile()
예제 #10
0
 def __init__(self,
              files,
              events=-1,
              load_emu_trees=False,
              load_reco_trees=True):
     from cmsl1t.utils.root_glob import glob
     input_files = []
     for f in files:
         if '*' in f:
             input_files.extend(glob(f))
         else:
             input_files.append(f)
     # this is not efficient
     self._trees = []
     self._names = []
     load_ROOT_library('L1TAnalysisDataformats.so')
     allTrees = get_trees(load_emu_trees, load_reco_trees)
     for name, path in allTrees.iteritems():
         try:
             chain = TreeChain(path, input_files, cache=True, events=events)
         except RuntimeError:
             logger.warn("Cannot find tree: {0} in input file".format(path))
             continue
         self._names.append(name)
         self._trees.append(chain)
def purge_bad_files(infiles):
    good_files = []
    for infile in infiles:
        try:
            _ = TreeChain('ntupler/tree', infile)
            good_files.append(infile)
        except:
            pass
    return good_files
예제 #12
0
    def __init__(self, files=None, type='MC', maxevents=-1):
        if type.upper() not in ['MC', 'DATA']:
            raise ValueError("Argument `type` need to be MC/DATA")
        self.Type = type
        self.MaxEvents = maxevents

        if not files: raise ValueError("Argument `files` need to be non-empty")
        if isinstance(files, str): files = [
                files,
        ]
        self.Chain = TreeChain('ffNtuplizer/ffNtuple', files)

        ## register collections ###
        self.Chain.define_collection('pvs', prefix='pv_', size='pv_n')
        self.Chain.define_collection('muons', prefix='muon_', size='muon_n')
        self.Chain.define_collection('dsamuons',
                                     prefix='dsamuon_',
                                     size='dsamuon_n')
        self.Chain.define_collection('ak4jets',
                                     prefix='akjet_ak4PFJetsCHS_',
                                     size='akjet_ak4PFJetsCHS_n')
        self.Chain.define_collection('leptonjets',
                                     prefix='pfjet_',
                                     size='pfjet_n',
                                     mix=LeptonJetMix)

        self.Chain.define_collection('trigobjs',
                                     prefix='trigobj_',
                                     size='trigobj_n')

        self.Chain.define_object('hlt', prefix='HLT_')
        self.Chain.define_object('metfilters', prefix='metfilters_')
        self.Chain.define_object('cosmicveto', prefix='cosmicveto_')

        self.Histos = {}

        self.LookupWeight = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/puWeights_10x_56ifb.root')).Get(
                    'puWeights')
        self.Scale = 1.
예제 #13
0
 def _load_trees(self):
     for treeName in self._treeNames:
         try:
             self._trees[treeName] = TreeChain(
                 treeName,
                 self.input_files,
                 cache=True,
                 events=self.nevents,
             )
         except RuntimeError:
             logger.warn(
                 "Cannot find tree: {0} in input file".format(treeName))
             continue
def copy_in_trigger_signal(in_files_name,
                           out_name,
                           tree_name,
                           prefix,
                           cdc_events,
                           cth_events,
                           rand_t=None):
    # Convert input lists to sets first
    set_cdc_events = set(cdc_events)
    set_cth_events = set(cth_events)

    # Define the chain of input trees
    in_chain = TreeChain(name=tree_name, files=in_files_name)
    # First create a new file to save the new tree in:
    out_file = root_open(out_name, "r+")

    # Add the time shift if we want it in the tree
    ExtraBranches = Tagged
    if rand_t is not None:
        ExtraBranches += Smeared

    # Get the new tree with its extra branches
    out_tree = Tree(tree_name, model=ExtraBranches.prefix(prefix))

    # This creates all the same branches in the new tree but
    # their addresses point to the same memory used by the original tree.
    out_tree.create_branches(in_chain._buffer)
    out_tree.update_buffer(in_chain._buffer)

    # Now loop over the original tree(s) and fill the new tree
    for entry in in_chain:
        # Add in the new values
        this_event_number = entry[prefix + "EventNumber"].value
        out_tree.__setattr__(prefix + "GoodTrack", this_event_number
                             in set_cdc_events)
        out_tree.__setattr__(prefix + "GoodTrig", this_event_number
                             in set_cth_events)
        if rand_t is not None:
            try:
                out_tree.__setattr__(prefix + "SmearTime",
                                     rand_t[this_event_number])
            except:
                for key, item in entry.iteritems():
                    print key, item
        # Fill, noting that most of the buffer is shared between the chain
        # and the output tree
        out_tree.Fill()
    # Close it up
    out_tree.Write()
    out_file.Close()
예제 #15
0
def test_chain_draw():

    chain = TreeChain('tree', FILE_PATHS)
    hist = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)

    # check that Draw can be repeated
    hist2 = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist2)
    assert_equal(hist.Integral(), hist2.Integral())
예제 #16
0
    def test_chain_draw(self):

        chain = TreeChain('tree', self.file_paths)
        hist = Hist(100, 0, 1)
        chain.draw('a_x', hist=hist)
        assert_equals(hist.Integral() > 0, True)

        # check that Draw can be repeated
        hist2 = Hist(100, 0, 1)
        chain.draw('a_x', hist=hist2)
        assert_equals(hist.Integral(), hist2.Integral())
예제 #17
0
def test_chain_draw():
    if sys.version_info[0] >= 3:
        raise SkipTest("Python 3 support not implemented")
    chain = TreeChain('tree', FILE_PATHS)
    hist = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)

    # check that Draw can be repeated
    hist2 = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist2)
    assert_equal(hist.Integral(), hist2.Integral())
예제 #18
0
def test_chain_draw():
    if sys.version_info[0] >= 3:
        raise SkipTest("Python 3 support not implemented")
    chain = TreeChain('tree', FILE_PATHS)
    hist = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)
    assert_equal(hist.GetEntries(), 300)

    # check that Draw can be repeated
    hist2 = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist2)
    assert_equal(hist.Integral(), hist2.Integral())

    # draw into a graph
    graph = chain.draw("a_x:a_y")
    assert_true(isinstance(graph, _GraphBase))
    assert_equal(len(graph), chain.GetEntries())
    assert_equal(len(graph), 300)
예제 #19
0
 def __init__(self, files, events=-1):
     from cmsl1t.utils.root_glob import glob
     input_files = []
     for f in files:
         if '*' in f:
             input_files.extend(glob(f))
         else:
             input_files.append(f)
     # this is not efficient
     self._trees = []
     self._names = []
     for name, path in ALL_TREE.iteritems():
         try:
             chain = TreeChain(path, input_files, cache=True, events=events)
         except RuntimeError:
             logger.warn("Cannot find tree: {0} in input file".format(path))
             continue
         self._names.append(name)
         self._trees.append(chain)
예제 #20
0
def mergeChannel(channel, fileList):
    '''
    Merge one channel for all files, return merged tree. 
    Should be run with a file open.
    '''
    if isinstance(fileList, str):
        fileList = [fileList]

    chain = TreeChain('{}/ntuple'.format(channel), fileList)
    out = Tree('ntuple'.format(channel))
    out.set_buffer(chain._buffer, create_branches=True)

    found = set()

    for ev in chain:
        evID = (ev.run, ev.lumi, ev.evt)
        if evID in found:
            continue

        out.fill()
        found.add(evID)

    return out
예제 #21
0
파일: test_tree.py 프로젝트: ndawe/rootpy
def test_chain_draw():
    if sys.version_info[0] >= 3:
        raise SkipTest("Python 3 support not implemented")
    chain = TreeChain('tree', FILE_PATHS)
    hist = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)
    assert_equal(hist.GetEntries(), 300)

    # check that Draw can be repeated
    hist2 = Hist(100, 0, 1)
    chain.draw('a_x', hist=hist2)
    assert_equal(hist.Integral(), hist2.Integral())

    # draw into a graph
    graph = chain.draw("a_x:a_y")
    assert_true(isinstance(graph, _GraphBase))
    assert_equal(len(graph), chain.GetEntries())
    assert_equal(len(graph), 300)
예제 #22
0
    def work(self):

        D4PD_model = RecoTauBlock + RecoJetBlock + EventVariables + TrueTauBlock + PartonBlock

        # initialize the TreeChain of all input files (each containing one tree named self.metadata.treename)
        tree = TreeChain(self.metadata.treename,
                         files=self.files,
                         events=self.events,
                         cache=True,
                         cache_size=10000000,
                         learn_entries=30)

        # create output tree
        self.output.cd()
        D4PD = Tree(name=self.metadata.name, model=D4PD_model)

        copied_variables = ['actualIntPerXing',
                            'averageIntPerXing']
        copied_variables += mc_triggers

        D4PD.set_buffer(tree.buffer, variables=copied_variables, create_branches=True, visible=False)
        tree.always_read(copied_variables)

        # set the event filters
        # passthrough for MC for trigger acceptance studies
        self.event_filters = EventFilterList([
            GRLFilter(self.grl, passthrough = self.metadata.datatype != datasets.DATA),
            PriVertex(),
            LArError(),
            JetCleaningLoose(passthrough = self.metadata.datatype != datasets.DATA),
            #JetCleaningMedium(passthrough = self.metadata.datatype != datasets.DATA),
            LArHole(),
            #JetCrackVeto(),
            ElectronVeto(),
            MuonVeto(),
            TauElectronVeto(),
            TauMuonVeto(),
            TauAuthorTrack(),
        ])

        self.event_filters.insert(1, MCTriggers())
        tree.filters += self.event_filters

        cutflow = Cutflow()

        # define tree collections
        tree.define_collection(name="taus", prefix="tau_", size="tau_n", mix=TauFourMomentum)
        # jet_eta etc is AntiKt4LCTopo in tau-perf D3PDs
        tree.define_collection(name="jets", prefix="jet_", size="jet_n", mix=FourMomentum)
        tree.define_collection(name="truetaus", prefix="trueTau_", size="trueTau_n", mix=MCTauFourMomentum)
        tree.define_collection(name="mc", prefix="mc_", size="mc_n", mix=MCParticle)
        tree.define_collection(name="muons", prefix="mu_staco_", size="mu_staco_n")
        tree.define_collection(name="electrons", prefix="el_", size="el_n")
        tree.define_collection(name="vertices", prefix="vxp_", size="vxp_n")

        # define tree objects
        D4PD.define_object(name='tau1', prefix='tau1_')
        D4PD.define_object(name='tau2', prefix='tau2_')
        D4PD.define_object(name='jet1', prefix='jet1_')
        D4PD.define_object(name='jet2', prefix='jet2_')

        """
        tree.define_association(origin='taus', target='truetaus', prefix='trueTauAssoc_', link='index')
        tree.define_association(origin='truetaus', target='taus', prefix='tauAssoc_', link='index')
        """

        # entering the main event loop...
        for event in tree:

            D4PD.reset()
            cutflow.reset()

            # tau selection
            event.taus.select(lambda tau: tau.pt > 15*GeV)
            # Jet selection
            event.jets.select(lambda jet: jet.fourvect.P() > 25*GeV and abs(jet.emscale_eta) < 4.5)

            """
            Get VBF jets
            """
            # get partons (already sorted by eta in hepmc)
            parton1, parton2 = hepmc.get_VBF_partons(event)
            PartonBlock.set(D4PD, parton1, parton2)
            D4PD.dR_quarks = parton1.fourvect.DeltaR(parton2.fourvect)

            """
            Get true taus
            """
            event.truetaus.select(lambda tau: tau.vis_Et > 10 * GeV and abs(tau.vis_eta) < 2.5)
            if len(event.truetaus) > 2:
                print "ERROR: too many true taus: %i" % len(event.truetaus)
                D4PD.error = 1
                D4PD.Fill()
                continue
            elif len(event.truetaus) < 2:
                print "ERROR: too few true taus: %i" % len(event.truetaus)
                D4PD.error = 2
                D4PD.Fill()
                continue

            """
            fourvects = []
            colors = []
            radii = []
            for thing in event.jets:
                fourvects.append(thing.fourvect)
                colors.append('blue')
                radii.append(.4)
            for parton in (parton1, parton2):
                fourvects.append(parton.fourvect)
                colors.append('green')
                radii.append(.1)
            for thing in event.taus:
                fourvects.append(thing.fourvect)
                colors.append('red')
                radii.append(.2)
            for tau in event.truetaus:
                fourvects.append(tau.fourvect)
                colors.append('purple')
                radii.append(.1)

            eventview.draw(event, fourvects, colors=colors, radii=radii)
            """

            TrueTauBlock.set(D4PD, 1, event.truetaus[0])
            TrueTauBlock.set(D4PD, 2, event.truetaus[1])

            D4PD.dR_truetaus = event.truetaus[0].fourvect.DeltaR(event.truetaus[1].fourvect)

            D4PD.dR_quark_tau = min([
                                    parton1.fourvect.DeltaR(event.truetaus[0].fourvect),
                                    parton2.fourvect.DeltaR(event.truetaus[0].fourvect),
                                    parton1.fourvect.DeltaR(event.truetaus[1].fourvect),
                                    parton2.fourvect.DeltaR(event.truetaus[1].fourvect),
                                    ])

            taus = []
            if event.taus:
                for truetau in event.truetaus:
                    closest_tau = min(event.taus, key=lambda tau: utils.dR(tau.seedCalo_eta, tau.seedCalo_phi, truetau.vis_eta, truetau.vis_phi))
                    if utils.dR(closest_tau.seedCalo_eta, closest_tau.seedCalo_phi, truetau.vis_eta, truetau.vis_phi) < 0.2:
                        if closest_tau in taus:
                            # collision
                            D4PD.error = 3
                            break
                        taus.append(closest_tau)
            if len(taus) < 2:
                # collision
                D4PD.Fill()
                continue

            """
            # Overlap removal between taus and jets
            event.jets.select(lambda jet: not any([tau for tau in taus if
                              utils.dR(jet.emscale_eta, jet.emscale_phi, tau.seedCalo_eta, tau.seedCalo_phi) < .2]))
            """

            jets = []
            if event.jets:
                for quark in (parton1, parton2):
                    closest_jet = min(event.jets, key=lambda jet: utils.dR(jet.eta, jet.phi, quark.eta, quark.phi))
                    if utils.dR(closest_jet.eta, closest_jet.phi, quark.eta, quark.phi) < 0.4:
                        if closest_jet in jets:
                            # collision
                            D4PD.error = 4
                            break
                        jets.append(closest_jet)
            if len(jets) < 2:
                # collision
                D4PD.Fill()
                continue

            """
            Jet variables
            """
            RecoJetBlock.set(D4PD, jets[0], jets[1])

            """
            Reco tau variables
            This must come after the RecoJetBlock is filled since
            that sets the jet_beta for boosting the taus
            """
            RecoTauBlock.set(event, D4PD, taus[0], taus[1])
            D4PD.true_Mvis_tau1_tau2 = (D4PD.trueTau1_fourvect_vis + D4PD.trueTau2_fourvect_vis).M()

            """
            MET
            """
            METx = event.MET_RefFinal_etx
            METy = event.MET_RefFinal_ety
            MET_vect = Vector2(METx, METy)
            MET_3vect = Vector3(METx, METy, 0.)

            D4PD.MET = event.MET_RefFinal_et
            D4PD.MET_phi = event.MET_RefFinal_phi

            # HT TODO: Fix
            sumET = event.MET_RefFinal_sumet
            D4PD.HT = sumET

            D4PD.numVertices = len([vtx for vtx in event.vertices if (vtx.type == 1 and vtx.nTracks >= 4) or
                                    (vtx.type == 3 and vtx.nTracks >= 2)])
            # fill output ntuple
            # use reset=True to reset all variables to their defaults after the fill
            # to avoid any values from this event carrying over into the next
            D4PD.cutflow = cutflow.int()
            D4PD.Fill()

        self.output.cd()
        D4PD.FlushBaskets()
        D4PD.Write()
예제 #23
0
    def __init__(self, fns, outn, outd, model):
        print "Initializing Container!"
        #self.tin = r.TChain('EcalVeto')
        #for fn in fns:
        #    self.tin.Add(fn)
        #self.tfile = root_open(fn,'r+')
        #with root_open(fn,'r+') as f:
        #self.tin = self.tfile.EcalVeto
        #self.tin.Print()
        self.tin = TreeChain('EcalVeto', fns)
        self.tin.create_branches({'discValue_gabrielle': 'F'})
        self.model = model

        self.outdir = outd
        self.outname = outn
        self.outfile = root_open(outn, 'RECREATE')
        self.tout = Tree('EcalVeto')
        self.tout.set_buffer(self.tin._buffer, create_branches=True)

        self.events = []
        #print self.tin.GetEntries()
        for event in self.tin:
            #if len(self.events)>10:
            #    continue
            evt = []
            ################################### Features #######################################
            evt.append(event.nReadoutHits)
            evt.append(event.summedDet)
            evt.append(event.summedTightIso)
            evt.append(event.maxCellDep)
            evt.append(event.showerRMS)
            evt.append(event.xStd)
            evt.append(event.yStd)
            evt.append(event.avgLayerHit)
            evt.append(event.deepestLayerHit)
            evt.append(event.stdLayerHit)
            #new features
            evt.append(event.ele68ContEnergy)
            evt.append(event.ele68x2ContEnergy)
            evt.append(event.ele68x3ContEnergy)
            evt.append(event.ele68x4ContEnergy)
            evt.append(event.ele68x5ContEnergy)
            evt.append(event.photon68ContEnergy)
            evt.append(event.photon68x2ContEnergy)
            evt.append(event.photon68x3ContEnergy)
            evt.append(event.photon68x4ContEnergy)
            evt.append(event.photon68x5ContEnergy)
            evt.append(event.outside68ContEnergy)
            evt.append(event.outside68x2ContEnergy)
            evt.append(event.outside68x3ContEnergy)
            evt.append(event.outside68x4ContEnergy)
            evt.append(event.outside68x5ContEnergy)
            evt.append(event.outside68ContNHits)
            evt.append(event.outside68x2ContNHits)
            evt.append(event.outside68x3ContNHits)
            evt.append(event.outside68x4ContNHits)
            evt.append(event.outside68x5ContNHits)
            evt.append(event.outside68ContXstd)
            evt.append(event.outside68x2ContXstd)
            evt.append(event.outside68x3ContXstd)
            evt.append(event.outside68x4ContXstd)
            evt.append(event.outside68x5ContXstd)
            evt.append(event.outside68ContYstd)
            evt.append(event.outside68x2ContYstd)
            evt.append(event.outside68x3ContYstd)
            evt.append(event.outside68x4ContYstd)
            evt.append(event.outside68x5ContYstd)
            evt.append(event.ecalBackEnergy)

            evtarray = np.array([evt])
            pred = float(model.predict(xgb.DMatrix(evtarray))[0])
            #print pred
            event.discValue_gabrielle = pred
            self.tout.Fill()
            ######################################################################################
            self.events.append(evt)

            if (len(self.events) % 10000 == 0 and len(self.events) > 0):
                print 'The shape of events = ', np.shape(self.events)

        self.outfile.cd()
        self.tout.Write()

        self.outfile.Close()
        #self.tfile.Close()
        print 'cp %s %s' % (self.outname, self.outdir)
        os.system('cp %s %s' % (self.outname, self.outdir))
예제 #24
0
파일: chain.py 프로젝트: zhufengGNSS/rootpy
tree = Tree("test")
tree.create_branches(branches)

for i in range(10000):
    tree.x = gauss(.5, 1.)
    tree.y = gauss(.3, 2.)
    tree.z = gauss(13., 42.)
    tree.i = i
    tree.fill()
tree.write()

# Make a histogram of the second tree
hist2 = Hist(100, -10, 10, name='hist2')
tree.Draw('x', 'y > 1', hist=hist2)
hist2.SetDirectory(0)  # memory resident
print("The second tree has {0:f} entries where y > 1".format(hist2.Integral()))
f.close()

combined_hist = hist1 + hist2

print("Building TreeChain")
chain = TreeChain('test', ['chaintest2.root', 'chaintest1.root'])
# Make the equivalent of the combined_hist
combined_hist_chain = Hist(100, -10, 10, name='combined')
chain.Draw('x', 'y > 1', hist=combined_hist_chain)

residual = combined_hist_chain - combined_hist
print("The combined histogram (separately) minus "
      "the combined from the chain has {0:f} entries".format(
          residual.Integral()))
예제 #25
0
 def __init__(self, files, events=-1):
     # this is not efficient
     self._trees = [
         TreeChain(name, files, cache=True, events=events)
         for name in TREE_NAMES
     ]
예제 #26
0
    def work(self):

        # trigger config tool to read trigger info in the ntuples
        trigger_config = get_trigger_config()

        OutputModel = (RecoTauBlock + EventVariables + SkimExtraModel +
            TrueTauBlock)

        onfilechange = []
        # update the trigger config maps on every file change
        onfilechange.append((update_trigger_config, (trigger_config,)))

        cutflow = Hist(2, 0, 2, name='cutflow', type='D')

        # initialize the TreeChain of all input files (each containing one tree named self.metadata.treename)
        chain = TreeChain(self.metadata.treename,
                         files=self.files,
                         events=self.events,
                         cache=True,
                         cache_size=10000000,
                         learn_entries=30,
                         onfilechange=onfilechange)

        # create output tree
        self.output.cd()
        tree = Tree(name='higgstautauhh', model=OutputModel)

        copied_variables = ['actualIntPerXing',
                            'averageIntPerXing',
                            'RunNumber',
                            'EventNumber',
                            'lbn']

        tree.set_buffer(
                chain.buffer,
                branches=copied_variables,
                create_branches=True,
                visible=False)
        chain.always_read(copied_variables)

        # set the event filters
        event_filters = EventFilterList([
            #Triggers(
            #    datatype=self.metadata.datatype,
            #    year=YEAR,
            #    skim=False),
            PriVertex(),
            LArError(),
            LArHole(datatype=self.metadata.datatype),
            JetCleaning(
                datatype=self.metadata.datatype,
                year=YEAR),
            TauAuthor(1),
            TauHasTrack(1),
            TauPT(1, thresh=25 * GeV),
            TauEta(1),
            TauCrack(1),
            TauLArHole(1),
            #TauTriggerMatch(
            #    config=trigger_config,
            #    year=YEAR,
            #    datatype=self.metadata.datatype,
            #    skim=False,
            #    tree=tree,
            #    min_taus=1),
        ])

        self.filters['event'] = event_filters
        chain.filters += event_filters

        # define tree collections
        chain.define_collection(name="taus", prefix="tau_", size="tau_n", mix=TauFourMomentum)
        chain.define_collection(name="taus_EF", prefix="trig_EF_tau_",
                                size="trig_EF_tau_n", mix=TauFourMomentum)

        # jet_* etc. is AntiKt4LCTopo_* in tau-perf D3PDs
        chain.define_collection(name="jets", prefix="jet_", size="jet_n", mix=FourMomentum)
        chain.define_collection(name="truetaus", prefix="trueTau_", size="trueTau_n", mix=MCTauFourMomentum)
        chain.define_collection(name="mc", prefix="mc_", size="mc_n", mix=MCParticle)
        chain.define_collection(name="muons", prefix="mu_staco_", size="mu_staco_n")
        chain.define_collection(name="electrons", prefix="el_", size="el_n")
        chain.define_collection(name="vertices", prefix="vxp_", size="vxp_n")

        from externaltools import PileupReweighting
        from ROOT import Root
        # Initialize the pileup reweighting tool
        pileup_tool = Root.TPileupReweighting()
        if YEAR == 2011:
            pileup_tool.AddConfigFile(PileupReweighting.get_resource('mc11b_defaults.prw.root'))
            pileup_tool.AddLumiCalcFile('lumi/2011/hadhad/ilumicalc_histograms_None_178044-191933.root')
        elif YEAR == 2012:
            pileup_tool.AddConfigFile(PileupReweighting.get_resource('mc12a_defaults.prw.root'))
            pileup_tool.SetDataScaleFactors(1./1.11)
            pileup_tool.AddLumiCalcFile('lumi/2012/hadhad/ilumicalc_histograms_None_200841-205113.root')
        else:
            raise ValueError('No pileup reweighting defined for year %d' %
                    YEAR)
        # discard unrepresented data (with mu not simulated in MC)
        pileup_tool.SetUnrepresentedDataAction(2)
        pileup_tool.Initialize()

        # entering the main event loop...
        for event in chain:
            tree.reset()

            event.vertices.select(vertex_selection)
            tree.number_of_good_vertices = len(event.vertices)

            # match only with visible true taus
            event.truetaus.select(lambda tau: tau.vis_Et > 10 * GeV and abs(tau.vis_eta) < 2.5)

            if len(event.truetaus) == 1:
                true_tau = event.truetaus[0]
                TrueTauBlock.set(tree, 1, true_tau)
            else:
                continue

            # Truth-matching
            matched_reco = None
            reco_index = true_tau.tauAssoc_index
            tau = event.taus.getitem(reco_index)
            if tau in event.taus:
                matched_reco = tau
            else:
                continue

            tree.MET = event.MET_RefFinal_BDTMedium_et

            # fill tau block
            RecoTauBlock.set(event, tree, matched_reco, None)

            # set the event weight
            tree.pileup_weight = pileup_tool.GetCombinedWeight(event.RunNumber,
                                                               event.mc_channel_number,
                                                               event.averageIntPerXing)
            tree.mc_weight = event.mc_event_weight
            tree.Fill(reset=True)

        self.output.cd()
        tree.FlushBaskets()
        tree.Write()
        total_events = event_filters[0].total
        cutflow[0] = total_events
        cutflow[1] = total_events
        cutflow.Write()
예제 #27
0
    def test_chain_draw_hist_init_first(self):

        hist = Hist(100, 0, 1)
        chain = TreeChain('tree', self.file_paths)
        chain.draw('a_x', hist=hist)
        assert_equals(hist.Integral() > 0, True)
예제 #28
0
class MuonTypeLJEvents(object):
    def __init__(self, files=None, type='MC', maxevents=-1):
        if type.upper() not in ['MC', 'DATA']:
            raise ValueError("Argument `type` need to be MC/DATA")
        self.Type = type
        self.MaxEvents = maxevents

        if not files: raise ValueError("Argument `files` need to be non-empty")
        if isinstance(files, str): files = [
                files,
        ]
        self.Chain = TreeChain('ffNtuplizer/ffNtuple', files)

        ## register collections ###
        self.Chain.define_collection('pvs', prefix='pv_', size='pv_n')
        self.Chain.define_collection('muons', prefix='muon_', size='muon_n')
        self.Chain.define_collection('dsamuons',
                                     prefix='dsamuon_',
                                     size='dsamuon_n')
        self.Chain.define_collection('ak4jets',
                                     prefix='akjet_ak4PFJetsCHS_',
                                     size='akjet_ak4PFJetsCHS_n')
        self.Chain.define_collection('leptonjets',
                                     prefix='pfjet_',
                                     size='pfjet_n',
                                     mix=LeptonJetMix)

        self.Chain.define_collection('trigobjs',
                                     prefix='trigobj_',
                                     size='trigobj_n')

        self.Chain.define_object('hlt', prefix='HLT_')
        self.Chain.define_object('metfilters', prefix='metfilters_')
        self.Chain.define_object('cosmicveto', prefix='cosmicveto_')

        self.Histos = {}

        self.LookupWeight = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/puWeights_10x_56ifb.root')).Get(
                    'puWeights')
        self.Scale = 1.

    def bookHisto(self, name, hist):
        self.Histos[name] = hist

    def setScale(self, scale):
        self.Scale = scale

    def process(self):

        for i, event in enumerate(self.Chain):

            if self.MaxEvents > 0 and i > self.MaxEvents: break

            ## event-level mask ##
            if not event.cosmicveto.result: continue
            if not event.metfilters.PrimaryVertexFilter: continue

            aux = {}
            ## event weight ##
            aux['wgt'] = self.Scale
            if self.Type == 'MC':
                aux['wgt'] *= event.weight  # gen weight
                aux['wgt'] *= self.LookupWeight.GetBinContent(
                    self.LookupWeight.GetXaxis().FindBin(
                        event.trueInteractionNum))  ## pileup correction

            leptonjets = [
                lj for lj in event.leptonjets if lj.passSelection(event)
            ]
            leptonjets.sort(key=lambda lj: lj.p4.pt(), reverse=True)
            leptonjets = leptonjets[:2]
            muontypeljs = [lj for lj in leptonjets if lj.isMuonType()]
            egmtypeljs = [lj for lj in leptonjets if lj.isEgmType()]
            # if not muontypeljs: continue

            aux['muontype'] = muontypeljs
            aux['egmtype'] = egmtypeljs
            aux['leptonjets'] = leptonjets

            looseMuonIdx, mediumMuonIdx = [], []
            for i, mu in enumerate(event.muons):
                if mu.p4.pt() < 5: continue
                if abs(mu.p4.eta()) > 2.4: continue
                if (mu.selectors & (1 << 0)) != (1 << 0): continue  # ID-loose
                if (mu.selectors & (1 << 7)) != (1 << 7): continue  # Iso-loose
                looseMuonIdx.append(i)
                if (mu.selectors & (1 << 1)) != (1 << 1): continue  # ID-medium
                mediumMuonIdx.append(i)

            if not mediumMuonIdx: continue
            muonpairs = []
            for m in mediumMuonIdx:
                for l in looseMuonIdx:
                    if m == l: continue
                    if (l, m) in muonpairs: continue
                    muonpairs.append((m, l))
            if not muonpairs: continue

            aux['muonpairs'] = muonpairs

            self.processEvent(event, aux)

    def processEvent(self, event, aux):
        """To be override by daughter class"""
        pass

    @property
    def histos(self):
        return self.Histos
예제 #29
0
    def work(self):
        """
        This is the one function that all "ATLASStudent"s must implement.
        """
        datatype = self.metadata.datatype
        year = self.metadata.year
        verbose = self.args.verbose

        OutputModel = C3POEvent

        if datatype == datasets.MC:
            # only create truth branches for MC
            OutputModel += (
                    FourVectModel.prefix('resonance_') +
                    TrueTau.prefix('truetau1_') +
                    TrueTau.prefix('truetau2_'))

        onfilechange = []
        count_funcs = {}

        if datatype in (datasets.MC, datasets.EMBED):

            def mc_weight_count(event):
                return event.mc_event_weight

            count_funcs = {
                'mc_weight': mc_weight_count,
            }

        trigger_config = None

        if datatype != datasets.EMBED:
            # trigger config tool to read trigger info in the ntuples
            trigger_config = get_trigger_config()

            # update the trigger config maps on every file change
            onfilechange.append((update_trigger_config, (trigger_config,)))

        if datatype == datasets.DATA:
            merged_grl = GRL()

            def update_grl(student, grl, name, file, tree):

                grl |= str(file.Get('Lumi/%s' % student.metadata.treename).GetString())

            onfilechange.append((update_grl, (self, merged_grl,)))

        if datatype == datasets.DATA:
            merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D')
        else:
            merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D')

        def update_cutflow(student, cutflow, name, file, tree):

            year = student.metadata.year
            datatype = student.metadata.datatype
            if datatype == datasets.MC:
                cutflow[0] += file.cutflow_event[0]
                cutflow[1] += file.cutflow_event_mc_weight[0]
            else:
                cutflow[0] += file.cutflow_event[0]

        onfilechange.append((update_cutflow, (self, merged_cutflow,)))

        # initialize the TreeChain of all input files
        # (each containing one tree named self.metadata.treename)
        chain = TreeChain(
                self.metadata.treename,
                files=self.files,
                events=self.events,
                read_branches_on_demand=True,
                cache=True,
                onfilechange=onfilechange)

        # create output tree
        self.output.cd()
        tree = Tree(name='higgstautauhh', model=OutputModel)

        copied_variables = [
                'actualIntPerXing',
                'averageIntPerXing',
                'RunNumber',
                'EventNumber',
                'lbn']

        tree.set_buffer(
                chain._buffer,
                branches=copied_variables,
                create_branches=True,
                visible=False)

        chain.always_read(copied_variables)

        # set the event filters
        event_filters = EventFilterList([
            CoreFlags(
                count_funcs=count_funcs),
            TauSelected(2,
                count_funcs=count_funcs),
            TruthMatching(
                passthrough=datatype != datasets.MC,
                count_funcs=count_funcs),
            MCWeight(
                datatype=datatype,
                tree=tree,
                passthrough=datatype != datasets.MC,
                count_funcs=count_funcs)
        ])

        self.filters['event'] = event_filters

        chain._filters += event_filters

        define_objects(chain, year, skim=False)

        # define tree objects
        taus = [
            tree.define_object(name='tau1', prefix='tau1_'),
            tree.define_object(name='tau2', prefix='tau2_')]

        if datatype == datasets.MC:
            truetaus = [
                tree.define_object(name='truetau1', prefix='truetau1_'),
                tree.define_object(name='truetau2', prefix='truetau2_')]

            tree.define_object(name='resonance', prefix='resonance_')

        # entering the main event loop...
        for event in chain:

            # sort taus and jets in decreasing order by pT
            event.taus.sort(key=lambda tau: tau.pt, reverse=True)

            tau1, tau2 = event.taus

            # MET
            METx = event.MET.etx
            METy = event.MET.ety
            MET_vect = Vector2(METx, METy)
            MET = event.MET.et
            MET_phi = event.MET.phi

            tree.MET = MET
            tree.MET_x = METx
            tree.MET_y = METy
            tree.MET_phi = MET_phi

            sumET = event.MET.sumet
            tree.sumET = sumET
            if sumET != 0:
                tree.MET_sig = ((2. * MET / GeV) /
                        (utils.sign(sumET) * sqrt(abs(sumET / GeV))))
            else:
                tree.MET_sig = -1.

            # use MMC values from skim
            mmc_mass = event.tau_MMC_mass
            mmc_resonance = event.tau_MMC_resonance
            mmc_met = Vector2(event.tau_MMC_MET_x, event.tau_MMC_MET_y)

            tree.mass_mmc_tau1_tau2 = mmc_mass
            tree.mmc_resonance.copy_from(mmc_resonance)
            if mmc_mass > 0:
                tree.mmc_resonance_pt = mmc_resonance.Pt()
            tree.MET_mmc = mmc_met.Mod()
            tree.MET_mmc_x = mmc_met.X()
            tree.MET_mmc_y = mmc_met.Y()
            tree.MET_mmc_phi = math.pi - mmc_met.Phi()

            # truth matching
            if datatype == datasets.MC:

                resonance, tau_decays = get_taus(event)

                if resonance is not None:

                    FourVectModel.set(tree.resonance, resonance)

                    matched_taus = []
                    decays = tau_decays[:]
                    for itau, tau in enumerate(event.taus):
                        for idecay, tau_decay in enumerate(decays):
                            if tau.matches_vect(tau_decay.fourvect_visible):
                                tau_decay.matched = True
                                tau_decay.matched_object = tau
                                tau.matched = True
                                tau.matched_object = tau_decay
                                TrueTau.set(truetaus[itau], tau_decay,
                                        verbose=verbose)
                                decays.pop(idecay)
                                matched_taus.append(itau)
                                break

                    if len(decays) > 0:
                        for idecay, decay in enumerate(decays):
                            reco_idx = -1
                            remaining_idx = range(2)
                            for imatched in remaining_idx:
                                if imatched not in matched_taus:
                                    reco_idx = imatched
                                    remaining_idx.remove(imatched)
                                    break
                            TrueTau.set(truetaus[reco_idx], tau_decay,
                                    verbose=verbose)

                    if len(tau_decays) == 2:
                        # write truth met
                        fourvect_missing = (tau_decays[0].fourvect_missing +
                                            tau_decays[1].fourvect_missing)

                        tree.MET_true = fourvect_missing.Pt()
                        tree.MET_phi_true = fourvect_missing.Phi()
                        tree.MET_x_true = tree.MET_true * math.cos(tree.MET_phi_true)
                        tree.MET_y_true = tree.MET_true * math.sin(tree.MET_phi_true)
                        tree.MET_phi_diff = Vector2.Phi_mpi_pi(tree.MET_phi_true - MET_phi)

            # tau - vertex association
            tree.tau_same_vertex = (
                    tau1.privtx_x == tau2.privtx_x and
                    tau1.privtx_y == tau2.privtx_y and
                    tau1.privtx_z == tau2.privtx_z)

            # fill tau block
            for outtau, intau in zip(taus, event.taus):
                RecoTau.set(outtau, intau, verbose=verbose)

            # fill output ntuple
            tree.Fill(reset=True)

        self.output.cd()
        tree.FlushBaskets()
        tree.Write()

        if datatype == datasets.DATA:
            xml_string = ROOT.TObjString(merged_grl.str())
            xml_string.Write('lumi')
        merged_cutflow.Write()
예제 #30
0
def test_chain_draw_hist_init_first():

    hist = Hist(100, 0, 1)
    chain = TreeChain('tree', FILE_PATHS)
    chain.draw('a_x', hist=hist)
    assert_equal(hist.Integral() > 0, True)
예제 #31
0
class Events(object):
    def __init__(self,
                 files=None,
                 outname=None,
                 type='MC',
                 dtag='',
                 maxevents=-1,
                 channel=['4mu', '2mu2e'],
                 ctau=None,
                 chargedlj=False):

        if type.upper() not in ['MC', 'DATA']:
            raise ValueError("Argument `type` need to be MC/DATA")
        self.OutName = outname
        self.Type = type.upper()
        self.ChargedLJ = chargedlj
        self.MaxEvents = maxevents
        self.Channel = channel
        self.Dtag = dtag
        self.Ctau = ctau
        __signal_sample_param = dict(
            [substr.split('-') for substr in self.Dtag.split('_')])
        self.SignalParam = {
            k.upper(): float(v.replace('p', '.'))
            for k, v in __signal_sample_param.items()
        }

        if not files: raise ValueError("Argument `files` need to be non-empty")
        if isinstance(files, str): files = [
                files,
        ]
        self.Chain = TreeChain('ffNtuplizer/ffNtuple', files)

        ### register collections ###
        # self.Chain.define_collection('pvs', prefix='pv_', size='pv_n')
        self.Chain.define_collection('electrons',
                                     prefix='electron_',
                                     size='electron_n')
        self.Chain.define_collection('muons', prefix='muon_', size='muon_n')
        self.Chain.define_collection('dsamuons',
                                     prefix='dsamuon_',
                                     size='dsamuon_n')
        self.Chain.define_collection('photons',
                                     prefix='photon_',
                                     size='photon_n')
        self.Chain.define_collection('ak4jets',
                                     prefix='akjet_ak4PFJetsCHS_',
                                     size='akjet_ak4PFJetsCHS_n')
        self.Chain.define_collection('hftagscores',
                                     prefix='hftagscore_',
                                     size='hftagscore_n')
        self.Chain.define_collection('leptonjets',
                                     prefix='pfjet_',
                                     size='pfjet_n',
                                     mix=LeptonJetMix)
        self.Chain.define_collection('ljsources',
                                     prefix='ljsource_',
                                     size='ljsource_n')
        self.Chain.define_collection('cosmicmuons',
                                     prefix='cosmicmuon_',
                                     size='cosmicmuon_n')

        self.Chain.define_collection('trigobjs',
                                     prefix='trigobj_',
                                     size='trigobj_n')
        #        self.Chain.define_collection('',)

        self.Chain.define_object('hlt', prefix='HLT_')
        self.Chain.define_object('metfilters', prefix='metfilters_')
        self.Chain.define_object('cosmicveto', prefix='cosmicveto_')

        self.Triggers = [
            "DoubleL2Mu23NoVtx_2Cha",
            "DoubleL2Mu23NoVtx_2Cha_NoL2Matched",
            "DoubleL2Mu23NoVtx_2Cha_CosmicSeed",
            "DoubleL2Mu23NoVtx_2Cha_CosmicSeed_NoL2Matched",
            "DoubleL2Mu25NoVtx_2Cha_Eta2p4",
            "DoubleL2Mu25NoVtx_2Cha_CosmicSeed_Eta2p4",
        ]

        #self.addTRG = [
        # "DoubleL2Mu23NoVtx_2Cha",]

        self.Histos = {}
        for chan in channel:
            self.Histos['{}/cutflow'.format(chan)] = ROOT.Hist(
                20, 0, 20, title='cutflow', drawstyle='hist')
        self.KeepCutFlow = False
        self.RawCutFlow = False

        self.LookupWeight = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/PUWeights_2018.root')).Get(
                    'puWeights')
        self.LookupMuonSFLowpT = root_open(
            os.path.join(os.getenv('CMSSW_BASE'),
                         'src/FireROOT/Analysis/data/mu_Loose_pt7.root')).Get(
                             'ratio_syst')
        self.LookupMuonSF = root_open(
            os.path.join(os.getenv('CMSSW_BASE'),
                         'src/FireROOT/Analysis/data/RunABCD_SF_ID.root')).Get(
                             'NUM_LooseID_DEN_TrackerMuons_pt_abseta_syst')
        self.LookupElectronSF = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/2018_ElectronLoose.root')).Get(
                    'EGamma_SF2D')
        self.LookupPhotonSF = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/2018_PhotonsLoose.root')).Get(
                    'EGamma_SF2D')
        self.Scale = 1.

    def setTriggers(self, triggers):
        self.Triggers = triggers

    def addTrigger(self, trigger):
        self.Triggers.append(trigger)

    def bookHisto(self, name, hist):
        self.Histos[name] = hist

    def setScale(self, scale):
        self.Scale = scale

    def process(self):

        for i, event in enumerate(self.Chain):
            if self.MaxEvents > 0 and i > self.MaxEvents: break

            ## event weight ##
            aux = {}
            aux['wgt'] = self.Scale
            if self.Type == 'MC':
                aux['wgt'] *= event.weight  # gen weight
                aux['wgt'] *= self.LookupWeight.GetBinContent(
                    self.LookupWeight.GetXaxis().FindBin(
                        event.trueInteractionNum))  ## pileup correction

            for ch in self.Channel:
                if self.RawCutFlow:
                    self.Histos['{}/cutflow'.format(ch)].Fill(0)
                else:
                    self.Histos['{}/cutflow'.format(ch)].Fill(0, aux['wgt'])

            ## trigger ##
            #if not any([getattr(event.hlt, t) for t in self.Triggers]): continue

            #if not any([getattr(event.hlt, t) for t in self.addTRG]): continue

            for ch in self.Channel:
                if self.RawCutFlow:
                    self.Histos['{}/cutflow'.format(ch)].Fill(1)
                else:
                    self.Histos['{}/cutflow'.format(ch)].Fill(1, aux['wgt'])

            ## 2 leptonjets in channel definition ##
            leptonjets = [lj for lj in event.leptonjets]
            if len(leptonjets) < 2: continue

            #print len(leptonjets)#js
            #LJ2 = none
            leptonjets.sort(key=lambda lj: lj.p4.pt(), reverse=True)
            LJ0 = leptonjets[0]
            LJ1 = leptonjets[1]
            #LJ2 = leptonjets[2]
            #if len(event.leptonjets)>2:#
            #LJ2 = leptonjets[2]#
            #aux['lj2'] = LJ2

            if self.ChargedLJ:
                if not LJ0.passChargedSelection(event): continue
                if not LJ1.passChargedSelection(event): continue
                #if not LJ2.passChargedSelection(event): continue
            else:
                if not LJ0.passSelection(event): continue
                if not LJ1.passSelection(event): continue
                #if not LJ2.passSelection(event): continue

            if LJ0.isMuonType() and LJ1.isMuonType(): aux['channel'] = '4mu'
            elif LJ0.isMuonType() and LJ1.isEgmType(): aux['channel'] = '2mu2e'
            elif LJ0.isEgmType() and LJ1.isMuonType(): aux['channel'] = '2mu2e'
            else: continue

            aux['lj0'] = LJ0
            aux['lj1'] = LJ1
            #aux['lj2'] = LJ2

            if self.Type == 'MC':
                aux['sf_electron'] = 1.
                aux['sf_electron_up'] = 1.
                aux['sf_electron_low'] = 1.

                aux['sf_photon'] = 1.
                aux['sf_photon_up'] = 1.
                aux['sf_photon_low'] = 1.

                aux['sf_pfmuon'] = 1.
                aux['sf_pfmuon_up'] = 1.
                aux['sf_pfmuon_low'] = 1.

                for lj in [LJ0, LJ1]:
                    ## muon scale factor
                    for i in lj.pfcand_pfmuonIdx:
                        _pfmu_p4 = event.muons[i].p4
                        pt, eta = _pfmu_p4.pt(), _pfmu_p4.eta()
                        if pt >= 20:
                            xbin = self.LookupMuonSF.xaxis.FindBin(pt)
                            xbin = min(max(xbin, 1),
                                       self.LookupMuonSF.nbins(0))
                            ybin = self.LookupMuonSF.yaxis.FindBin(abs(eta))
                            sf = self.LookupMuonSF.GetBinContent(xbin, ybin)
                            err_up = self.LookupMuonSF.GetBinErrorUp(
                                xbin, ybin)
                            err_lo = self.LookupMuonSF.GetBinErrorLow(
                                xbin, ybin)
                        else:
                            for i in range(self.LookupMuonSFLowpT.num_points):
                                x = self.LookupMuonSFLowpT.x(i)
                                xh = x + self.LookupMuonSFLowpT.xerrh(i)
                                xl = x - self.LookupMuonSFLowpT.xerrl(i)
                                if xl <= eta and eta <= xh:
                                    sf = self.LookupMuonSFLowpT.y(i)
                                    err_up = self.LookupMuonSFLowpT.yerrh(i)
                                    err_lo = self.LookupMuonSFLowpT.yerrl(i)
                                    break
                        aux['wgt'] *= sf

                        aux['sf_pfmuon'] *= sf
                        aux['sf_pfmuon_up'] *= sf + err_up
                        aux['sf_pfmuon_low'] *= sf - err_lo

                    ## NOTE DSA scale factor, nothing for now
                    for i in lj.pfcand_dsamuonIdx:
                        #print dsamu.pt.P4()
                        sf = 1.
                        aux['wgt'] *= sf

                    ## electron scale factor
                    for i in lj.pfcand_electronIdx:
                        _electron = event.electrons[i]
                        xbin = self.LookupElectronSF.xaxis.FindBin(
                            _electron.scEta)
                        ybin = self.LookupElectronSF.xaxis.FindBin(
                            _electron.p4.pt())
                        ybin = min(max(ybin, 1),
                                   self.LookupElectronSF.nbins(1))
                        sf = self.LookupElectronSF.GetBinContent(xbin, ybin)
                        aux['wgt'] *= sf

                        aux['sf_electron'] *= sf
                        aux['sf_electron_up'] *= sf + self.LookupElectronSF.GetBinErrorUp(
                            xbin, ybin)
                        aux['sf_electron_low'] *= sf - self.LookupElectronSF.GetBinErrorLow(
                            xbin, ybin)

                    ## photon scale factor
                    for i in lj.pfcand_photonIdx:
                        _photon = event.photons[i]
                        xbin = self.LookupPhotonSF.xaxis.FindBin(_photon.scEta)
                        ybin = self.LookupPhotonSF.xaxis.FindBin(
                            _photon.p4.pt())
                        ybin = min(max(ybin, 1), self.LookupPhotonSF.nbins(1))
                        sf = self.LookupPhotonSF.GetBinContent(xbin, ybin)
                        aux['wgt'] *= sf

                        aux['sf_photon'] *= sf
                        aux['sf_photon_up'] *= sf + self.LookupPhotonSF.GetBinErrorUp(
                            xbin, ybin)
                        aux['sf_photon_low'] *= sf - self.LookupPhotonSF.GetBinErrorLow(
                            xbin, ybin)

                aux['wgt_electron_up'] = aux['wgt'] / aux['sf_electron'] * aux[
                    'sf_electron_up']
                aux['wgt_electron_low'] = aux['wgt'] / aux[
                    'sf_electron'] * aux['sf_electron_low']
                aux['wgt_photon_up'] = aux['wgt'] / aux['sf_photon'] * aux[
                    'sf_photon_up']
                aux['wgt_photon_low'] = aux['wgt'] / aux['sf_photon'] * aux[
                    'sf_photon_low']
                aux['wgt_pfmuon_up'] = aux['wgt'] / aux['sf_pfmuon'] * aux[
                    'sf_pfmuon_up']
                aux['wgt_pfmuon_low'] = aux['wgt'] / aux['sf_pfmuon'] * aux[
                    'sf_pfmuon_low']

                # for t, pt, eta in zip(list(lj.pfcand_type), list(lj.pfcand_pt), list(lj.pfcand_eta)):
                #     ## muon scale factor, DSA same as muon for now
                #     if t==3 or t==8:
                #         xbin = self.LookupMuonSF.xaxis.FindBin(pt)
                #         xbin = min(max(xbin, 1), self.LookupMuonSF.nbins(0))
                #         ybin = self.LookupMuonSF.yaxis.FindBin(abs(eta))
                #         sf = self.LookupMuonSF.GetBinContent(xbin, ybin)
                #         aux['wgt'] *= sf
                #     ## electron scale factor, using eta instead of SC eta for now
                #     if t==2:
                #         xbin = self.LookupElectronSF.xaxis.FindBin(eta)
                #         ybin = self.LookupElectronSF.xaxis.FindBin(pt)
                #         ybin = min(max(ybin, 1), self.LookupElectronSF.nbins(1))
                #         sf = self.LookupElectronSF.GetBinContent(xbin, ybin)
                #         aux['wgt'] *= sf
                #     ## photon scale factor, using eta instead of SC eta for now
                #     if t==4:
                #         xbin = self.LookupPhotonSF.xaxis.FindBin(eta)
                #         ybin = self.LookupPhotonSF.xaxis.FindBin(pt)
                #         ybin = min(max(ybin, 1), self.LookupPhotonSF.nbins(1))
                #         sf = self.LookupPhotonSF.GetBinContent(xbin, ybin)
                #         aux['wgt'] *= sf

            for ch in self.Channel:
                if self.RawCutFlow:
                    self.Histos['{}/cutflow'.format(ch)].Fill(2)
                else:
                    self.Histos['{}/cutflow'.format(ch)].Fill(2, aux['wgt'])

            ## event-level mask ##
            if not event.metfilters.PrimaryVertexFilter: continue
            for ch in self.Channel:
                if self.RawCutFlow:
                    self.Histos['{}/cutflow'.format(ch)].Fill(3)
                else:
                    self.Histos['{}/cutflow'.format(ch)].Fill(3, aux['wgt'])

            nppCOSMIC, cosmicShowerTagged = globalCosmicShower(
                event.cosmicmuons, aux['channel'])
            if cosmicShowerTagged: continue

            for ch in self.Channel:
                if self.RawCutFlow:
                    self.Histos['{}/cutflow'.format(ch)].Fill(4)
                else:
                    self.Histos['{}/cutflow'.format(ch)].Fill(4, aux['wgt'])

            self.processEvent(event, aux)

    def processEvent(self, event, aux):
        """To be override by daughter class"""
        pass

    def postProcess(self):
        if self.KeepCutFlow:
            labels = [
                'total', 'trigger_pass', 'leptonjet_ge2', 'pv_good',
                'cosmicveto_pass'
            ]
            for ch in self.Channel:
                xaxis = self.Histos['{}/cutflow'.format(ch)].axis(0)
                for i, s in enumerate(labels, start=1):
                    xaxis.SetBinLabel(i, s)
                    # binNum., labAngel, labSize, labAlign, labColor, labFont, labText
                    xaxis.ChangeLabel(i, 315, -1, 11, -1, -1, s)
        else:
            for ch in self.Channel:
                self.Histos.pop('{}/cutflow'.format(ch))

    @property
    def histos(self):
        return self.Histos

    @property
    def channel(self):
        return self.Channel
    pass


# ______________________________________________________________________________
# Analyzer

mystate = 0

verbose = False

batch_mode = True

# Open file
if mystate == 0:
    tree_name_i = '/home/jlow/L1MuonTrigger/CRAB3/P2_9_2_3_patch1/crab_projects/crab_ntuple_SingleMuon_PositiveEndCap/results/ntuple_SingleMuon_PositiveEndCap_%i.root'
    tree = TreeChain('ntupler/tree',
                     [(tree_name_i % (i + 1)) for i in range(8)])
    #tree_name_i = '/home/jlow/L1MuonTrigger/CRAB3/P2_9_2_3_patch1/crab_projects/crab_ntuple_SingleMuon_PositiveEndCap_PU200/results/ntuple_SingleMuon_PositiveEndCap_PU200_%i.root'
    #tree = TreeChain('ntupler/tree', [(tree_name_i % (i+1)) for i in range(50)])

elif mystate == 1:
    tree_name_i = '/home/jlow/L1MuonTrigger/CRAB3/P2_9_2_3_patch1/crab_projects/crab_ntuple_SingleNeutrino_PU200/results/ntuple_SingleNeutrino_PU200_%i.root'
    tree = TreeChain('ntupler/tree',
                     [(tree_name_i % (i + 1)) for i in range(50)])
else:
    raise Exception("Unexpected state: %i" % mystate)

maxEvents = -1
#maxEvents = 2000

# ROOT globals
gROOT.SetBatch(batch_mode)
예제 #33
0
    def __init__(self,
                 files=None,
                 outname=None,
                 type='MC',
                 dtag='',
                 maxevents=-1,
                 channel=['4mu', '2mu2e'],
                 ctau=None,
                 chargedlj=False):

        if type.upper() not in ['MC', 'DATA']:
            raise ValueError("Argument `type` need to be MC/DATA")
        self.OutName = outname
        self.Type = type.upper()
        self.ChargedLJ = chargedlj
        self.MaxEvents = maxevents
        self.Channel = channel
        self.Dtag = dtag
        self.Ctau = ctau
        __signal_sample_param = dict(
            [substr.split('-') for substr in self.Dtag.split('_')])
        self.SignalParam = {
            k.upper(): float(v.replace('p', '.'))
            for k, v in __signal_sample_param.items()
        }

        if not files: raise ValueError("Argument `files` need to be non-empty")
        if isinstance(files, str): files = [
                files,
        ]
        self.Chain = TreeChain('ffNtuplizer/ffNtuple', files)

        ### register collections ###
        # self.Chain.define_collection('pvs', prefix='pv_', size='pv_n')
        self.Chain.define_collection('electrons',
                                     prefix='electron_',
                                     size='electron_n')
        self.Chain.define_collection('muons', prefix='muon_', size='muon_n')
        self.Chain.define_collection('dsamuons',
                                     prefix='dsamuon_',
                                     size='dsamuon_n')
        self.Chain.define_collection('photons',
                                     prefix='photon_',
                                     size='photon_n')
        self.Chain.define_collection('ak4jets',
                                     prefix='akjet_ak4PFJetsCHS_',
                                     size='akjet_ak4PFJetsCHS_n')
        self.Chain.define_collection('hftagscores',
                                     prefix='hftagscore_',
                                     size='hftagscore_n')
        self.Chain.define_collection('leptonjets',
                                     prefix='pfjet_',
                                     size='pfjet_n',
                                     mix=LeptonJetMix)
        self.Chain.define_collection('ljsources',
                                     prefix='ljsource_',
                                     size='ljsource_n')
        self.Chain.define_collection('cosmicmuons',
                                     prefix='cosmicmuon_',
                                     size='cosmicmuon_n')

        self.Chain.define_collection('trigobjs',
                                     prefix='trigobj_',
                                     size='trigobj_n')
        #        self.Chain.define_collection('',)

        self.Chain.define_object('hlt', prefix='HLT_')
        self.Chain.define_object('metfilters', prefix='metfilters_')
        self.Chain.define_object('cosmicveto', prefix='cosmicveto_')

        self.Triggers = [
            "DoubleL2Mu23NoVtx_2Cha",
            "DoubleL2Mu23NoVtx_2Cha_NoL2Matched",
            "DoubleL2Mu23NoVtx_2Cha_CosmicSeed",
            "DoubleL2Mu23NoVtx_2Cha_CosmicSeed_NoL2Matched",
            "DoubleL2Mu25NoVtx_2Cha_Eta2p4",
            "DoubleL2Mu25NoVtx_2Cha_CosmicSeed_Eta2p4",
        ]

        #self.addTRG = [
        # "DoubleL2Mu23NoVtx_2Cha",]

        self.Histos = {}
        for chan in channel:
            self.Histos['{}/cutflow'.format(chan)] = ROOT.Hist(
                20, 0, 20, title='cutflow', drawstyle='hist')
        self.KeepCutFlow = False
        self.RawCutFlow = False

        self.LookupWeight = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/PUWeights_2018.root')).Get(
                    'puWeights')
        self.LookupMuonSFLowpT = root_open(
            os.path.join(os.getenv('CMSSW_BASE'),
                         'src/FireROOT/Analysis/data/mu_Loose_pt7.root')).Get(
                             'ratio_syst')
        self.LookupMuonSF = root_open(
            os.path.join(os.getenv('CMSSW_BASE'),
                         'src/FireROOT/Analysis/data/RunABCD_SF_ID.root')).Get(
                             'NUM_LooseID_DEN_TrackerMuons_pt_abseta_syst')
        self.LookupElectronSF = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/2018_ElectronLoose.root')).Get(
                    'EGamma_SF2D')
        self.LookupPhotonSF = root_open(
            os.path.join(
                os.getenv('CMSSW_BASE'),
                'src/FireROOT/Analysis/data/2018_PhotonsLoose.root')).Get(
                    'EGamma_SF2D')
        self.Scale = 1.
def load_tree_multiple(infiles):
    print('[INFO] Opening file: %s' % ' '.join(infiles))
    tree = TreeChain('ntupler/tree', infiles)
    define_collections(tree)
    return tree
예제 #35
0
파일: hhskim.py 프로젝트: ndawe/hhntup
    def work(self):
        # get argument values
        local = self.args.local
        syst_terms = self.args.syst_terms
        datatype = self.metadata.datatype
        year = self.metadata.year
        verbose = self.args.student_verbose
        very_verbose = self.args.student_very_verbose
        redo_selection = self.args.redo_selection
        nominal_values = self.args.nominal_values

        # get the dataset name
        dsname = os.getenv('INPUT_DATASET_NAME', None)
        if dsname is None:
            # attempt to guess dsname from dirname
            if self.files:
                dsname = os.path.basename(os.path.dirname(self.files[0]))

        # is this a signal sample?
        # if so we will also keep some truth information in the output below
        is_signal = datatype == datasets.MC and (
            '_VBFH' in dsname or
            '_ggH' in dsname or
            '_ZH' in dsname or
            '_WH' in dsname or
            '_ttH' in dsname)
        log.info("DATASET: {0}".format(dsname))
        log.info("IS SIGNAL: {0}".format(is_signal))

        # is this an inclusive signal sample for overlap studies?
        is_inclusive_signal = is_signal and '_inclusive' in dsname

        # is this a BCH-fixed sample? (temporary)
        is_bch_sample = 'r5470_r4540_p1344' in dsname
        if is_bch_sample:
            log.warning("this is a BCH-fixed r5470 sample")

        # onfilechange will contain a list of functions to be called as the
        # chain rolls over to each new file
        onfilechange = []
        count_funcs = {}

        if datatype != datasets.DATA:
            # count the weighted number of events
            if local:
                def mc_weight_count(event):
                    return event.hh_mc_weight
            else:
                def mc_weight_count(event):
                    return event.mc_event_weight

            count_funcs = {
                'mc_weight': mc_weight_count,
            }

        # three instances of the pileup reweighting tool are created to write
        # out the nominal, high and low pileup weights
        pileup_tool = None
        pileup_tool_high = None
        pileup_tool_low = None

        if local:
            # local means running on the skims, the output of this script
            # running on the grid
            if datatype == datasets.DATA:
                # merge the GRL fragments
                merged_grl = goodruns.GRL()

                def update_grl(student, grl, name, file, tree):
                    grl |= str(file.Get('Lumi/%s' % student.metadata.treename).GetString())

                onfilechange.append((update_grl, (self, merged_grl,)))

            if datatype == datasets.DATA:
                merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D')
            else:
                merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D')

            def update_cutflow(student, cutflow, name, file, tree):
                # record a cut-flow
                year = student.metadata.year
                datatype = student.metadata.datatype
                cutflow[1].value += file.cutflow_event[1].value
                if datatype != datasets.DATA:
                    cutflow[2].value += file.cutflow_event_mc_weight[1].value

            onfilechange.append((update_cutflow, (self, merged_cutflow,)))

        else:
            # get pileup reweighting tool
            pileup_tool = get_pileup_reweighting_tool(
                year=year,
                use_defaults=True)
            pileup_tool_high = get_pileup_reweighting_tool(
                year=year,
                use_defaults=True,
                systematic='high')
            pileup_tool_low = get_pileup_reweighting_tool(
                year=year,
                use_defaults=True,
                systematic='low')

            if datatype not in (datasets.EMBED, datasets.MCEMBED):
                # merge TrigConfTrees
                metadirname = '%sMeta' % self.metadata.treename
                trigconfchain = ROOT.TChain('%s/TrigConfTree' % metadirname)
                map(trigconfchain.Add, self.files)
                metadir = self.output.mkdir(metadirname)
                metadir.cd()
                trigconfchain.Merge(self.output, -1, 'fast keep')
                self.output.cd()

            if datatype == datasets.DATA:
                # merge GRL XML strings
                merged_grl = goodruns.GRL()
                for fname in self.files:
                    with root_open(fname) as f:
                        for key in f.Lumi.keys():
                            merged_grl |= goodruns.GRL(
                                str(key.ReadObj().GetString()),
                                from_string=True)
                lumi_dir = self.output.mkdir('Lumi')
                lumi_dir.cd()
                xml_string= ROOT.TObjString(merged_grl.str())
                xml_string.Write(self.metadata.treename)
                self.output.cd()

        self.output.cd()

        # create the output tree
        model = get_model(datatype, dsname,
                          prefix=None if local else 'hh_',
                          is_inclusive_signal=is_inclusive_signal)
        log.info("Output Model:\n\n{0}\n\n".format(model))
        outtree = Tree(name=self.metadata.treename,
                       model=model)

        if local:
            tree = outtree
        else:
            tree = outtree.define_object(name='tree', prefix='hh_')

        tree.define_object(name='tau', prefix='tau_')
        tree.define_object(name='tau1', prefix='tau1_')
        tree.define_object(name='tau2', prefix='tau2_')
        tree.define_object(name='truetau1', prefix='truetau1_')
        tree.define_object(name='truetau2', prefix='truetau2_')
        tree.define_object(name='jet1', prefix='jet1_')
        tree.define_object(name='jet2', prefix='jet2_')
        tree.define_object(name='jet3', prefix='jet3_')

        mmc_objects = [
            tree.define_object(name='mmc0', prefix='mmc0_'),
            tree.define_object(name='mmc1', prefix='mmc1_'),
            tree.define_object(name='mmc2', prefix='mmc2_'),
        ]

        for mmc_obj in mmc_objects:
            mmc_obj.define_object(name='resonance', prefix='resonance_')

        trigger_emulation = TauTriggerEmulation(
            year=year,
            passthrough=local or datatype != datasets.MC or year > 2011,
            count_funcs=count_funcs)

        if not trigger_emulation.passthrough:
            onfilechange.append(
                (update_trigger_trees, (self, trigger_emulation,)))

        trigger_config = None

        if datatype not in (datasets.EMBED, datasets.MCEMBED):
            # trigger config tool to read trigger info in the ntuples
            trigger_config = get_trigger_config()
            # update the trigger config maps on every file change
            onfilechange.append((update_trigger_config, (trigger_config,)))

        # define the list of event filters
        if local and syst_terms is None and not redo_selection:
            event_filters = None
        else:
            tau_ntrack_recounted_use_ntup = False
            if year > 2011:
                # peek at first tree to determine if the extended number of
                # tracks is already stored
                with root_open(self.files[0]) as test_file:
                    test_tree = test_file.Get(self.metadata.treename)
                    tau_ntrack_recounted_use_ntup = (
                        'tau_out_track_n_extended' in test_tree)

            event_filters = EventFilterList([
                GRLFilter(
                    self.grl,
                    passthrough=(
                        local or (
                            datatype not in (datasets.DATA, datasets.EMBED))),
                    count_funcs=count_funcs),
                CoreFlags(
                    passthrough=local,
                    count_funcs=count_funcs),
                EmbeddingPileupPatch(
                    passthrough=(
                        local or year > 2011 or datatype != datasets.EMBED),
                    count_funcs=count_funcs),
                averageIntPerXingPatch(
                    passthrough=(
                        local or year < 2012 or datatype != datasets.MC),
                    count_funcs=count_funcs),
                PileupTemplates(
                    year=year,
                    passthrough=(
                        local or is_bch_sample or datatype not in (
                            datasets.MC, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                RandomSeed(
                    datatype=datatype,
                    count_funcs=count_funcs),
                BCHSampleRunNumber(
                    passthrough=not is_bch_sample,
                    count_funcs=count_funcs),
                RandomRunNumber(
                    tree=tree,
                    datatype=datatype,
                    pileup_tool=pileup_tool,
                    passthrough=local,
                    count_funcs=count_funcs),
                trigger_emulation,
                Triggers(
                    year=year,
                    tree=tree,
                    datatype=datatype,
                    passthrough=datatype in (datasets.EMBED, datasets.MCEMBED),
                    count_funcs=count_funcs),
                PileupReweight(
                    year=year,
                    tool=pileup_tool,
                    tool_high=pileup_tool_high,
                    tool_low=pileup_tool_low,
                    tree=tree,
                    passthrough=(
                        local or (
                            datatype not in (datasets.MC, datasets.MCEMBED))),
                    count_funcs=count_funcs),
                PriVertex(
                    passthrough=local,
                    count_funcs=count_funcs),
                LArError(
                    passthrough=local,
                    count_funcs=count_funcs),
                TileError(
                    passthrough=local,
                    count_funcs=count_funcs),
                TileTrips(
                    passthrough=(
                        local or datatype in (datasets.MC, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                JetCopy(
                    tree=tree,
                    passthrough=local,
                    count_funcs=count_funcs),
                # IMPORTANT!
                # JetCalibration MUST COME BEFORE ANYTHING THAT REFERS TO
                # jet.fourvect since jet.fourvect IS CACHED!
                JetCalibration(
                    datatype=datatype,
                    year=year,
                    verbose=very_verbose,
                    passthrough=local or nominal_values,
                    count_funcs=count_funcs),
                # in situ TES shift for 2012 data
                TauEnergyShift(
                    passthrough=(
                        local or datatype != datasets.DATA
                        or year < 2012 or nominal_values),
                    count_funcs=count_funcs),
                # truth matching must come before systematics due to
                # TES_TRUE/FAKE
                TruthMatching(
                    passthrough=datatype == datasets.DATA,
                    count_funcs=count_funcs),
                NvtxJets(
                    tree=tree,
                    count_funcs=count_funcs),
                # PUT THE SYSTEMATICS "FILTER" BEFORE
                # ANY FILTERS THAT REFER TO OBJECTS
                # BUT AFTER CALIBRATIONS
                # Systematics must also come before anything that refers to
                # thing.fourvect since fourvect is cached!
                Systematics(
                    terms=syst_terms,
                    year=year,
                    datatype=datatype,
                    tree=tree,
                    verbose=verbose,
                    passthrough=not syst_terms,
                    count_funcs=count_funcs),
                JetIsPileup(
                    passthrough=(
                        local or year < 2012 or
                        datatype not in (datasets.MC, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                LArHole(
                    tree=tree,
                    passthrough=year > 2011,
                    count_funcs=count_funcs),
                JetCleaning(
                    datatype=datatype,
                    year=year,
                    count_funcs=count_funcs),
                ElectronVeto(
                    count_funcs=count_funcs),
                MuonVeto(
                    year=year,
                    count_funcs=count_funcs),
                TauPT(2,
                    thresh=20 * GeV,
                    count_funcs=count_funcs),
                TauHasTrack(2,
                    count_funcs=count_funcs),
                TauEta(2,
                    count_funcs=count_funcs),
                TauElectronVeto(2,
                    count_funcs=count_funcs),
                TauMuonVeto(2,
                    count_funcs=count_funcs),
                TauAuthor(2,
                    count_funcs=count_funcs),
                TauCrack(2,
                    count_funcs=count_funcs),
                TauLArHole(2,
                    tree=tree,
                    passthrough=year > 2011,
                    count_funcs=count_funcs),
                # before selecting the leading and subleading taus
                # be sure to only consider good candidates
                TauIDMedium(2,
                    count_funcs=count_funcs),
                #TauTriggerMatchIndex(
                #    config=trigger_config,
                #    year=year,
                #    datatype=datatype,
                #    passthrough=datatype == datasets.EMBED,
                #    count_funcs=count_funcs),
                # Select two leading taus at this point
                # 25 and 35 for data
                # 20 and 30 for MC to leave room for TES uncertainty
                TauLeadSublead(
                    lead=(
                        35 * GeV if datatype == datasets.DATA or local
                        else 30 * GeV),
                    sublead=(
                        25 * GeV if datatype == datasets.DATA or local
                        else 20 * GeV),
                    count_funcs=count_funcs),
                # taus are sorted (in decreasing order) by pT from here on
                TauIDSelection(
                    tree=tree,
                    count_funcs=count_funcs),
                TaudR(3.2,
                    count_funcs=count_funcs),
                #TauTriggerMatchThreshold(
                #    datatype=datatype,
                #    tree=tree,
                #    count_funcs=count_funcs),
                TauTriggerEfficiency(
                    year=year,
                    datatype=datatype,
                    tree=tree,
                    tes_systematic=self.args.syst_terms and (
                        Systematics.TES_TERMS & self.args.syst_terms),
                    passthrough=datatype == datasets.DATA,
                    count_funcs=count_funcs),
                PileupScale(
                    tree=tree,
                    year=year,
                    datatype=datatype,
                    passthrough=local,
                    count_funcs=count_funcs),
                TauIDScaleFactors(
                    year=year,
                    passthrough=datatype == datasets.DATA,
                    count_funcs=count_funcs),
                TauFakeRateScaleFactors(
                    year=year,
                    datatype=datatype,
                    tree=tree,
                    tes_up=(self.args.syst_terms is not None and
                        (Systematics.TES_FAKE_TOTAL_UP in self.args.syst_terms or
                         Systematics.TES_FAKE_FINAL_UP in self.args.syst_terms)),
                    tes_down=(self.args.syst_terms is not None and
                        (Systematics.TES_FAKE_TOTAL_DOWN in self.args.syst_terms or
                         Systematics.TES_FAKE_FINAL_DOWN in self.args.syst_terms)),
                    passthrough=datatype in (datasets.DATA, datasets.EMBED),
                    count_funcs=count_funcs),
                HiggsPT(
                    year=year,
                    tree=tree,
                    passthrough=not is_signal or local,
                    count_funcs=count_funcs),
                TauTrackRecounting(
                    year=year,
                    use_ntup_value=tau_ntrack_recounted_use_ntup,
                    passthrough=local,
                    count_funcs=count_funcs),
                MCWeight(
                    datatype=datatype,
                    tree=tree,
                    passthrough=local or datatype == datasets.DATA,
                    count_funcs=count_funcs),
                EmbeddingIsolation(
                    tree=tree,
                    passthrough=(
                        local or year < 2012 or
                        datatype not in (datasets.EMBED, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                EmbeddingCorrections(
                    tree=tree,
                    year=year,
                    passthrough=(
                        local or
                        datatype not in (datasets.EMBED, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                EmbeddingTauSpinner(
                    year=year,
                    tree=tree,
                    passthrough=(
                        local or datatype not in (
                            datasets.EMBED, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                # put MET recalculation after tau selection but before tau-jet
                # overlap removal and jet selection because of the RefAntiTau
                # MET correction
                METRecalculation(
                    terms=syst_terms,
                    year=year,
                    tree=tree,
                    refantitau=not nominal_values,
                    verbose=verbose,
                    very_verbose=very_verbose,
                    count_funcs=count_funcs),
                TauJetOverlapRemoval(
                    count_funcs=count_funcs),
                JetPreselection(
                    count_funcs=count_funcs),
                NonIsolatedJet(
                    tree=tree,
                    count_funcs=count_funcs),
                JetSelection(
                    year=year,
                    count_funcs=count_funcs),
                RecoJetTrueTauMatching(
                    passthrough=datatype == datasets.DATA or local,
                    count_funcs=count_funcs),
                BCHCleaning(
                    tree=tree,
                    passthrough=year == 2011 or local,
                    datatype=datatype,
                    count_funcs=count_funcs),
                ClassifyInclusiveHiggsSample(
                    tree=tree,
                    passthrough=not is_inclusive_signal,
                    count_funcs=count_funcs),
            ])

            # set the event filters
            self.filters['event'] = event_filters

        # peek at first tree to determine which branches to exclude
        with root_open(self.files[0]) as test_file:
            test_tree = test_file.Get(self.metadata.treename)
            ignore_branches = test_tree.glob(
                hhbranches.REMOVE,
                exclude=hhbranches.KEEP)
            ignore_branches_output = test_tree.glob(
                hhbranches.REMOVE_OUTPUT,
                exclude=hhbranches.KEEP_OUTPUT)

        # initialize the TreeChain of all input files
        chain = TreeChain(
            self.metadata.treename,
            files=self.files,
            ignore_branches=ignore_branches,
            events=self.events,
            onfilechange=onfilechange,
            filters=event_filters,
            cache=True,
            cache_size=50000000,
            learn_entries=100)

        if local:
            copied = [
                'EventNumber',
            ]

            hh_buffer = TreeBuffer()
            buffer = TreeBuffer()
            for name, value in chain._buffer.items():
                if name.startswith('hh_'):
                    hh_buffer[name[3:]] = value
                elif name in copied:
                    buffer[name] = value
            outtree.set_buffer(
                hh_buffer,
                create_branches=False,
                visible=True)
            outtree.set_buffer(
                buffer,
                create_branches=True,
                visible=False)

        else:
            # additional decorations on existing objects
            if year > 2011 and datatype in (datasets.MC, datasets.MCEMBED):
                class Decorations(TreeModel):
                    jet_ispileup = stl.vector('bool')

                chain.set_buffer(Decorations(), create_branches=True)

            # include the branches in the input chain in the output tree
            # set branches to be removed in ignore_branches
            outtree.set_buffer(
                chain._buffer,
                ignore_branches=ignore_branches + ignore_branches_output,
                create_branches=True,
                ignore_duplicates=True,
                transfer_objects=True,
                visible=False)

        # define tree objects
        define_objects(chain, year)

        # create the MMC
        mmc = mass.MMC(year=year)

        # report which packages have been loaded
        externaltools.report()

        self.output.cd()

        # The main event loop
        # the event filters above are automatically run for each event and only
        # the surviving events are looped on
        for event in chain:

            if local and syst_terms is None and not redo_selection:
                outtree.Fill()
                continue

            # sort taus and jets in decreasing order by pT
            event.taus.sort(key=lambda tau: tau.pt, reverse=True)
            event.jets.sort(key=lambda jet: jet.pt, reverse=True)

            # tau1 is the leading tau
            # tau2 is the subleading tau
            tau1, tau2 = event.taus
            jets = list(event.jets)
            jet1, jet2, jet3 = None, None, None
            beta = None

            if len(jets) >= 2:
                jet1, jet2 = jets[:2]

                # determine boost of system
                # determine jet CoM frame
                beta = (jet1.fourvect + jet2.fourvect).BoostVector()
                tree.jet_beta.copy_from(beta)

                jet1.fourvect_boosted.copy_from(jet1.fourvect)
                jet2.fourvect_boosted.copy_from(jet2.fourvect)
                jet1.fourvect_boosted.Boost(beta * -1)
                jet2.fourvect_boosted.Boost(beta * -1)

                tau1.fourvect_boosted.copy_from(tau1.fourvect)
                tau2.fourvect_boosted.copy_from(tau2.fourvect)
                tau1.fourvect_boosted.Boost(beta * -1)
                tau2.fourvect_boosted.Boost(beta * -1)

                tau1.min_dr_jet = min(
                    tau1.fourvect.DeltaR(jet1.fourvect),
                    tau1.fourvect.DeltaR(jet2.fourvect))
                tau2.min_dr_jet = min(
                    tau2.fourvect.DeltaR(jet1.fourvect),
                    tau2.fourvect.DeltaR(jet2.fourvect))

                #sphericity, aplanarity = eventshapes.sphericity_aplanarity(
                #    [tau1.fourvect,
                #     tau2.fourvect,
                #     jet1.fourvect,
                #     jet2.fourvect])

                # sphericity
                #tree.sphericity = sphericity
                # aplanarity
                #tree.aplanarity = aplanarity

                #sphericity_boosted, aplanarity_boosted = eventshapes.sphericity_aplanarity(
                #    [tau1.fourvect_boosted,
                #     tau2.fourvect_boosted,
                #     jet1.fourvect_boosted,
                #     jet2.fourvect_boosted])

                # sphericity
                #tree.sphericity_boosted = sphericity_boosted
                # aplanarity
                #tree.aplanarity_boosted = aplanarity_boosted

                # tau centrality (degree to which they are between the two jets)
                tau1.centrality = eventshapes.eta_centrality(
                    tau1.fourvect.Eta(),
                    jet1.fourvect.Eta(),
                    jet2.fourvect.Eta())

                tau2.centrality = eventshapes.eta_centrality(
                    tau2.fourvect.Eta(),
                    jet1.fourvect.Eta(),
                    jet2.fourvect.Eta())

                # boosted tau centrality
                tau1.centrality_boosted = eventshapes.eta_centrality(
                    tau1.fourvect_boosted.Eta(),
                    jet1.fourvect_boosted.Eta(),
                    jet2.fourvect_boosted.Eta())

                tau2.centrality_boosted = eventshapes.eta_centrality(
                    tau2.fourvect_boosted.Eta(),
                    jet1.fourvect_boosted.Eta(),
                    jet2.fourvect_boosted.Eta())

                # 3rd leading jet
                if len(jets) >= 3:
                    jet3 = jets[2]
                    jet3.fourvect_boosted.copy_from(jet3.fourvect)
                    jet3.fourvect_boosted.Boost(beta * -1)

            elif len(jets) == 1:
                jet1 = jets[0]

                tau1.min_dr_jet = tau1.fourvect.DeltaR(jet1.fourvect)
                tau2.min_dr_jet = tau2.fourvect.DeltaR(jet1.fourvect)

                #sphericity, aplanarity = eventshapes.sphericity_aplanarity(
                #    [tau1.fourvect,
                #     tau2.fourvect,
                #     jet1.fourvect])

                # sphericity
                #tree.sphericity = sphericity
                # aplanarity
                #tree.aplanarity = aplanarity

            RecoJetBlock.set(tree, jet1, jet2, jet3, local=local)

            # mass of ditau + leading jet system
            if jet1 is not None:
                tree.mass_tau1_tau2_jet1 = (
                    tau1.fourvect + tau2.fourvect + jet1.fourvect).M()

            # full sphericity and aplanarity
            #sphericity_full, aplanarity_full = eventshapes.sphericity_aplanarity(
            #    [tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets])

            #tree.sphericity_full = sphericity_full
            #tree.aplanarity_full = aplanarity_full

            #####################################
            # number of tracks from PV minus taus
            #####################################
            ntrack_pv = 0
            ntrack_nontau_pv = 0
            for vxp in event.vertices:
                # primary vertex
                if vxp.type == 1:
                    ntrack_pv = vxp.nTracks
                    ntrack_nontau_pv = ntrack_pv - tau1.numTrack - tau2.numTrack
                    break
            tree.ntrack_pv = ntrack_pv
            tree.ntrack_nontau_pv = ntrack_nontau_pv

            #########################
            # MET variables
            #########################
            METx = event.MET.etx
            METy = event.MET.ety
            MET = event.MET.et
            MET_vect = Vector2(METx, METy)
            MET_4vect = LorentzVector()
            MET_4vect.SetPxPyPzE(METx, METy, 0., MET)
            MET_4vect_boosted = LorentzVector()
            MET_4vect_boosted.copy_from(MET_4vect)
            if beta is not None:
                MET_4vect_boosted.Boost(beta * -1)

            tree.MET_et = MET
            tree.MET_etx = METx
            tree.MET_ety = METy
            tree.MET_phi = event.MET.phi
            dPhi_tau1_tau2 = abs(tau1.fourvect.DeltaPhi(tau2.fourvect))
            dPhi_tau1_MET = abs(tau1.fourvect.DeltaPhi(MET_4vect))
            dPhi_tau2_MET = abs(tau2.fourvect.DeltaPhi(MET_4vect))
            tree.dPhi_tau1_tau2 = dPhi_tau1_tau2
            tree.dPhi_tau1_MET = dPhi_tau1_MET
            tree.dPhi_tau2_MET = dPhi_tau2_MET
            tree.dPhi_min_tau_MET = min(dPhi_tau1_MET, dPhi_tau2_MET)
            tree.MET_bisecting = is_MET_bisecting(
                dPhi_tau1_tau2,
                dPhi_tau1_MET,
                dPhi_tau2_MET)

            sumET = event.MET.sumet
            tree.MET_sumet = sumET
            if sumET != 0:
                tree.MET_sig = ((2. * MET / GeV) /
                    (utils.sign(sumET) * sqrt(abs(sumET / GeV))))
            else:
                tree.MET_sig = -1.

            tree.MET_centrality = eventshapes.phi_centrality(
                tau1.fourvect,
                tau2.fourvect,
                MET_vect)
            tree.MET_centrality_boosted = eventshapes.phi_centrality(
                tau1.fourvect_boosted,
                tau2.fourvect_boosted,
                MET_4vect_boosted)

            tree.number_of_good_vertices = len(event.vertices)

            ##########################
            # Jet and sum pt variables
            ##########################
            tree.numJets = len(event.jets)

            # sum pT with only the two leading jets
            tree.sum_pt = sum(
                [tau1.pt, tau2.pt] +
                [jet.pt for jet in jets[:2]])

            # sum pT with all selected jets
            tree.sum_pt_full = sum(
                [tau1.pt, tau2.pt] +
                [jet.pt for jet in jets])

            # vector sum pT with two leading jets and MET
            tree.vector_sum_pt = sum(
                [tau1.fourvect, tau2.fourvect] +
                [jet.fourvect for jet in jets[:2]] +
                [MET_4vect]).Pt()

            # vector sum pT with all selected jets and MET
            tree.vector_sum_pt_full = sum(
                [tau1.fourvect, tau2.fourvect] +
                [jet.fourvect for jet in jets] +
                [MET_4vect]).Pt()

            # resonance pT
            tree.resonance_pt = sum(
                [tau1.fourvect, tau2.fourvect, MET_4vect]).Pt()

            #############################
            # tau <-> vertex association
            #############################
            tree.tau_same_vertex = (
                tau1.privtx_x == tau2.privtx_x and
                tau1.privtx_y == tau2.privtx_y and
                tau1.privtx_z == tau2.privtx_z)

            tau1.vertex_prob = ROOT.TMath.Prob(
                tau1.privtx_chiSquared,
                int(tau1.privtx_numberDoF))

            tau2.vertex_prob = ROOT.TMath.Prob(
                tau2.privtx_chiSquared,
                int(tau2.privtx_numberDoF))

            ##########################
            # MMC Mass
            ##########################
            mmc_result = mmc.mass(
                tau1, tau2,
                METx, METy, sumET,
                njets=len(event.jets))

            for mmc_method, mmc_object in enumerate(mmc_objects):
                mmc_mass, mmc_resonance, mmc_met = mmc_result[mmc_method]
                if verbose:
                    log.info("MMC (method %d): %f" % (mmc_method, mmc_mass))

                mmc_object.mass = mmc_mass
                mmc_object.MET_et = mmc_met.Mod()
                mmc_object.MET_etx = mmc_met.X()
                mmc_object.MET_ety = mmc_met.Y()
                mmc_object.MET_phi = math.pi - mmc_met.Phi()
                if mmc_mass > 0:
                    FourMomentum.set(mmc_object.resonance, mmc_resonance)

            ############################
            # collinear and visible mass
            ############################
            vis_mass, collin_mass, tau1_x, tau2_x = mass.collinearmass(
                tau1, tau2, METx, METy)

            tree.mass_vis_tau1_tau2 = vis_mass
            tree.mass_collinear_tau1_tau2 = collin_mass
            tau1.collinear_momentum_fraction = tau1_x
            tau2.collinear_momentum_fraction = tau2_x

            ###########################
            # Match jets to VBF partons
            ###########################
            #if datatype == datasets.MC and 'VBF' in dsname and year == 2011:
            #    # get partons (already sorted by eta in hepmc) FIXME!!!
            #    parton1, parton2 = hepmc.get_VBF_partons(event)
            #    tree.mass_true_quark1_quark2 = (parton1.fourvect + parton2.fourvect).M()
            #    # order here needs to be revised since jets are no longer
            #    # sorted by eta but instead by pT
            #    PartonBlock.set(tree, parton1, parton2)
            #    if len(jets) >= 2:
            #        jet1, jet2 = jets[:2]
            #        for i, jet in zip((1, 2), (jet1, jet2)):
            #            for parton in (parton1, parton2):
            #                if utils.dR(jet.eta, jet.phi, parton.eta, parton.phi) < .8:
            #                    setattr(tree, 'jet%i_matched' % i, True)

            # Fill the tau block
            # This must come after the RecoJetBlock is filled since
            # that sets the jet_beta for boosting the taus
            RecoTauBlock.set(event, tree, datatype, tau1, tau2, local=local)
            if datatype != datasets.DATA:
                TrueTauBlock.set(tree, tau1, tau2)

            # fill the output tree
            outtree.Fill(reset=True)

        externaltools.report()

        # flush any baskets remaining in memory to disk
        self.output.cd()
        outtree.FlushBaskets()
        outtree.Write()

        if local:
            if datatype == datasets.DATA:
                xml_string = ROOT.TObjString(merged_grl.str())
                xml_string.Write('lumi')
            merged_cutflow.Write()
예제 #36
0
파일: test_tree.py 프로젝트: suvayu/rootpy
def test_chain_draw_hist_init_first():

    hist = Hist(100, 0, 1)
    chain = TreeChain("tree", FILE_PATHS)
    chain.draw("a_x", hist=hist)
    assert_equal(hist.Integral() > 0, True)
예제 #37
0
        # fill the tree
        for i in range(100):
            tree.x = gauss(.5, 1.)
            tree.y = gauss(.3, 2.)
            tree.z = gauss(13., 42.)
            tree.i = i
            tree.fill()
        tree.write()
"""
This section below takes the example trees and copies it while overwriting a
branch with new values.
"""

# first define the chain of trees
chain = TreeChain(name="test", files=fnames)

# Now we want to copy the tree above into a new file while overwriting a branch
# First create a new file to save the new tree in:
f_copy = root_open("test_copy.root", "recreate")

# You may not know the entire model of the original tree but only the branches
# you intend to overwrite, so I am not specifying the model=Event below as an
# example of how to deal with this in general:
tree_copy = Tree("test_copy")

# If the original tree was not handed to you through rootpy don't forget to:
# >>> from rootpy import asrootpy
# >>> tree = asrootpy(tree)

# Here we specify the buffer for the new tree to use. We use the same buffer as
예제 #38
0
    def test_chain_draw_hist_init_first(self):

        hist = Hist(100, 0, 1)
        chain = TreeChain('tree', self.file_paths)
        chain.draw('a_x', hist=hist)
        assert_equals(hist.Integral() > 0, True)
예제 #39
0
    def work(self):
        # get argument values
        local = self.args.local
        syst_terms = self.args.syst_terms
        datatype = self.metadata.datatype
        year = self.metadata.year
        verbose = self.args.student_verbose
        very_verbose = self.args.student_very_verbose
        redo_selection = self.args.redo_selection
        nominal_values = self.args.nominal_values

        # get the dataset name
        dsname = os.getenv('INPUT_DATASET_NAME', None)
        if dsname is None:
            # attempt to guess dsname from dirname
            if self.files:
                dsname = os.path.basename(os.path.dirname(self.files[0]))

        # is this a signal sample?
        # if so we will also keep some truth information in the output below
        is_signal = datatype == datasets.MC and (
            '_VBFH' in dsname or '_ggH' in dsname or '_ZH' in dsname
            or '_WH' in dsname or '_ttH' in dsname)
        log.info("DATASET: {0}".format(dsname))
        log.info("IS SIGNAL: {0}".format(is_signal))

        # is this an inclusive signal sample for overlap studies?
        is_inclusive_signal = is_signal and '_inclusive' in dsname

        # is this a BCH-fixed sample? (temporary)
        is_bch_sample = 'r5470_r4540_p1344' in dsname
        if is_bch_sample:
            log.warning("this is a BCH-fixed r5470 sample")

        # onfilechange will contain a list of functions to be called as the
        # chain rolls over to each new file
        onfilechange = []
        count_funcs = {}

        if datatype != datasets.DATA:
            # count the weighted number of events
            if local:

                def mc_weight_count(event):
                    return event.hh_mc_weight
            else:

                def mc_weight_count(event):
                    return event.mc_event_weight

            count_funcs = {
                'mc_weight': mc_weight_count,
            }

        # three instances of the pileup reweighting tool are created to write
        # out the nominal, high and low pileup weights
        pileup_tool = None
        pileup_tool_high = None
        pileup_tool_low = None

        if local:
            # local means running on the skims, the output of this script
            # running on the grid
            if datatype == datasets.DATA:
                # merge the GRL fragments
                merged_grl = goodruns.GRL()

                def update_grl(student, grl, name, file, tree):
                    grl |= str(
                        file.Get('Lumi/%s' %
                                 student.metadata.treename).GetString())

                onfilechange.append((update_grl, (
                    self,
                    merged_grl,
                )))

            if datatype == datasets.DATA:
                merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D')
            else:
                merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D')

            def update_cutflow(student, cutflow, name, file, tree):
                # record a cut-flow
                year = student.metadata.year
                datatype = student.metadata.datatype
                cutflow[1].value += file.cutflow_event[1].value
                if datatype != datasets.DATA:
                    cutflow[2].value += file.cutflow_event_mc_weight[1].value

            onfilechange.append((update_cutflow, (
                self,
                merged_cutflow,
            )))

        else:
            # get pileup reweighting tool
            pileup_tool = get_pileup_reweighting_tool(year=year,
                                                      use_defaults=True)
            pileup_tool_high = get_pileup_reweighting_tool(year=year,
                                                           use_defaults=True,
                                                           systematic='high')
            pileup_tool_low = get_pileup_reweighting_tool(year=year,
                                                          use_defaults=True,
                                                          systematic='low')

            if datatype not in (datasets.EMBED, datasets.MCEMBED):
                # merge TrigConfTrees
                metadirname = '%sMeta' % self.metadata.treename
                trigconfchain = ROOT.TChain('%s/TrigConfTree' % metadirname)
                map(trigconfchain.Add, self.files)
                metadir = self.output.mkdir(metadirname)
                metadir.cd()
                trigconfchain.Merge(self.output, -1, 'fast keep')
                self.output.cd()

            if datatype == datasets.DATA:
                # merge GRL XML strings
                merged_grl = goodruns.GRL()
                for fname in self.files:
                    with root_open(fname) as f:
                        for key in f.Lumi.keys():
                            merged_grl |= goodruns.GRL(str(
                                key.ReadObj().GetString()),
                                                       from_string=True)
                lumi_dir = self.output.mkdir('Lumi')
                lumi_dir.cd()
                xml_string = ROOT.TObjString(merged_grl.str())
                xml_string.Write(self.metadata.treename)
                self.output.cd()

        self.output.cd()

        # create the output tree
        model = get_model(datatype,
                          dsname,
                          prefix=None if local else 'hh_',
                          is_inclusive_signal=is_inclusive_signal)
        log.info("Output Model:\n\n{0}\n\n".format(model))
        outtree = Tree(name=self.metadata.treename, model=model)

        if local:
            tree = outtree
        else:
            tree = outtree.define_object(name='tree', prefix='hh_')

        tree.define_object(name='tau', prefix='tau_')
        tree.define_object(name='tau1', prefix='tau1_')
        tree.define_object(name='tau2', prefix='tau2_')
        tree.define_object(name='truetau1', prefix='truetau1_')
        tree.define_object(name='truetau2', prefix='truetau2_')
        tree.define_object(name='jet1', prefix='jet1_')
        tree.define_object(name='jet2', prefix='jet2_')
        tree.define_object(name='jet3', prefix='jet3_')

        mmc_objects = [
            tree.define_object(name='mmc0', prefix='mmc0_'),
            tree.define_object(name='mmc1', prefix='mmc1_'),
            tree.define_object(name='mmc2', prefix='mmc2_'),
        ]

        for mmc_obj in mmc_objects:
            mmc_obj.define_object(name='resonance', prefix='resonance_')

        trigger_emulation = TauTriggerEmulation(year=year,
                                                passthrough=local
                                                or datatype != datasets.MC
                                                or year > 2011,
                                                count_funcs=count_funcs)

        if not trigger_emulation.passthrough:
            onfilechange.append((update_trigger_trees, (
                self,
                trigger_emulation,
            )))

        trigger_config = None

        if datatype not in (datasets.EMBED, datasets.MCEMBED):
            # trigger config tool to read trigger info in the ntuples
            trigger_config = get_trigger_config()
            # update the trigger config maps on every file change
            onfilechange.append((update_trigger_config, (trigger_config, )))

        # define the list of event filters
        if local and syst_terms is None and not redo_selection:
            event_filters = None
        else:
            tau_ntrack_recounted_use_ntup = False
            if year > 2011:
                # peek at first tree to determine if the extended number of
                # tracks is already stored
                with root_open(self.files[0]) as test_file:
                    test_tree = test_file.Get(self.metadata.treename)
                    tau_ntrack_recounted_use_ntup = ('tau_out_track_n_extended'
                                                     in test_tree)

            event_filters = EventFilterList([
                averageIntPerXingPatch(
                    passthrough=(local or year < 2012
                                 or datatype != datasets.MC),
                    count_funcs=count_funcs),
                PileupTemplates(
                    year=year,
                    passthrough=(local or is_bch_sample or datatype
                                 not in (datasets.MC, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                RandomSeed(datatype=datatype, count_funcs=count_funcs),
                RandomRunNumber(tree=tree,
                                datatype=datatype,
                                pileup_tool=pileup_tool,
                                passthrough=local,
                                count_funcs=count_funcs),
                PileupReweight(
                    year=year,
                    tool=pileup_tool,
                    tool_high=pileup_tool_high,
                    tool_low=pileup_tool_low,
                    tree=tree,
                    passthrough=(local
                                 or (datatype
                                     not in (datasets.MC, datasets.MCEMBED))),
                    count_funcs=count_funcs),
                TruthMatching(passthrough=datatype == datasets.DATA,
                              count_funcs=count_funcs),
                JetIsPileup(
                    passthrough=(local or year < 2012 or datatype
                                 not in (datasets.MC, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                HiggsPT(year=year,
                        tree=tree,
                        passthrough=not is_signal or local,
                        count_funcs=count_funcs),
                MCWeight(datatype=datatype,
                         tree=tree,
                         passthrough=local or datatype == datasets.DATA,
                         count_funcs=count_funcs),
                ClassifyInclusiveHiggsSample(
                    tree=tree,
                    passthrough=not is_inclusive_signal,
                    count_funcs=count_funcs),
            ])

            # set the event filters
            self.filters['event'] = event_filters

        # peek at first tree to determine which branches to exclude
        with root_open(self.files[0]) as test_file:
            test_tree = test_file.Get(self.metadata.treename)
            ignore_branches = test_tree.glob(hhbranches.REMOVE,
                                             exclude=hhbranches.KEEP)
            ignore_branches_output = test_tree.glob(
                hhbranches.REMOVE_OUTPUT, exclude=hhbranches.KEEP_OUTPUT)

        # initialize the TreeChain of all input files
        chain = TreeChain(self.metadata.treename,
                          files=self.files,
                          ignore_branches=ignore_branches,
                          events=self.events,
                          onfilechange=onfilechange,
                          filters=event_filters,
                          cache=True,
                          cache_size=50000000,
                          learn_entries=100)

        if local:
            copied = [
                'EventNumber',
            ]

            hh_buffer = TreeBuffer()
            buffer = TreeBuffer()
            for name, value in chain._buffer.items():
                if name.startswith('hh_'):
                    hh_buffer[name[3:]] = value
                elif name in copied:
                    buffer[name] = value
            outtree.set_buffer(hh_buffer, create_branches=False, visible=True)
            outtree.set_buffer(buffer, create_branches=True, visible=False)

        else:
            # additional decorations on existing objects
            if year > 2011 and datatype in (datasets.MC, datasets.MCEMBED):

                class Decorations(TreeModel):
                    jet_ispileup = stl.vector('bool')

                chain.set_buffer(Decorations(), create_branches=True)

            # include the branches in the input chain in the output tree
            # set branches to be removed in ignore_branches
            outtree.set_buffer(chain._buffer,
                               ignore_branches=ignore_branches +
                               ignore_branches_output,
                               create_branches=True,
                               ignore_duplicates=True,
                               transfer_objects=True,
                               visible=False)

        # define tree objects
        define_objects(chain, year)

        # create the MMC
        mmc = mass.MMC(year=year)

        # report which packages have been loaded
        externaltools.report()

        self.output.cd()

        # The main event loop
        # the event filters above are automatically run for each event and only
        # the surviving events are looped on
        for event in chain:

            if local and syst_terms is None and not redo_selection:
                outtree.Fill()
                continue

            # sort taus and jets in decreasing order by pT
            event.taus.sort(key=lambda tau: tau.pt, reverse=True)
            event.jets.sort(key=lambda jet: jet.pt, reverse=True)

            # tau1 is the leading tau
            # tau2 is the subleading tau
            taus = list(event.taus)
            if len(taus) >= 2:
                tau1, tau2 = taus[0], taus[1]
                jets = list(event.jets)
                jet1, jet2, jet3 = None, None, None
                beta = None

                if len(jets) >= 2:
                    jet1, jet2 = jets[:2]

                    # determine boost of system
                    # determine jet CoM frame
                    beta = (jet1.fourvect + jet2.fourvect).BoostVector()
                    tree.jet_beta.copy_from(beta)

                    jet1.fourvect_boosted.copy_from(jet1.fourvect)
                    jet2.fourvect_boosted.copy_from(jet2.fourvect)
                    jet1.fourvect_boosted.Boost(beta * -1)
                    jet2.fourvect_boosted.Boost(beta * -1)

                    tau1.fourvect_boosted.copy_from(tau1.fourvect)
                    tau2.fourvect_boosted.copy_from(tau2.fourvect)
                    tau1.fourvect_boosted.Boost(beta * -1)
                    tau2.fourvect_boosted.Boost(beta * -1)

                    tau1.min_dr_jet = min(tau1.fourvect.DeltaR(jet1.fourvect),
                                          tau1.fourvect.DeltaR(jet2.fourvect))
                    tau2.min_dr_jet = min(tau2.fourvect.DeltaR(jet1.fourvect),
                                          tau2.fourvect.DeltaR(jet2.fourvect))

                    # sphericity, aplanarity = eventshapes.sphericity_aplanarity(
                    #    [tau1.fourvect,
                    #     tau2.fourvect,
                    #     jet1.fourvect,
                    #     jet2.fourvect])

                    # sphericity
                    # tree.sphericity = sphericity
                    # aplanarity
                    # tree.aplanarity = aplanarity

                    # sphericity_boosted, aplanarity_boosted = eventshapes.sphericity_aplanarity(
                    #    [tau1.fourvect_boosted,
                    #     tau2.fourvect_boosted,
                    #     jet1.fourvect_boosted,
                    #     jet2.fourvect_boosted])

                    # sphericity
                    # tree.sphericity_boosted = sphericity_boosted
                    # aplanarity
                    # tree.aplanarity_boosted = aplanarity_boosted

                    # tau centrality (degree to which they are between the two jets)
                    tau1.centrality = eventshapes.eta_centrality(
                        tau1.fourvect.Eta(), jet1.fourvect.Eta(),
                        jet2.fourvect.Eta())

                    tau2.centrality = eventshapes.eta_centrality(
                        tau2.fourvect.Eta(), jet1.fourvect.Eta(),
                        jet2.fourvect.Eta())

                    # boosted tau centrality
                    tau1.centrality_boosted = eventshapes.eta_centrality(
                        tau1.fourvect_boosted.Eta(),
                        jet1.fourvect_boosted.Eta(),
                        jet2.fourvect_boosted.Eta())

                    tau2.centrality_boosted = eventshapes.eta_centrality(
                        tau2.fourvect_boosted.Eta(),
                        jet1.fourvect_boosted.Eta(),
                        jet2.fourvect_boosted.Eta())

                    # 3rd leading jet
                    if len(jets) >= 3:
                        jet3 = jets[2]
                        jet3.fourvect_boosted.copy_from(jet3.fourvect)
                        jet3.fourvect_boosted.Boost(beta * -1)

                elif len(jets) == 1:
                    jet1 = jets[0]

                    tau1.min_dr_jet = tau1.fourvect.DeltaR(jet1.fourvect)
                    tau2.min_dr_jet = tau2.fourvect.DeltaR(jet1.fourvect)

                    # sphericity, aplanarity = eventshapes.sphericity_aplanarity(
                    #    [tau1.fourvect,
                    #     tau2.fourvect,
                    #     jet1.fourvect])

                    # sphericity
                    # tree.sphericity = sphericity
                    # aplanarity
                    #tree.aplanarity = aplanarity

                    RecoJetBlock.set(tree, jet1, jet2, jet3, local=local)

                # mass of ditau + leading jet system
                if jet1 is not None:
                    tree.mass_tau1_tau2_jet1 = (tau1.fourvect + tau2.fourvect +
                                                jet1.fourvect).M()

                # full sphericity and aplanarity
                # sphericity_full, aplanarity_full = eventshapes.sphericity_aplanarity(
                #    [tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets])

                # tree.sphericity_full = sphericity_full
                # tree.aplanarity_full = aplanarity_full

                # ####################################
                # number of tracks from PV minus taus
                # ####################################
                ntrack_pv = 0
                ntrack_nontau_pv = 0
                for vxp in event.vertices:
                    # primary vertex
                    if vxp.type == 1:
                        ntrack_pv = vxp.nTracks
                        ntrack_nontau_pv = ntrack_pv - tau1.numTrack - tau2.numTrack
                        break
                tree.ntrack_pv = ntrack_pv
                tree.ntrack_nontau_pv = ntrack_nontau_pv

                # ########################
                # MET variables
                # ########################
                METx = event.MET.etx
                METy = event.MET.ety
                MET = event.MET.et
                MET_vect = Vector2(METx, METy)
                MET_4vect = LorentzVector()
                MET_4vect.SetPxPyPzE(METx, METy, 0., MET)
                MET_4vect_boosted = LorentzVector()
                MET_4vect_boosted.copy_from(MET_4vect)
                if beta is not None:
                    MET_4vect_boosted.Boost(beta * -1)

                tree.MET_et = MET
                tree.MET_etx = METx
                tree.MET_ety = METy
                tree.MET_phi = event.MET.phi
                dPhi_tau1_tau2 = abs(tau1.fourvect.DeltaPhi(tau2.fourvect))
                dPhi_tau1_MET = abs(tau1.fourvect.DeltaPhi(MET_4vect))
                dPhi_tau2_MET = abs(tau2.fourvect.DeltaPhi(MET_4vect))
                tree.dPhi_tau1_tau2 = dPhi_tau1_tau2
                tree.dPhi_tau1_MET = dPhi_tau1_MET
                tree.dPhi_tau2_MET = dPhi_tau2_MET
                tree.dPhi_min_tau_MET = min(dPhi_tau1_MET, dPhi_tau2_MET)
                tree.MET_bisecting = is_MET_bisecting(dPhi_tau1_tau2,
                                                      dPhi_tau1_MET,
                                                      dPhi_tau2_MET)

                sumET = event.MET.sumet
                tree.MET_sumet = sumET
                if sumET != 0:
                    tree.MET_sig = (
                        (2. * MET / GeV) /
                        (utils.sign(sumET) * sqrt(abs(sumET / GeV))))
                else:
                    tree.MET_sig = -1.

                tree.MET_centrality = eventshapes.phi_centrality(
                    tau1.fourvect, tau2.fourvect, MET_vect)
                tree.MET_centrality_boosted = eventshapes.phi_centrality(
                    tau1.fourvect_boosted, tau2.fourvect_boosted,
                    MET_4vect_boosted)

                tree.number_of_good_vertices = len(event.vertices)

                # #########################
                # Jet and sum pt variables
                # #########################
                tree.numJets = len(event.jets)

                # sum pT with only the two leading jets
                tree.sum_pt = sum([tau1.pt, tau2.pt] +
                                  [jet.pt for jet in jets[:2]])

                # sum pT with all selected jets
                tree.sum_pt_full = sum([tau1.pt, tau2.pt] +
                                       [jet.pt for jet in jets])

                # vector sum pT with two leading jets and MET
                tree.vector_sum_pt = sum([tau1.fourvect, tau2.fourvect] +
                                         [jet.fourvect for jet in jets[:2]] +
                                         [MET_4vect]).Pt()

                # vector sum pT with all selected jets and MET
                tree.vector_sum_pt_full = sum([tau1.fourvect, tau2.fourvect] +
                                              [jet.fourvect for jet in jets] +
                                              [MET_4vect]).Pt()

                # resonance pT
                tree.resonance_pt = sum(
                    [tau1.fourvect, tau2.fourvect, MET_4vect]).Pt()

                # ############################
                # tau <-> vertex association
                # ############################
                tree.tau_same_vertex = (tau1.privtx_x == tau2.privtx_x
                                        and tau1.privtx_y == tau2.privtx_y
                                        and tau1.privtx_z == tau2.privtx_z)

                tau1.vertex_prob = ROOT.TMath.Prob(tau1.privtx_chiSquared,
                                                   int(tau1.privtx_numberDoF))

                tau2.vertex_prob = ROOT.TMath.Prob(tau2.privtx_chiSquared,
                                                   int(tau2.privtx_numberDoF))

                # #########################
                # MMC Mass
                # #########################
                mmc_result = mmc.mass(tau1,
                                      tau2,
                                      METx,
                                      METy,
                                      sumET,
                                      njets=len(event.jets))

                for mmc_method, mmc_object in enumerate(mmc_objects):
                    mmc_mass, mmc_resonance, mmc_met = mmc_result[mmc_method]
                    if verbose:
                        log.info("MMC (method %d): %f" %
                                 (mmc_method, mmc_mass))

                    mmc_object.mass = mmc_mass
                    mmc_object.MET_et = mmc_met.Mod()
                    mmc_object.MET_etx = mmc_met.X()
                    mmc_object.MET_ety = mmc_met.Y()
                    mmc_object.MET_phi = math.pi - mmc_met.Phi()
                    if mmc_mass > 0:
                        FourMomentum.set(mmc_object.resonance, mmc_resonance)

                # ###########################
                # collinear and visible mass
                # ###########################
                vis_mass, collin_mass, tau1_x, tau2_x = mass.collinearmass(
                    tau1, tau2, METx, METy)

                tree.mass_vis_tau1_tau2 = vis_mass
                tree.mass_collinear_tau1_tau2 = collin_mass
                tau1.collinear_momentum_fraction = tau1_x
                tau2.collinear_momentum_fraction = tau2_x

                ###########################
                # Match jets to VBF partons
                ###########################
                #if datatype == datasets.MC and 'VBF' in dsname and year == 2011:
                #    # get partons (already sorted by eta in hepmc) FIXME!!!
                #    parton1, parton2 = hepmc.get_VBF_partons(event)
                #    tree.mass_true_quark1_quark2 = (parton1.fourvect + parton2.fourvect).M()
                #    # order here needs to be revised since jets are no longer
                #    # sorted by eta but instead by pT
                #    PartonBlock.set(tree, parton1, parton2)
                #    if len(jets) >= 2:
                #        jet1, jet2 = jets[:2]
                #        for i, jet in zip((1, 2), (jet1, jet2)):
                #            for parton in (parton1, parton2):
                #                if utils.dR(jet.eta, jet.phi, parton.eta, parton.phi) < .8:
                #                    setattr(tree, 'jet%i_matched' % i, True)

                # Fill the tau block
                # This must come after the RecoJetBlock is filled since
                # that sets the jet_beta for boosting the taus
                RecoTauBlock.set(event,
                                 tree,
                                 datatype,
                                 tau1,
                                 tau2,
                                 local=local)
                if datatype != datasets.DATA:
                    TrueTauBlock.set(tree, tau1, tau2)

            # fill the output tree
            outtree.Fill(reset=True)

        externaltools.report()

        # flush any baskets remaining in memory to disk
        self.output.cd()
        outtree.FlushBaskets()
        outtree.Write()

        if local:
            if datatype == datasets.DATA:
                xml_string = ROOT.TObjString(merged_grl.str())
                xml_string.Write('lumi')
            merged_cutflow.Write()
예제 #40
0
    TrueTaus(count_funcs=count_funcs),
    ClassifyDecay(count_funcs=count_funcs, tree=outtree),
    TrueJets(count_funcs=count_funcs),
])

files = [args.input]

# peek at first tree to determine which branches to exclude
with root_open(files[0]) as test_file:
    test_tree = test_file.Get(args.tree_name)
    ignore_branches = test_tree.glob(branches.REMOVE)

chain = TreeChain(args.tree_name,
                  files=files,
                  filters=event_filters,
                  cache=True,
                  cache_size=50000000,
                  learn_entries=100,
                  ignore_branches=ignore_branches)

define_objects(chain)

for event in chain:

    outtree.runnumber = event.RunNumber
    outtree.evtnumber = event.EventNumber
    outtree.weight = event.mc_event_weight
    # sort taus and jets in decreasing order by pT
    event.taus.sort(key=lambda tau: tau.decay.fourvect_vis.Pt(), reverse=True)
    event.jets.sort(key=lambda jet: jet.pt, reverse=True)