CoffeaTeam / coffea

Basic tools and wrappers for enabling not-too-alien syntax when running columnar Collider HEP analysis.
https://coffeateam.github.io/coffea/
BSD 3-Clause "New" or "Revised" License
132 stars 126 forks source link

Adding the CorrT1METJet class to nanoaod.py #1119

Open 9GaoHong opened 4 months ago

9GaoHong commented 4 months ago

Is your feature request related to a problem? Please describe. When propagating JEC to Type1 corrections for raw MET, it is necessary to consider the contributions of both jet and soft jets, but currently the script does not support the use of soft jets branch CorrT1METJet.

Describe the solution you'd like Fix the /coffea/nanoevents/methods/nanoaod.py, add these code for CorrT1METJet:

@awkward.mixin_class(behavior)
class CorrT1METJet(vector.PtEtaPhiMLorentzVector, base.NanoCollection, base.Systematic):
    """NanoAOD narrow radius jet object"""
    @property
    def pt(self):
        return self["rawPt"]

    @property
    def rawFactor(self):
        self["rawFactor"] = awkward.zeros_like(self["rawPt"])
        return self["rawFactor"]

    @property
    def mass(self):
        self["mass"] = awkward.zeros_like(self["rawPt"])
        return self["mass"]

    @property
    def chEmEF(self):
        self["chEmEF"] = awkward.zeros_like(self["rawPt"])
        return self["chEmEF"]

    @property
    def neEmEF(self):
        self["neEmEF"] = awkward.zeros_like(self["rawPt"])
        return self["neEmEF"]

_set_repr_name("CorrT1METJet")

Fix the /coffea/nanoevents/schemas/nanoaod.py, add this key in 'mixins' dictionary:

"CorrT1METJet":"CorrT1METJet",

lucalavezzo commented 3 months ago

+1 having the same problem. @9GaoHong do you have an working implementation of the correction to MET from CorrT1METJets? I can't seem to be able to propagate the CorrT1METJets through the CorrectedMETFactory...

9GaoHong commented 3 months ago

@lucalavezzo Yes, I solved the problem by changing the nanoaod.py of coffee framework in local with the above code, and you can refer to the following code for specific instructions on how to use CorrT1METJet:

        rho = event.fixedGridRhoFastjetAll
        cache = event.caches[0]
        if is_data: 
            softjet_gen_pt = None
        else:
            softjet_gen_pt = find_best_match(event.CorrT1METJet,event.GenJet)
        softjets_shift_L123 = self._jmeu.corrected_jets_L123(event.CorrT1METJet, rho, cache, softjet_gen_pt)
        softjets_shift_L1 = self._jmeu.corrected_jets_L1(event.CorrT1METJet, rho, cache, softjet_gen_pt)
        jets_shift_L123 = self._jmeu.corrected_jets_L123(event.Jet, rho, cache)
        jets_shift_L1 = self._jmeu.corrected_jets_L1(event.Jet, rho, cache)

        jets_col_shift_L123 = ak.concatenate([jets_shift_L123, softjets_shift_L123],axis=1)
        jets_col_shift_L1 = ak.concatenate([jets_shift_L1, softjets_shift_L1],axis=1)

        raw_met = event.RawMET
        met_to_correct = event.MET
        met_to_correct["pt"] = raw_met.pt
        met_to_correct["phi"] = raw_met.phi
        jets = self._jmeu.corrected_jets_jer(event.Jet, event.fixedGridRhoFastjetAll, event.caches[0])
        met = self._jmeu.corrected_met(met_to_correct, jets_col_shift_L123, jets_col_shift_L1, event.fixedGridRhoFastjetAll, event.caches[0])

You can try in the same way. If you encounter any problems, feel free to ask me.

ikrommyd commented 3 months ago

@lgray CorrT1METJet appears to be part of regular nano. Should we add it in the schema?