split up cutflow, change NN->BDT in histogram name, clean up iso/quality cut logic

Author: Moelf

src/WWZ_ana.jl

@@ -26,7 +26,7 @@ function WWZ_chi2(Z_pair, W_pair, v_l_pid, v_l_tlv)
 end
 
 function WWZ_Cut(
-    Z_pair, W_pair, v_l_pid, v_l_order, v_l_tlv
+    Z_pair, W_pair, v_l_pid, v_l_order, v_l_tlv, dict, cutflow_ptr, NN_hist
 )
     chi2 = WWZ_chi2(Z_pair, W_pair, v_l_pid, v_l_tlv)
     FAIL = (false, Inf, W_pair)
@@ -37,10 +37,15 @@ function WWZ_Cut(
             WWZ_dilepton_mass < 12 && return FAIL
         end
     end
+    cutflow_ptr[] += 1
+    if NN_hist push!(dict[:CutFlow], cutflow_ptr[]) end
+
     @inbounds pt(v_l_tlv[v_l_order[1]]) < 30 && return FAIL
     @inbounds pt(v_l_tlv[v_l_order[2]]) < 15 && return FAIL
     @inbounds pt(v_l_tlv[v_l_order[3]]) < 8 &&  return FAIL
     @inbounds pt(v_l_tlv[v_l_order[4]]) < 6 &&  return FAIL
+    cutflow_ptr[] += 1
+    if NN_hist push!(dict[:CutFlow], cutflow_ptr[]) end
 
     # selected lepton min dR
     @inbounds for i in 1:4
@@ -49,7 +54,8 @@ function WWZ_Cut(
             dR < 0.1 && return FAIL
         end
     end
-
+    cutflow_ptr[] += 1
+    if NN_hist push!(dict[:CutFlow], cutflow_ptr[]) end
 
     return true, chi2, W_pair
 end

src/analysis_utils.jl

@@ -8,19 +8,19 @@ const ANALYSIS_DIR = Ref("/data/jiling/WVZ/v2.3_hists")
 const ONNX_MODEL_PATH = Ref("/data/grabanal/NN/NN_08_23.onnx")
 const BDT_MODEL_PATH = Ref("/home/rjacobse/BDT/kfolding/")
 
- function init_ONNX()
-     model=ONNX.load(ONNX_MODEL_PATH[], zeros(Float32, 30, 1))
-     rescaling_parameters = joinpath(dirname(@__DIR__), "config/NN_08_23_rescaling_parameters.json") |> read |> JSON3.read
-     NN_order = ("HT", "MET", "METPhi", "METSig", "Njet", "Wlep1_dphi", "Wlep1_eta",
-                 "Wlep1_phi", "Wlep1_pt", "Wlep2_dphi", "Wlep2_eta", "Wlep2_phi",
-                 "Wlep2_pt", "Zcand_mass", "Zlep1_dphi", "Zlep1_eta", "Zlep1_phi",
-                 "Zlep1_pt", "Zlep2_dphi", "Zlep2_eta", "Zlep2_phi", "Zlep2_pt",
-                 "leptonic_HT", "mass_4l", "other_mass", "pt_4l", "total_HT",
-                 "sr_SF_inZ", "sr_SF_noZ", "sr_DF")
-     return model, 
-     [rescaling_parameters["scale"][name] for name in NN_order],
-     [rescaling_parameters["min"][name] for name in NN_order]
- end
+function init_ONNX()
+    model=ONNX.load(ONNX_MODEL_PATH[], zeros(Float32, 30, 1))
+    rescaling_parameters = joinpath(dirname(@__DIR__), "config/NN_08_23_rescaling_parameters.json") |> read |> JSON3.read
+    NN_order = ("HT", "MET", "METPhi", "METSig", "Njet", "Wlep1_dphi", "Wlep1_eta",
+                "Wlep1_phi", "Wlep1_pt", "Wlep2_dphi", "Wlep2_eta", "Wlep2_phi",
+                "Wlep2_pt", "Zcand_mass", "Zlep1_dphi", "Zlep1_eta", "Zlep1_phi",
+                "Zlep1_pt", "Zlep2_dphi", "Zlep2_eta", "Zlep2_phi", "Zlep2_pt",
+                "leptonic_HT", "mass_4l", "other_mass", "pt_4l", "total_HT",
+                "sr_SF_inZ", "sr_SF_noZ", "sr_DF")
+    return model, 
+    [rescaling_parameters["scale"][name] for name in NN_order],
+    [rescaling_parameters["min"][name] for name in NN_order]
+end
 
 function init_BDT()
     SFinZs = Tuple(

src/constants.jl

@@ -1,6 +1,19 @@
+const CUTFLOW_NAMES = [
+    :in_minitree,
+    :pass_trigger,
+    :fourl_with_eta_mask,
+    :has_Z_candidate,
+    :pass_lepton_quality,
+    :pass_lepton_iso,
+    :SFOS_m_12GeVplus,
+    :lepton_minimum_pts,
+    :dilepton_dR_0p1plus,
+    :no_Bjet,
+    :MET_10GeVplus
+    ]
 const SIG_TAGS = ("Signal", )
 const BKG_TAGS = ("ZZ", "Zjets", "Zgamma", "ttbar", "WZ", "tZ", "ttZ", "tWZ", "VBS", "VH", "Others")
-const ALL_TAGS = (SIG_TAGS..., BKG_TAGS...)
+const ALL_TAGS = [SIG_TAGS...; BKG_TAGS...]
 const Z_m = 91.1876 # everything in GeV
 const e_mass = 0.51099885 / 1000
 const m_mass = 105.65837 / 1000

src/hist_exfiltration.jl

@@ -1,53 +1,31 @@
-function NN_hist_init(; sfsys, shape_variation)
-    if sfsys && (shape_variation != "NOMINAL")
-        error("can't do sf systematics and shape systematics at the same time")
-    end
-    _dict = Dict{Symbol, Hist1D}()
-    bins = 0:0.01:1
-    for n in (:SFinZ__NN, :SFnoZ__NN, :DF__NN)
-        _dict[Symbol(n, :__, shape_variation)] = Hist1D(Float64; bins, overflow=true)
-
+function populate_hist!(dict, shape_variation, symbols, bins, sfsys)
+    for n in symbols
+        dict[Symbol(n, :__, shape_variation)] = Hist1D(Float64; bins, overflow=true)
         !sfsys && continue
         for (_,vs) in SF_BRANCH_DICT
             for v in vs
-                _dict[Symbol(n, :__, v)] = Hist1D(Float64; bins, overflow=true)
+                dict[Symbol(n, :__, v)] = Hist1D(Float64; bins, overflow=true)
             end
         end
     end
-    bins = 0:5:300
-    for n in (:SFinZ__MET, :SFnoZ__MET, :DF__MET)
-        _dict[Symbol(n, :__, shape_variation)] = Hist1D(Float64; bins, overflow=true)
+end
 
-        !sfsys && continue
-        for (_,vs) in SF_BRANCH_DICT
-            for v in vs
-                _dict[Symbol(n, :__, v)] = Hist1D(Float64; bins, overflow=true)
-            end
-        end
+function NN_hist_init(; sfsys, shape_variation)
+    if sfsys && (shape_variation != "NOMINAL")
+        error("can't do sf systematics and shape systematics at the same time")
     end
-    bins = 0:5
-    for n in (:SFinZ__Njet, :SFnoZ__Njet, :DF__Njet)
-        _dict[Symbol(n, :__, shape_variation)] = Hist1D(Float64; bins, overflow=true)
+    _dict = Dict{Symbol, Hist1D}()
 
-        !sfsys && continue
-        for (_,vs) in SF_BRANCH_DICT
-            for v in vs
-                _dict[Symbol(n, :__, v)] = Hist1D(Float64; bins, overflow=true)
-            end
-        end
-    end
-    for n in (:ZZCR__Njet, :ttZCR__Njet)
-        _dict[Symbol(n, :__, shape_variation)] = Hist1D(Float64; bins, overflow=true)
+    bins = 0:0.01:1
+    populate_hist!(_dict, shape_variation, (:SFinZ__BDT, :SFnoZ__BDT, :DF__BDT), bins, sfsys)
 
-        !sfsys && continue
-        for (_,vs) in SF_BRANCH_DICT
-            for v in vs
-                _dict[Symbol(n, :__, v)] = Hist1D(Float64; bins, overflow=true)
-            end
-        end
-    end
+    bins = 0:5:300
+    populate_hist!(_dict, shape_variation, (:SFinZ__MET, :SFnoZ__MET, :DF__MET), bins, sfsys)
+
+    bins = 0:5
+    populate_hist!(_dict, shape_variation, (:SFinZ__Njet, :SFnoZ__Njet, :DF__Njet, :ZZCR__Njet, :ttZCR__Njet), bins, sfsys)
 
-    _dict[:CutFlow] = Hist1D(Int; bins=0:20)
+    _dict[:CutFlow] = Hist1D(Int; bins=1:20)
 
     return _dict
 end

src/mainlooper.jl

@@ -38,17 +38,22 @@ function main_looper(mytree, sumWeight, dict, pusher!, models,
     model = models # for BDT
     for evt in mytree
         ### initial_cut
-        cutflow_ptr = 0
-        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
+        cutflow_ptr = Ref(1)
+        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
+
+        !(evt.passTrig) && continue
+        cutflow_ptr[] += 1
+        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
+
         v_m_eta_orig, v_e_caloeta_orig = evt.v_m_eta, evt.v_e_caloeta
         e_etamask = [abs(η) < 2.47 && (abs(η)<1.37 || abs(η)>1.52) for η in v_e_caloeta_orig]
         m_etamask = [abs(η) < 2.5 for η in v_m_eta_orig]
         v_l_pid = @views Vcat(evt.v_e_pid[e_etamask], evt.v_m_pid[m_etamask])
-        
         Nlep = length(v_l_pid)
         Nlep != 4 && continue
-        cutflow_ptr += 1
-        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
+        cutflow_ptr[] += 1
+        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
+
         Nelec = count(e_etamask)
         Nmuon = Nlep - Nelec
         (; v_e_pt, v_e_eta, v_e_phi,
@@ -61,65 +66,61 @@ function main_looper(mytree, sumWeight, dict, pusher!, models,
         v_e_tlv = @views LorentzVectorCyl.(v_e_pt[e_etamask], v_e_eta[e_etamask], v_e_phi[e_etamask], v_e_m[e_etamask])
         v_m_tlv = @views LorentzVectorCyl.(v_m_pt[m_etamask], v_m_eta[m_etamask], v_m_phi[m_etamask], v_m_m[m_etamask])
         v_l_tlv = Vcat(v_e_tlv, v_m_tlv)
+
         Z_pair, W_pair, best_Z_mass = Find_Z_Pairs(v_l_pid, v_l_tlv)
         isinf(best_Z_mass) && continue
-        cutflow_ptr += 1
-        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
         other_mass = mass(v_l_tlv[W_pair[1]] + v_l_tlv[W_pair[2]])
         abs(best_Z_mass - Z_m) > 20 && continue
-        cutflow_ptr += 1
-        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
+        cutflow_ptr[] += 1
+        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
+
         mass_4l = mass(sum(v_l_tlv))
         mass_4l < 0.0 && continue
-        cutflow_ptr += 1
-        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
         ### end of initial_cut
-        !(evt.passTrig) && continue
-        cutflow_ptr += 1
-        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
+        
         v_l_order = sortperm(v_l_tlv; by=pt, rev=true)
-        v_l_medium = @views Vcat(evt.v_e_LHMedium[e_etamask] , evt.v_m_medium[m_etamask]) #quality
 
-        ############## use PLIV for W lepton ISO #################
-        v_l_PLTight = Vcat(evt.v_e_passIso_PLImprovedTight[e_etamask], evt.v_m_passIso_PLImprovedTight[m_etamask])
-        # if controlregion == :Zjets
-        #     failed_PLTight = any(==(true), v_l_PLTight[W_pair])
-        #     failed_PLTight && continue
-        #     v_l_Loose = Vcat(evt.v_e_passIso_Loose_VarRad, evt.v_m_passIso_PflowLoose_VarRad)
-        #     failed_Loose = all(==(true), v_l_Loose[W_pair])
-        #     failed_Loose && continue
-        # else
-            failed_PLTight = any(==(false), v_l_PLTight[W_pair])
-            failed_PLTight && continue
-            cutflow_ptr += 1
-            if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
-        # end
-        (;
-         v_e_passIso_Loose_VarRad,
-         v_m_passIso_PflowLoose_VarRad,
-        ) = evt
+        ### QUALITY
+        v_l_medium = @views Vcat(evt.v_e_LHMedium[e_etamask] , evt.v_m_medium[m_etamask])
+        # for W leptons, require medium quality
+        !v_l_medium[W_pair[1]] && continue
+        !v_l_medium[W_pair[2]] && continue
+        # for Z leptons only Loose requirement, no additional quality cut
+        #
+        cutflow_ptr[] += 1
+        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
+
+        ### ISOLATION
+        # for W leptons, require PLIVTight Isolation
+        v_l_PLTight = @views Vcat(evt.v_e_passIso_PLImprovedTight[e_etamask], evt.v_m_passIso_PLImprovedTight[m_etamask])
+        !v_l_PLTight[W_pair[1]] && continue
+        !v_l_PLTight[W_pair[2]] && continue
+        # for Z leptons require Loose isolation
+        (;v_e_passIso_Loose_VarRad, v_m_passIso_PflowLoose_VarRad) = evt
         v_l_passIso = @views Vcat(v_e_passIso_Loose_VarRad[e_etamask], v_m_passIso_PflowLoose_VarRad[m_etamask])
+        !v_l_passIso[Z_pair[1]] && continue
+        !v_l_passIso[Z_pair[2]] && continue
+
+        cutflow_ptr[] += 1
+        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
+
 
         pass_WWZ_cut, chisq, W_id = WWZ_Cut(
                                             Z_pair,
                                             W_pair,
                                             v_l_pid,
                                             v_l_order,
-                                            v_l_tlv
+                                            v_l_tlv,
+                                            dict,
+                                            cutflow_ptr,
+                                            NN_hist
                                            )
-        for i in 1:2
-            ### for Z leptons isolation: Loose(e) and Loose(mu)
-            !v_l_passIso[Z_pair[i]] && (pass_WWZ_cut = false)
-
-            ### for W leptons, require medium quality and PLIV tight
-            !v_l_medium[W_pair[i]] && (pass_WWZ_cut = false)
-        end
 
         !pass_WWZ_cut && continue
-        cutflow_ptr += 1
-        if NN_hist pusher!(dict[:CutFlow], cutflow_ptr) end
 
         has_b = any(evt.v_j_btag77)
+        cutflow_ptr[] += 1
+        if !has_b && NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
 
         (; MET, METSig, METPhi) = evt
         MET /= 1000
@@ -133,7 +134,7 @@ function main_looper(mytree, sumWeight, dict, pusher!, models,
             make_sfsys_wgt!(evt, wgt_dict,
                             :weight; sfsys, pre_mask=1)
             make_sfsys_wgt!(evt, wgt_dict, 
-                            :v_j_wgt_btag77, : ; sfsys)
+                            :v_j_wgt_btag77, Colon() ; sfsys)
             # I hate indexing
             Z_pair_in_e = filter(<=(Nelec), Z_pair)
             Z_pair_in_m = filter!(>(0), Z_pair .- Nelec)
@@ -162,14 +163,12 @@ function main_looper(mytree, sumWeight, dict, pusher!, models,
             make_sfsys_wgt!(evt, wgt_dict,
                             :v_m_wgtIso_PflowLoose_VarRad,
                             Z_pair_in_m; pre_mask = m_etamask, sfsys)
-            # if controlregion != :ZJets
-                make_sfsys_wgt!(evt, wgt_dict,
-                                :v_e_wgtIso_PLImprovedTight_Medium, 
-                                W_pair_in_e; pre_mask = e_etamask, sfsys)
-                make_sfsys_wgt!(evt, wgt_dict,
-                                :v_m_wgtIso_PLImprovedTight, 
-                                W_pair_in_m; pre_mask = m_etamask, sfsys)
-            # end
+            make_sfsys_wgt!(evt, wgt_dict,
+                            :v_e_wgtIso_PLImprovedTight_Medium, 
+                            W_pair_in_e; pre_mask = e_etamask, sfsys)
+            make_sfsys_wgt!(evt, wgt_dict,
+                            :v_m_wgtIso_PLImprovedTight, 
+                            W_pair_in_m; pre_mask = m_etamask, sfsys)
         end
 
         lep1_pid, lep2_pid, lep3_pid, lep4_pid = @view v_l_pid[v_l_order]
@@ -238,6 +237,8 @@ function main_looper(mytree, sumWeight, dict, pusher!, models,
         if cr_ZZ || cr_ttZ
             SR = -1
         end
+        cutflow_ptr[] += 1
+        if MET>10 && NN_hist pusher!(dict[:CutFlow], cutflow_ptr[]) end
 
         if NN_hist && !arrow_making
             region_prefix = if sr_SF_inZ
@@ -256,9 +257,7 @@ function main_looper(mytree, sumWeight, dict, pusher!, models,
                 elseif cr_ttZ
                     pusher!(dict[Symbol(:ttZCR, :__Njet, :__, k)], Njet, v)
                 elseif SR >= 0
-                    cutflow_ptr += 1
-                    pusher!(dict[:CutFlow], cutflow_ptr)
-                    pusher!(dict[Symbol(region_prefix, :__NN, :__, k)], NN_score, v)
+                    pusher!(dict[Symbol(region_prefix, :__BDT, :__, k)], NN_score, v)
                     pusher!(dict[Symbol(region_prefix, :__MET, :__, k)], MET, v)
                     pusher!(dict[Symbol(region_prefix, :__Njet, :__, k)], Njet, v)
                 end