Incorrect MET correction method

[9GaoHong]

Hi @lgray

There are two questions about MET correction：

1. It seems that CorrectedMETFactory.py has not been updated according to the official correction formula.

In this script, propagate JEC effect to pfMET through the corrected_polar_met function:

$$ \rm E_T^{miss,corr} = E_T^{miss}+\sum_{jets}({p_{T,jet}^{JEC}-p_{T,jet}^{orig})} $$

def corrected_polar_met(
    met_pt, met_phi, jet_pt, jet_phi, jet_pt_orig, positive=None, dx=None, dy=None
):
    sj, cj = numpy.sin(jet_phi), numpy.cos(jet_phi)
    x = met_pt * numpy.cos(met_phi) + awkward.sum((jet_pt - jet_pt_orig) * cj, axis=1)
    y = met_pt * numpy.sin(met_phi) + awkward.sum((jet_pt - jet_pt_orig) * sj, axis=1)
    if positive is not None and dx is not None and dy is not None:
        x = x + dx if positive else x - dx
        y = y + dy if positive else y - dy

    return awkward.zip(
        {"pt": numpy.hypot(x, y), "phi": numpy.arctan2(y, x)}, depth_limit=1
    )

Can you illustrate where this correction method was referenced from？

And by checking this twiki link, I found that the formula for MET Type-I correction(mostly used when processing with NanoAOD) should be like:

$$ \rm E_T^{miss,corr} = E_T^{miss}-\sum_{jets}({p_{T,jet}^{L123}-p_{T,jet}^{L1})} $$

However, it is evident that in CorrectedMETFactory.py , jet_pt is the ${p_{T,\text{jet}}^{L123}}$, and jet_pt_orig is pTRaw instead of ${p_{T,\text{jet}}^{L1}}$. This will result in incorrect MET correction.

2. The target object for CorrectedMETFactory.py use with is MET, not RawMET.
According to the following code, it can be seen that this script is applicable to MET, not RawMET.
Because in NanoAOD samples, only MET has branchs related to UnClusteredEnergyDeltaX/Y, and RawMET does not have.

out_dict["MET_UnclusteredEnergy"] = dask_awkward.map_partitions(
            create_variants,
            MET,
            corrected_jets,
            MET[self.name_map["UnClusteredEnergyDeltaX"]],
            MET[self.name_map["UnClusteredEnergyDeltaY"]],
            label="UnclusteredEnergy_met",
)

But in NanoAOD samples, the branch names and what they correspond to are as follows:

• RawMET: raw (i.e., uncorrected) PF MET
• MET: Type-1 corrected PF MET
• RawPuppiMET: raw PUPPI MET
• PuppiMET: Type-1 corrected PUPPI MET

So it is unreasonable to use MET because it has already undergone Type-I correction, and RawMET should be used instead.
How do you think?

[lgray]

Please go ahead and supply a fix. I'll be happy to review associated code changes.

[9GaoHong]

Sure, I have submitted the PR, please have a look.
How to use the CorrectedMETFactory function after the update? The example code is as follows:

from coffea.nanoevents import NanoEventsFactory, NanoAODSchema
from coffea.lookup_tools import extractor
from coffea.jetmet_tools import FactorizedJetCorrector
from coffea.jetmet_tools import JetResolution
from coffea.jetmet_tools import JECStack
from coffea.jetmet_tools import JetCorrectionUncertainty
from coffea.jetmet_tools import JetResolutionScaleFactor
from coffea.jetmet_tools import CorrectedJetsFactory
from coffea.jetmet_tools import CorrectedMETFactory

from coffea.lookup_tools import dense_lookup
import awkward as ak
import numpy as np
import os


jec_name_map = {
    'JetPt': 'pt',
    'JetMass': 'mass',
    'JetEta': 'eta',
    'JetA': 'area',
    'ptRaw': 'pt_raw',
    'massRaw': 'mass_raw',
    'Rho': 'rho',
    'METpt': 'pt',
    'METphi': 'phi',
    'JetPhi': 'phi',
    'UnClusteredEnergyDeltaX': 'MetUnclustEnUpDeltaX',
    'UnClusteredEnergyDeltaY': 'MetUnclustEnUpDeltaY',
    'ptGenJet': 'pt_gen',
}

def update_collection(event, coll):
    out = event
    for name, value in coll.items():
        out = ak.with_field(out, value, name)
    return out


def add_jme_variables(jets , events_rho):
    jets['pt_raw'  ] = (1 - jets.rawFactor) * jets.pt * (1- jets.muonSubtrFactor)
    jets['mass_raw'] = (1 - jets.rawFactor) * jets.mass * (1- jets.muonSubtrFactor)

    if hasattr(jets, 'matched_gen'):
        jets['pt_gen'] = ak.values_astype(ak.fill_none(jets.matched_gen.pt, 0), np.float32)
    else:
        jets['pt_gen'] = ak.Array(np.zeros(len(jets), dtype=np.float32))
    
    jets['rho'     ] = ak.broadcast_arrays(events_rho, jets.pt)[0]
    return jets

class JMEUncertainty:
    def __init__(
        self,
        jec_tag: str = 'Summer19UL18_V5_MC',
        jer_tag: str = 'Summer19UL18_JRV2_MC',
        era: str = "2018",
        is_mc: bool = True
    ):
        _data_path = os.path.join(os.path.dirname(__file__), 'data/jme')
        extract = extractor()
        extract_L1 = extractor()
        correction_list = [
            # Jet Energy Correction
            f'* * {_data_path}/{era}/{jec_tag}_L1FastJet_AK4PFchs.jec.txt',
            f'* * {_data_path}/{era}/{jec_tag}_L2L3Residual_AK4PFchs.jec.txt',
            f'* * {_data_path}/{era}/{jec_tag}_L2Relative_AK4PFchs.jec.txt',
            f'* * {_data_path}/{era}/{jec_tag}_L3Absolute_AK4PFchs.jec.txt',
        ]
        correction_list_L1 = [
            # Jet Energy Correction
            f'* * {_data_path}/{era}/{jec_tag}_L1FastJet_AK4PFchs.jec.txt',
        ]
        if is_mc:
            correction_list += [    
                # Jet Energy Resolution
                f'* * {_data_path}/{era}/RegroupedV2_{jec_tag}_UncertaintySources_AK4PFchs.junc.txt',
                f'* * {_data_path}/{era}/{jer_tag}_PtResolution_AK4PFchs.jr.txt',
                f'* * {_data_path}/{era}/{jer_tag}_SF_AK4PFchs.jersf.txt',
            ]
            
        
        jec_name_map.update({'ptGenJet': 'pt_gen'})
        
        extract_L1.add_weight_sets(correction_list_L1)
        extract_L1.finalize()
        evaluator_L1 = extract_L1.make_evaluator()
        jec_inputs_L1 = {
            name: evaluator_L1[name] for name in dir(evaluator_L1)
        }
        self.jec_stack_L1 = JECStack(jec_inputs_L1)
        self.jec_factory_L1 = CorrectedJetsFactory(jec_name_map, self.jec_stack_L1)
        
        extract.add_weight_sets(correction_list)
        extract.finalize()
        evaluator = extract.make_evaluator()
        jec_inputs = {
            name: evaluator[name] for name in dir(evaluator)
        }
        self.jec_stack = JECStack(jec_inputs)
        self.jec_factory = CorrectedJetsFactory(jec_name_map, self.jec_stack)
        self.met_factory = CorrectedMETFactory(jec_name_map)

    def corrected_jets(self, jets, event_rho, lazy_cache):
        jets_L123 = self.jec_factory.build(
            add_jme_variables(jets, event_rho),
            lazy_cache #events.caches[0]
        )
        jets_L1 = self.jec_factory_L1.build(
            add_jme_variables(jets, event_rho),
            lazy_cache #events.caches[0]
        )
        emFraction = jets_L123.chEmEF + jets_L123.neEmEF
        mask_jec = (jets_L123.pt > 15) & (emFraction <= 0.9)
        selected_jets_L123 = ak.mask(jets_L123, mask_jec)
        jets_jec = selected_jets_L123
        jets_jec['pt'] = selected_jets_L123.pt - jets_L1.pt
        return jets_jec
    def corrected_jets_jec(self, jets, event_rho, lazy_cache):
        return self.jec_factory.build(
            add_jme_variables(jets, event_rho),
            lazy_cache #events.caches[0]
        )
    def corrected_met(self, rawmet, met, jets, event_rho, lazy_cache):
        return self.met_factory.build(
            rawmet,met,
            add_jme_variables(jets, event_rho),
            lazy_cache=lazy_cache # events.caches[0]
        )

self._jmeu = JMEUncertainty(jec_tag, jer_tag, era=self._era, is_mc=(len(run_period)==0))
met  = self._jmeu.corrected_met(event.RawMET,event.MET, jets_col, event.fixedGridRhoFastjetAll, event.caches[0])

If there are any issues with these codes, please feel free to point them out.

[nsmith-]

@mcremone can you check this?
I think the original intention is that this CorrectedMETFactory would not create the Type-1-corrected MET, but rather would take as input a Type-1 corrected MET, and then add only the delta corresponding to alternative JEC variations, i.e. subtract the L1-corrected jets and then add the (L1+systematic)-corrected jets.