Esempio n. 1
0
    def process(self, df):
        output = self.accumulator.identity()
        dataset = df['dataset']

        ## construct weights ##
        wgts = processor.Weights(df.size)
        if self.data_type != 'data':
            wgts.add('genw', df['weight'])
            npv = df['trueInteractionNum']
            wgts.add('pileup', *(f(npv) for f in self.pucorrs))

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)

        weight = wgts.weight()
        ########################

        ak4jets = JaggedCandidateArray.candidatesfromcounts(
            df['akjet_ak4PFJetsCHS_p4'],
            px=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fX'].content,
            py=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fY'].content,
            pz=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fZ'].content,
            energy=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fT'].content,
            jetid=df['akjet_ak4PFJetsCHS_jetid'].content,
        )
        ak4jets = ak4jets[ak4jets.jetid & (ak4jets.pt > 30) &
                          (np.abs(ak4jets.eta) < 2.4)]

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            sumtkpt=df['pfjet_tkPtSum05'].content,
            pfiso=df['pfjet_pfIsolationNoPU05'].content,
            mintkdist=df['pfjet_pfcands_minTwoTkDist'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label, ndsa=ndsa)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))
        leptonjets.add_attributes(
            mucharged=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum != 0))))
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(ljdsamuSubset.sum() == 0))
        leptonjets = leptonjets[(leptonjets.nocosmic) & (leptonjets.pt > 30) &
                                (leptonjets.mintkdist < 50)]

        ## __ twoleptonjets__
        twoleptonjets = leptonjets.counts >= 2
        dileptonjets = leptonjets[twoleptonjets]
        ak4jets = ak4jets[twoleptonjets]
        wgt = weight[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ak4jets = ak4jets[ak4jets.pt > (lj0.pt.flatten())]
        ak4jetCounts = (ak4jets.counts > 0).astype(int)
        minpfiso = ((lj0.pfiso > lj1.pfiso).astype(int) * lj1.pfiso +
                    (lj0.pfiso < lj1.pfiso).astype(int) * lj0.pfiso).flatten()
        ljneutrality = (
            (lj0.isneutral & lj1.isneutral).astype(int) * 1 +
            (lj0.mucharged & lj1.mucharged).astype(int) * 2).flatten()

        ## channel def ##
        #### 2mu2e
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = singleMuljEvents & muljInLeading2Events

        output['chan-2mu2e'].fill(dataset=dataset,
                                  iso=minpfiso[channel_2mu2e],
                                  val=ak4jetCounts[channel_2mu2e],
                                  weight=wgt[channel_2mu2e])

        #### 4mu
        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = doubleMuljEvents & muljIsLeading2Events

        output['chan-4mu'].fill(dataset=dataset,
                                iso=minpfiso[channel_4mu],
                                val=ljneutrality[channel_4mu],
                                weight=wgt[channel_4mu])

        ###########

        return output
    def process(self, df):
        output = self.accumulator.identity()
        if df.size==0: return output

        dataset = df['dataset']

        ## construct weights ##
        wgts = processor.Weights(df.size)
        if self.data_type!='data':
            wgts.add('genw', df['weight'])
            npv = df['trueInteractionNum']
            wgts.add('pileup', *(f(npv) for f in self.pucorrs))

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)
        # ...bla bla, other weights goes here

        weight = wgts.weight()
        ########################

        ak4jets = JaggedCandidateArray.candidatesfromcounts(
            df['akjet_ak4PFJetsCHS_p4'],
            px=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fX'].content,
            py=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fY'].content,
            pz=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fZ'].content,
            energy=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fT'].content,
            hadfrac=df['akjet_ak4PFJetsCHS_hadronEnergyFraction'].content,
            jetid=df['akjet_ak4PFJetsCHS_jetid'].content,
            deepcsv=df['hftagscore_DeepCSV_b'].content,
        )
        deepcsv_tight = np.bitwise_and(ak4jets.deepcsv, 1<<2)==(1<<2)
        ak4jets.add_attributes(deepcsvTight=deepcsv_tight)
        ak4jets=ak4jets[ak4jets.jetid&(ak4jets.pt>20)&(np.abs(ak4jets.eta)<2.5)]

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            ncands=df['pfjet_pfcands_n'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype==3).sum()
        ndsa = (ljdautype==8).sum()
        isegammajet = (npfmu==0)&(ndsa==0)
        ispfmujet = (npfmu>=2)&(ndsa==0)
        isdsajet = ndsa>0
        label = isegammajet.astype(int)*1+ispfmujet.astype(int)*2+isdsajet.astype(int)*3
        leptonjets.add_attributes(label=label)
        nmu = ((ljdautype==3)|(ljdautype==8)).sum()
        leptonjets.add_attributes(ismutype=(nmu>=2), iseltype=(nmu==0))

        ljpfmuId = fromNestNestIndexArray(df['muon_selectors'], awkward.fromiter(df['pfjet_pfcand_pfmuonIdx']))
        leptonjets.add_attributes(nloose=(ljpfmuId&(1<<0)==(1<<0)).sum(),
                                  nmedium=(ljpfmuId&(1<<1)==(1<<1)).sum(),
                                  ntight=(ljpfmuId&(1<<3)==(1<<3)).sum(),
                                  nisoloose=(ljpfmuId&(1<<7)==(1<<7)).sum(),
                                  nisomedium=(ljpfmuId&(1<<8)==(1<<8)).sum(),
                                  nisotight=(ljpfmuId&(1<<9)==(1<<9)).sum(),
                                  )

        ## __ twoleptonjets__
        twoleptonjets = leptonjets.counts>=2
        dileptonjets = leptonjets[twoleptonjets]
        ak4jets = ak4jets[twoleptonjets]
        wgt = weight[twoleptonjets]

        if dileptonjets.size==0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum()==1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents&muljInLeading2Events).astype(int)*1

        doubleMuljEvents = dileptonjets.ismutype.sum()==2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents&muljIsLeading2Events).astype(int)*2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        cuts = [
            np.ones_like(wgt).astype(bool),                        # all
            (np.abs(lj0.p4.delta_phi(lj1.p4))>np.pi/2).flatten(),  # dphi > pi/2
            ak4jets.counts<4,                                      # N(jets) < 4
            ak4jets[(ak4jets.pt>30)&(np.abs(ak4jets.eta)<2.4)&ak4jets.deepcsvTight].counts==0, # N(tightB)==0
            (~channel_2mu2e.astype(bool)) | (channel_2mu2e.astype(bool)&(((lj0.iseltype)&(lj0.pt>40)) | ((lj1.iseltype)&(lj1.pt>40))).flatten() ), # EGMpt0>40
            ( (lj0.ismutype&(lj0.pt>40)) | ((~lj0.ismutype)&(lj1.ismutype&(lj1.pt>40))) ).flatten(), # Mupt0>40
            ( (~(channel_==2)) | (channel_==2)&((lj1.pt>30).flatten()) ), # Mupt1>30
        ]

        if self.region == 'CR':
            cuts[1] = ~cuts[1]

        totcut = np.logical_and.reduce(cuts)

        dileptonjets = dileptonjets[totcut]
        wgt = wgt[totcut]
        channel_ = channel_[totcut]

        ljpfmu = dileptonjets[dileptonjets.label==2]
        ljpfmuones = ljpfmu.pt.ones_like()
        if ljpfmu.size==0: return output

        output['nloose'].fill(dataset=dataset, cnt=ljpfmu.nloose.flatten(), weight=(wgt*ljpfmuones).flatten(), channel=(channel_*ljpfmuones).flatten())
        output['nmedium'].fill(dataset=dataset, cnt=ljpfmu.nmedium.flatten(), weight=(wgt*ljpfmuones).flatten(), channel=(channel_*ljpfmuones).flatten())
        output['ntight'].fill(dataset=dataset, cnt=ljpfmu.ntight.flatten(), weight=(wgt*ljpfmuones).flatten(), channel=(channel_*ljpfmuones).flatten())
        output['nisoloose'].fill(dataset=dataset, cnt=ljpfmu.nisoloose.flatten(), weight=(wgt*ljpfmuones).flatten(), channel=(channel_*ljpfmuones).flatten())
        output['nisomedium'].fill(dataset=dataset, cnt=ljpfmu.nisomedium.flatten(), weight=(wgt*ljpfmuones).flatten(), channel=(channel_*ljpfmuones).flatten())
        output['nisotight'].fill(dataset=dataset, cnt=ljpfmu.nisotight.flatten(), weight=(wgt*ljpfmuones).flatten(), channel=(channel_*ljpfmuones).flatten())

        return output
Esempio n. 3
0
    def process(self, df):
        output = self.accumulator.identity()
        if df.size == 0: return output

        dataset = df['dataset']
        ## construct weights ##
        wgts = processor.Weights(df.size)
        if self.data_type != 'data':
            wgts.add('genw', df['weight'])
            npv = df['trueInteractionNum']
            wgts.add('pileup', *(f(npv) for f in self.pucorrs))

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)
        # ...bla bla, other weights goes here

        weight = wgts.weight()
        ########################

        ak4jets = JaggedCandidateArray.candidatesfromcounts(
            df['akjet_ak4PFJetsCHS_p4'],
            px=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fX'].content,
            py=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fY'].content,
            pz=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fZ'].content,
            energy=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fT'].content,
            jetid=df['akjet_ak4PFJetsCHS_jetid'].content,
            muefrac=df['akjet_ak4PFJetsCHS_muonEnergyFraction'].content,
            chaemefrac=df['akjet_ak4PFJetsCHS_chaEmEnergyFraction'].content,
            emefrac=df['akjet_ak4PFJetsCHS_emEnergyFraction'].content,
            hadfrac=df['akjet_ak4PFJetsCHS_hadronEnergyFraction'].content,
            chahadfrac=df['akjet_ak4PFJetsCHS_chaHadEnergyFraction'].content,
            deepcsv=df['hftagscore_DeepCSV_b'].content,
        )
        deepcsv_tight = np.bitwise_and(ak4jets.deepcsv, 1 << 2) == (1 << 2)
        ak4jets.add_attributes(deepcsvTight=deepcsv_tight, )
        ak4jets = ak4jets[ak4jets.jetid & (ak4jets.pt > 30) &
                          (np.abs(ak4jets.eta) < 2.4)]

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            vx=df['pfjet_klmvtx.fCoordinates.fX'].content,
            vy=df['pfjet_klmvtx.fCoordinates.fY'].content,
            vz=df['pfjet_klmvtx.fCoordinates.fZ'].content,
            mintkdist=df['pfjet_pfcands_minTwoTkDist'].content,
        )
        leptonjets.add_attributes(vxy=np.hypot(leptonjets.vx, leptonjets.vy))
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))
        leptonjets.add_attributes(
            displaced=((leptonjets.vxy >= 5) |
                       (np.isnan(leptonjets.vxy) & leptonjets.ismutype)
                       ))  # non-vertex treated as displaced too
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(ljdsamuSubset.sum() == 0))

        leptonjets = leptonjets[(leptonjets.isneutral) & (leptonjets.nocosmic)
                                & (leptonjets.pt > 30) &
                                (leptonjets.mintkdist < 50)]

        # mask_ = ak4jets.match(leptonjets, deltaRCut=0.4)
        # ak4jets = ak4jets[~mask_]

        ## __ twoleptonjets__
        twoleptonjets = (leptonjets.counts >= 2) & (leptonjets.ismutype.sum()
                                                    >= 1)
        dileptonjets = leptonjets[twoleptonjets]
        ak4jets = ak4jets[twoleptonjets]
        wgt = weight[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents
                         & muljInLeading2Events).astype(int) * 1

        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents & muljIsLeading2Events).astype(int) * 2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        ak4jets = ak4jets[ak4jets.pt > (lj0.pt.flatten())]

        output['njets'].fill(dataset=dataset,
                             cnt=ak4jets.counts,
                             weight=wgt,
                             channel=channel_)
        if ak4jets.flatten().size != 0:
            ak4jets = ak4jets[(ak4jets.pt > 30) & (np.abs(ak4jets.eta) < 2.4) &
                              (ak4jets.deepcsvTight)]
        output['ntightb'].fill(dataset=dataset,
                               cnt=ak4jets.counts,
                               weight=wgt,
                               channel=channel_)

        return output
Esempio n. 4
0
    def process(self, df):
        output = self.accumulator.identity()
        dataset = df['dataset']

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            sumtkpt=df['pfjet_tkPtSum05'].content,
            pfiso=df['pfjet_pfIsolationNoPU05'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(ljdsamuSubset.sum() == 0))
        leptonjets = leptonjets[(leptonjets.nocosmic) & (leptonjets.isneutral)]

        ## __ twoleptonjets__
        twoleptonjets = leptonjets.counts >= 2
        dileptonjets = leptonjets[twoleptonjets]
        # wgt = weight[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents
                         & muljInLeading2Events).astype(int) * 1

        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents & muljIsLeading2Events).astype(int) * 2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        ljones = dileptonjets.pt.ones_like()
        output['sumpt'].fill(
            dataset=dataset,
            sumpt=dileptonjets.sumtkpt.flatten(),
        )
        output['pfiso'].fill(
            dataset=dataset,
            iso=dileptonjets.pfiso.flatten(),
        )

        return output
Esempio n. 5
0
    def process(self, df):
        output = self.accumulator.identity()
        if df.size == 0: return output

        dataset = df['dataset']
        ## construct weights ##
        wgts = processor.Weights(df.size)
        if self.data_type != 'data':
            wgts.add('genw', df['weight'])
            npv = df['trueInteractionNum']
            wgts.add('pileup', *(f(npv) for f in self.pucorrs))

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)
        # ...bla bla, other weights goes here

        weight = wgts.weight()
        ########################

        ak4jets = JaggedCandidateArray.candidatesfromcounts(
            df['akjet_ak4PFJetsCHS_p4'],
            px=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fX'].content,
            py=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fY'].content,
            pz=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fZ'].content,
            energy=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fT'].content,
            jetid=df['akjet_ak4PFJetsCHS_jetid'].content,
            deepcsv=df['hftagscore_DeepCSV_b'].content,
        )
        deepcsv_tight = np.bitwise_and(ak4jets.deepcsv, 1 << 2) == (1 << 2)
        ak4jets.add_attributes(deepcsvTight=deepcsv_tight)
        ak4jets = ak4jets[ak4jets.jetid & (ak4jets.pt > 20) &
                          (np.abs(ak4jets.eta) < 2.5)]

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            lxy=df['pfjet_klmvtx_lxy'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label, ndsa=ndsa)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))
        leptonjets.add_attributes(
            displaced=((np.abs(leptonjets.lxy) >= 5) |
                       (np.isnan(leptonjets.lxy) & leptonjets.ismutype)
                       ))  # non-vertex treated as displaced too
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(ljdsamuSubset.sum() == 0))

        leptonjets = leptonjets[(leptonjets.nocosmic) & (leptonjets.pt > 30)]

        ## __twoleptonjets__ AND >=1 displaced
        twoleptonjets = (leptonjets.counts >= 2) & (
            leptonjets.ismutype.sum() >= 1) & (leptonjets.displaced.sum() >= 1)
        dileptonjets = leptonjets[twoleptonjets]
        ak4jets = ak4jets[twoleptonjets]
        wgt = weight[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents
                         & muljInLeading2Events).astype(int) * 1

        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents & muljIsLeading2Events).astype(int) * 2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        cuts = [
            ((lj0.isneutral) & (lj1.isneutral)).flatten(),  # both 'neutral'
            (np.abs(lj0.p4.delta_phi(lj1.p4)) >
             np.pi / 2).flatten(),  # dphi > pi/2
            (~channel_2mu2e.astype(bool)) |
            (channel_2mu2e.astype(bool) &
             (((lj0.iseltype) & (lj0.pt > 60)) |
              ((lj1.iseltype) & (lj1.pt > 60))).flatten()),  # EGMpt0>60
            ak4jets.counts < 3,  # N(jets) < 4
            ak4jets[(ak4jets.pt > 30) & (np.abs(ak4jets.eta) < 2.4)
                    & ak4jets.deepcsvTight].counts == 0,  # N(tightB)==0
        ]

        if self.region == 'CR':
            cuts[1] = ~cuts[1]
        if self.enforceNeutral == False:
            cuts[0] = ~cuts[0]

        for i, c in enumerate(itertools.accumulate(cuts, np.logical_and)):
            output['count'].fill(dataset=dataset,
                                 cnt=np.ones_like(wgt[c]) * i,
                                 weight=wgt[c],
                                 channel=channel_[c])

        return output
Esempio n. 6
0
    def process(self, df):
        output = self.accumulator.identity()
        dataset = df['dataset']

        ## construct weights ##
        wgts = processor.Weights(df.size)
        if self.data_type != 'data':
            wgts.add('genw', df['weight'])
            npv = df['trueInteractionNum']
            wgts.add('pileup', *(f(npv) for f in self.pucorrs))

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)
        # ...bla bla, other weights goes here

        weight = wgts.weight()
        ########################

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            sumtkpt=df['pfjet_tkPtSum05'].content,
            pfiso=df['pfjet_pfIsolationNoPU05'].content,
            isodbeta=df['pfjet_pfiso'].content,
            mintkdist=df['pfjet_pfcands_minTwoTkDist'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))
        leptonjets.add_attributes(
            mucharged=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum != 0))))
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(ljdsamuSubset.sum() == 0))
        leptonjets = leptonjets[(leptonjets.nocosmic) & (leptonjets.pt > 30) &
                                (leptonjets.mintkdist < 50)]

        ## __ twoleptonjets__
        twoleptonjets = leptonjets.counts >= 2
        dileptonjets = leptonjets[twoleptonjets]
        wgt = weight[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents
                         & muljInLeading2Events).astype(int) * 1

        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents & muljIsLeading2Events).astype(int) * 2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        # isControl = (np.abs(lj0.p4.delta_phi(lj1.p4))<np.pi/2).flatten()
        # if self.data_type!='data':
        #     dileptonjets = dileptonjets[isControl]
        #     channel_ = channel_[isControl]
        #     wgt = wgt[isControl]

        mask_ = (lj0.isneutral & lj1.isneutral).flatten()
        if self.bothNeutral is False:
            mask_ = (lj0.mucharged & lj1.mucharged).flatten()
            # mask_ = ((channel_==2)&((~lj0.isneutral&(~lj1.isneutral)).flatten())) | ((channel_==1)&mask_)

        channel_ = channel_[mask_]
        wgt = wgt[mask_]
        dileptonjets = dileptonjets[mask_]

        minpfiso = (lj0.pfiso > lj1.pfiso).astype(int) * lj1.pfiso + (
            lj0.pfiso < lj1.pfiso).astype(int) * lj0.pfiso
        output['minpfiso'].fill(dataset=dataset,
                                iso=minpfiso[mask_].flatten(),
                                channel=channel_,
                                weight=wgt)
        maxpfiso = (lj0.pfiso > lj1.pfiso).astype(int) * lj0.pfiso + (
            lj0.pfiso < lj1.pfiso).astype(int) * lj1.pfiso
        output['maxpfiso'].fill(dataset=dataset,
                                iso=maxpfiso[mask_].flatten(),
                                channel=channel_,
                                weight=wgt)

        output['lj0pfiso'].fill(dataset=dataset,
                                iso=lj0.pfiso[mask_].flatten(),
                                channel=channel_,
                                weight=wgt)

        ljones = dileptonjets.pt.ones_like()
        output['sumpt'].fill(dataset=dataset,
                             sumpt=dileptonjets.sumtkpt.flatten(),
                             channel=(channel_ * ljones).flatten(),
                             weight=(wgt * ljones).flatten())
        output['pfiso'].fill(dataset=dataset,
                             iso=dileptonjets.pfiso.flatten(),
                             channel=(channel_ * ljones).flatten(),
                             weight=(wgt * ljones).flatten())
        output['isodbeta'].fill(dataset=dataset,
                                iso=dileptonjets.isodbeta.flatten(),
                                channel=(channel_ * ljones).flatten(),
                                weight=(wgt * ljones).flatten())

        return output
Esempio n. 7
0
    def process(self, df):
        output = self.accumulator.identity()
        if df.size == 0: return output

        dataset = df['dataset']
        run = df['run']
        lumi = df['lumi']
        event = df['event']

        ## construct weights ##
        wgts = processor.Weights(df.size)
        if self.data_type != 'data':
            wgts.add('genw', df['weight'])
            npv = df['trueInteractionNum']
            wgts.add('pileup', *(f(npv) for f in self.pucorrs))

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)
        # ...bla bla, other weights goes here

        weight = wgts.weight()
        ########################

        ak4jets = JaggedCandidateArray.candidatesfromcounts(
            df['akjet_ak4PFJetsCHS_p4'],
            px=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fX'].content,
            py=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fY'].content,
            pz=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fZ'].content,
            energy=df['akjet_ak4PFJetsCHS_p4.fCoordinates.fT'].content,
            jetid=df['akjet_ak4PFJetsCHS_jetid'].content,
            deepcsv=df['hftagscore_DeepCSV_b'].content,
        )
        deepcsv_tight = np.bitwise_and(ak4jets.deepcsv, 1 << 2) == (1 << 2)
        ak4jets.add_attributes(deepcsvTight=deepcsv_tight)
        ak4jets = ak4jets[ak4jets.jetid & (ak4jets.pt > 20) &
                          (np.abs(ak4jets.eta) < 2.5)]

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            vx=df['pfjet_klmvtx.fCoordinates.fX'].content,
            vy=df['pfjet_klmvtx.fCoordinates.fY'].content,
            vz=df['pfjet_klmvtx.fCoordinates.fZ'].content,
            mintkdist=df['pfjet_pfcands_minTwoTkDist'].content,
        )
        ## attribute: `label, ismutype, iseltype`
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label, ndsa=ndsa)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ## attribute: `isneutral`
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))
        ## attribute: `displaced`
        leptonjets.add_attributes(vxy=np.hypot(leptonjets.vx, leptonjets.vy))
        leptonjets.add_attributes(
            displaced=((leptonjets.vxy >= 5) |
                       (np.isnan(leptonjets.vxy) & leptonjets.ismutype)
                       ))  # non-vertex treated as displaced too
        ## attribute: `nocosmic`
        ljdsamuFoundOppo = fromNestNestIndexArray(
            df['dsamuon_hasOppositeMuon'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        dtcscTime = fromNestNestIndexArray(
            df['dsamuon_timeDiffDTCSC'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        rpcTime = fromNestNestIndexArray(
            df['dsamuon_timeDiffRPC'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        if len(dtcscTime.flatten().flatten()):
            dtcscTime = dtcscTime[ljdsamuFoundOppo]
            rpcTime = rpcTime[ljdsamuFoundOppo]
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(((dtcscTime < -20).sum() == 0)
                                            & ((rpcTime < -7.5).sum() == 0)
                                            & (ljdsamuSubset.sum() == 0)))

        leptonjets = leptonjets[(leptonjets.nocosmic) & (leptonjets.pt > 30) &
                                (leptonjets.mintkdist < 30)]

        ## __twoleptonjets__
        twoleptonjets = (leptonjets.counts >= 2) & (
            leptonjets.ismutype.sum() >= 1)  #&(leptonjets.displaced.sum()>=1)
        dileptonjets = leptonjets[twoleptonjets]
        ak4jets = ak4jets[twoleptonjets]
        wgt = weight[twoleptonjets]

        run = run[twoleptonjets]
        lumi = lumi[twoleptonjets]
        event = event[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents
                         & muljInLeading2Events).astype(int) * 1

        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents & muljIsLeading2Events).astype(int) * 2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        # cuts = [
        #     wgt.astype(bool),
        #     ((lj0.isneutral)&(lj1.isneutral)).flatten(), # both 'neutral'
        #     (np.abs(lj0.p4.delta_phi(lj1.p4)) > np.pi / 2).flatten(),  # dphi > pi/2
        #     (~channel_2mu2e.astype(bool)) | (channel_2mu2e.astype(bool)&(((lj0.iseltype)&(lj0.pt>60)) | ((lj1.iseltype)&(lj1.pt>60))).flatten() ), # EGMpt0>60
        #     ak4jets.counts<4,                                      # N(jets) < 4
        #     ak4jets[(ak4jets.pt > 30) & (np.abs(ak4jets.eta) < 2.4) & ak4jets.deepcsvTight].counts == 0,  # N(tightB)==0
        # ]
        # if self.data_type == 'data':
        #     cuts[1] = ~cuts[1]

        cuts = [
            wgt.astype(bool),
            ((~lj0.isneutral) & (~lj1.isneutral)).flatten(),
            (np.abs(lj0.p4.delta_phi(lj1.p4)) > 2.98).flatten(),
        ]

        totcut = np.logical_and.reduce(cuts)

        output['run_1'] += processor.column_accumulator(run[totcut
                                                            & (channel_ == 1)])
        output['lumi_1'] += processor.column_accumulator(
            lumi[totcut & (channel_ == 1)])
        output['event_1'] += processor.column_accumulator(
            event[totcut & (channel_ == 1)])
        output['run_2'] += processor.column_accumulator(run[totcut
                                                            & (channel_ == 2)])
        output['lumi_2'] += processor.column_accumulator(
            lumi[totcut & (channel_ == 2)])
        output['event_2'] += processor.column_accumulator(
            event[totcut & (channel_ == 2)])

        if self.data_type == 'data':
            era = np.ones_like(run) * list('ABCD').index(dataset)
            output['era_1'] += processor.column_accumulator(
                era[totcut & (channel_ == 1)])
            output['era_2'] += processor.column_accumulator(
                era[totcut & (channel_ == 2)])

        return output
    def process(self, df):
        output = self.accumulator.identity()
        dataset = df['dataset']

        ## construct weights ##
        wgts = processor.Weights(df.size)

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)

        weight = wgts.weight()
        ########################

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            pfiso=df['pfjet_pfIsolationNoPU05'].content,
            mintkdist=df['pfjet_pfcands_minTwoTkDist'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label, ndsa=ndsa)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))
        ## attribute: `nocosmic`
        ljdsamuFoundOppo = fromNestNestIndexArray(
            df['dsamuon_hasOppositeMuon'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        dtcscTime = fromNestNestIndexArray(
            df['dsamuon_timeDiffDTCSC'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        rpcTime = fromNestNestIndexArray(
            df['dsamuon_timeDiffRPC'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        if len(dtcscTime.flatten().flatten()):
            dtcscTime = dtcscTime[ljdsamuFoundOppo]
            rpcTime = rpcTime[ljdsamuFoundOppo]
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(((dtcscTime < -20).sum() == 0)
                                            & ((rpcTime < -7.5).sum() == 0)
                                            & (ljdsamuSubset.sum() == 0)))
        leptonjets = leptonjets[(leptonjets.nocosmic) & (leptonjets.pt > 30) &
                                (leptonjets.mintkdist < 30)]

        ## __ twoleptonjets__
        twoleptonjets = leptonjets.counts >= 2
        dileptonjets = leptonjets[twoleptonjets]
        wgt = weight[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents
                         & muljInLeading2Events).astype(int) * 1

        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = doubleMuljEvents & muljIsLeading2Events
        ###########

        dphi = np.abs(lj0.p4.delta_phi(lj1.p4))

        if self.category == '00':
            channel_4mu = channel_4mu & (
                (lj0.isneutral & lj1.isneutral).flatten())
        elif self.category == '01':
            channel_4mu = channel_4mu & ((lj0.isneutral &
                                          (~lj1.isneutral)).flatten())
        elif self.category == '10':
            channel_4mu = channel_4mu & ((
                (~lj0.isneutral) & lj1.isneutral).flatten())
        elif self.category == '11':
            channel_4mu = channel_4mu & (((~lj0.isneutral) &
                                          (~lj1.isneutral)).flatten())

        dphi = dphi[channel_4mu]
        wgt = wgt[channel_4mu]

        output['dphi'] += processor.column_accumulator(
            dphi[wgt.astype(bool)].flatten())

        return output
Esempio n. 9
0
    def process(self, df):
        output = self.accumulator.identity()
        dataset = df['dataset']

        ## construct weights ##
        wgts = processor.Weights(df.size)
        if self.data_type!='data':
            wgts.add('genw', df['weight'])
            npv = df['trueInteractionNum']
            wgts.add('pileup', *(f(npv) for f in self.pucorrs))

        triggermask = np.logical_or.reduce([df[t] for t in Triggers])
        wgts.add('trigger', triggermask)
        cosmicpairmask = df['cosmicveto_result']
        wgts.add('cosmicveto', cosmicpairmask)
        pvmask = df['metfilters_PrimaryVertexFilter']
        wgts.add('primaryvtx', pvmask)
        # ...bla bla, other weights goes here

        weight = wgts.weight()
        ########################

        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            mintkdist=df['pfjet_pfcands_minTwoTkDist'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype==3).sum()
        ndsa = (ljdautype==8).sum()
        isegammajet = (npfmu==0)&(ndsa==0)
        ispfmujet = (npfmu>=2)&(ndsa==0)
        isdsajet = ndsa>0
        label = isegammajet.astype(int)*1+ispfmujet.astype(int)*2+isdsajet.astype(int)*3
        leptonjets.add_attributes(label=label)
        nmu = ((ljdautype==3)|(ljdautype==8)).sum()
        leptonjets.add_attributes(ismutype=(nmu>=2), iseltype=(nmu==0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(isneutral=(leptonjets.iseltype | (leptonjets.ismutype&(leptonjets.qsum==0))))
        leptonjets.add_attributes(mucharged=(leptonjets.iseltype | (leptonjets.ismutype&(leptonjets.qsum!=0))))
        ljdsamuSubset = fromNestNestIndexArray(df['dsamuon_isSubsetFilteredCosmic1Leg'], awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        leptonjets.add_attributes(nocosmic=(ljdsamuSubset.sum()==0))
        leptonjets = leptonjets[(leptonjets.nocosmic)&(leptonjets.pt>30)&(leptonjets.mintkdist<50)]

        ## __ twoleptonjets__
        twoleptonjets = leptonjets.counts>=2
        dileptonjets = leptonjets[twoleptonjets]
        wgt = weight[twoleptonjets]

        if dileptonjets.size==0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum()==1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents&muljInLeading2Events).astype(int)*1

        doubleMuljEvents = dileptonjets.ismutype.sum()==2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents&muljIsLeading2Events).astype(int)*2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        ljBothNeutral = ((lj0.isneutral)&(lj1.isneutral)).flatten()
        ljBothCharged = ((~lj0.isneutral).sum()+(~lj1.isneutral).sum())==((lj0.ismutype).sum()+(lj1.ismutype).sum())
        lj4mu0cha = (channel_==2)&((lj0.mucharged&lj1.isneutral).flatten())
        lj4mu1cha = (channel_==2)&((lj0.isneutral&(lj1.mucharged)).flatten())
        lj4mu0or1cha = (channel_==2)&((lj0.isneutral&(lj1.mucharged)).flatten()|(lj0.mucharged&lj1.isneutral).flatten())
        output['dphi-neu'].fill(dataset=dataset, dphi=np.abs(lj0.p4.delta_phi(lj1.p4)[ljBothNeutral].flatten()), channel=channel_[ljBothNeutral], weight=wgt[ljBothNeutral])
        output['dphi-cha'].fill(dataset=dataset, dphi=np.abs(lj0.p4.delta_phi(lj1.p4)[ljBothCharged].flatten()), channel=channel_[ljBothCharged], weight=wgt[ljBothCharged])
        output['dphi-0mucha'].fill(dataset=dataset, dphi=np.abs(lj0.p4.delta_phi(lj1.p4)[lj4mu0cha].flatten()), channel=channel_[lj4mu0cha], weight=wgt[lj4mu0cha])
        output['dphi-1mucha'].fill(dataset=dataset, dphi=np.abs(lj0.p4.delta_phi(lj1.p4)[lj4mu1cha].flatten()), channel=channel_[lj4mu1cha], weight=wgt[lj4mu1cha])
        output['dphi-01mucha'].fill(dataset=dataset, dphi=np.abs(lj0.p4.delta_phi(lj1.p4)[lj4mu0or1cha].flatten()), channel=channel_[lj4mu0or1cha], weight=wgt[lj4mu0or1cha])

        return output
Esempio n. 10
0
    def process(self, df):
        output = self.accumulator.identity()
        dataset = df['dataset']

        parallelpairs = df['cosmicveto_parallelpairs']
        leptonjets = JaggedCandidateArray.candidatesfromcounts(
            df['pfjet_p4'],
            px=df['pfjet_p4.fCoordinates.fX'].content,
            py=df['pfjet_p4.fCoordinates.fY'].content,
            pz=df['pfjet_p4.fCoordinates.fZ'].content,
            energy=df['pfjet_p4.fCoordinates.fT'].content,
            mintkdist=df['pfjet_pfcands_minTwoTkDist'].content,
        )
        ljdautype = awkward.fromiter(df['pfjet_pfcand_type'])
        npfmu = (ljdautype == 3).sum()
        ndsa = (ljdautype == 8).sum()
        isegammajet = (npfmu == 0) & (ndsa == 0)
        ispfmujet = (npfmu >= 2) & (ndsa == 0)
        isdsajet = ndsa > 0
        label = isegammajet.astype(int) * 1 + ispfmujet.astype(
            int) * 2 + isdsajet.astype(int) * 3
        leptonjets.add_attributes(label=label)
        nmu = ((ljdautype == 3) | (ljdautype == 8)).sum()
        leptonjets.add_attributes(ismutype=(nmu >= 2), iseltype=(nmu == 0))
        ljdaucharge = awkward.fromiter(df['pfjet_pfcand_charge']).sum()
        leptonjets.add_attributes(qsum=ljdaucharge)
        leptonjets.add_attributes(
            isneutral=(leptonjets.iseltype
                       | (leptonjets.ismutype & (leptonjets.qsum == 0))))

        ljdsamuFoundOppo = fromNestNestIndexArray(
            df['dsamuon_hasOppositeMuon'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))
        dtcscTime = fromNestNestIndexArray(
            df['dsamuon_timeDiffDTCSC'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))[ljdsamuFoundOppo]
        rpcTime = fromNestNestIndexArray(
            df['dsamuon_timeDiffRPC'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))[ljdsamuFoundOppo]
        ljdsamuSubset = fromNestNestIndexArray(
            df['dsamuon_isSubsetFilteredCosmic1Leg'],
            awkward.fromiter(df['pfjet_pfcand_dsamuonIdx']))

        leptonjets_ = leptonjets[leptonjets.isneutral]
        ljdsamuFoundOppo = ljdsamuFoundOppo[leptonjets.isneutral]
        dtcscTime = dtcscTime[leptonjets.isneutral]
        rpcTime = rpcTime[leptonjets.isneutral]
        ljdsamuSubset = ljdsamuSubset[leptonjets.isneutral]

        ## __ twoleptonjets__
        twoleptonjets = leptonjets_.counts >= 2
        dileptonjets = leptonjets_[twoleptonjets]
        parallelpairs = parallelpairs[twoleptonjets]
        ljdsamuFoundOppo = ljdsamuFoundOppo[twoleptonjets]
        dtcscTime = dtcscTime[twoleptonjets]
        rpcTime = rpcTime[twoleptonjets]
        ljdsamuSubset = ljdsamuSubset[twoleptonjets]

        if dileptonjets.size == 0: return output
        lj0 = dileptonjets[dileptonjets.pt.argmax()]
        lj1 = dileptonjets[dileptonjets.pt.argsort()[:, 1:2]]

        ## channel def ##
        singleMuljEvents = dileptonjets.ismutype.sum() == 1
        muljInLeading2Events = (lj0.ismutype | lj1.ismutype).flatten()
        channel_2mu2e = (singleMuljEvents
                         & muljInLeading2Events).astype(int) * 1

        doubleMuljEvents = dileptonjets.ismutype.sum() == 2
        muljIsLeading2Events = (lj0.ismutype & lj1.ismutype).flatten()
        channel_4mu = (doubleMuljEvents & muljIsLeading2Events).astype(int) * 2

        channel_ = channel_2mu2e + channel_4mu
        ###########

        output['npairs'].fill(dataset=dataset, num=parallelpairs[channel_ > 0])
        output['ljdsaOppo'].fill(
            dataset=dataset,
            frac=ljdsamuFoundOppo[channel_ > 0].flatten().flatten())
        output['dtcscTime'].fill(dataset=dataset,
                                 t=dtcscTime[channel_ > 0].flatten().flatten())
        output['rpcTime'].fill(dataset=dataset,
                               t=rpcTime[channel_ > 0].flatten().flatten())
        output['ljdsaSubset'].fill(
            dataset=dataset,
            frac=ljdsamuSubset[channel_ > 0].flatten().flatten())
        output['tkDist'].fill(
            dataset=dataset,
            dist=dileptonjets.mintkdist[channel_ > 0].flatten())

        return output