fails with a dropped column in column optimization (particularly the broadcasted rho column!)

jec_dropped_column.py

import time

import awkward as ak
import dask_awkward as dak
import numpy as np

import os

from coffea.lookup_tools import extractor
from coffea.jetmet_tools import FactorizedJetCorrector

extract = extractor()

extract.add_weight_sets(
    [
        "* * tests/samples/Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi.jec.txt.gz",
        "* * tests/samples/Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi.jec.txt.gz",
        "* * tests/samples/Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi.jec.txt.gz",
        "* * tests/samples/Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi.jec.txt.gz",
    ]
)

extract.finalize()

evaluator =  extract.make_evaluator()

from coffea.jetmet_tools import CorrectedJetsFactory, CorrectedMETFactory, JECStack
from coffea.nanoevents import NanoEventsFactory

events = NanoEventsFactory.from_root(
    os.path.abspath("tests/samples/nano_dy.root"),
    treepath="Events",
    metadata={},
    permit_dask=True,
)

jec_stack_names = [
    "Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi",
    "Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi",
    "Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi",
    "Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi",
]

jec_inputs = {name: evaluator[name] for name in jec_stack_names}
jec_stack = JECStack(jec_inputs)

name_map = jec_stack.blank_name_map
name_map["JetPt"] = "pt"
name_map["JetMass"] = "mass"
name_map["JetEta"] = "eta"
name_map["JetA"] = "area"

jets = events.Jet

jets = dak.with_field(jets, (1 - jets["rawFactor"]) * jets.pt, "pt_raw")
jets = dak.with_field(jets, (1 - jets["rawFactor"]) * jets.mass, "mass_raw")
jets = dak.with_field(
    jets,
    dak.values_astype(dak.fill_none(jets.matched_gen.pt, 0), np.float32),
    "pt_gen",
)
jets = dak.with_field(
    jets, dak.broadcast_arrays(events.fixedGridRhoFastjetAll, jets.pt)[0], "rho"
)
name_map["ptGenJet"] = "pt_gen"
name_map["ptRaw"] = "pt_raw"
name_map["massRaw"] = "mass_raw"
name_map["Rho"] = "rho"

corrector = FactorizedJetCorrector(
    **{name: evaluator[name] for name in jec_stack_names}
)

lazy_correction = corrector.getCorrection(
    JetEta=jets.eta,
    Rho=jets.rho,
    JetPt=jets.pt_raw,
    JetA=jets.area,
)

print(dak.necessary_columns(lazy_correction))

check_corrs = lazy_correction.compute()