Add yield cutflow and refactor the insanity (#97)

* fix CutFlow duplications

* add some tests

Co-authored-by: Rafael Jacobsen <[email protected]>
Co-authored-by: Rafael Antonio Jacobsen <[email protected]>

Author: 3 people

README.md

@@ -31,6 +31,6 @@ julia> addprocs(HTCManager(10); extrajdl=["+queue=\"short\""], exeflags = `-e 'i
 Waiting for 10 workers: 1 2 3 4 5 6 7 8 9 10 .
 
 julia> foreach(WVZAnalysis.ALL_TAGS) do tag
-           WVZAnalysis.hist_root_pmap(tag; scouting=false, output_dir="/data/jiling/WVZ/v2.3_hists_uproot_nov14_andMET/");
+           WVZAnalysis.hist_root(tag; scouting=false, output_dir="/data/jiling/WVZ/v2.3_hists_uproot_nov14_andMET/");
        end
 ```

src/WWZ_ana.jl

@@ -24,38 +24,3 @@ function WWZ_chi2(Z_pair, W_pair, v_l_pid, v_l_tlv)
     end
     return chi2
 end
-
-function WWZ_Cut(
-    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)
-    @inbounds for i in eachindex(v_l_tlv)
-        for j in (i + 1):length(v_l_tlv)
-            v_l_pid[i] + v_l_pid[j] != 0 && continue
-            WWZ_dilepton_mass = mass(v_l_tlv[i] + v_l_tlv[j])
-            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
-        for j in (i + 1):4
-            dR = deltar(v_l_tlv[v_l_order[i]], v_l_tlv[v_l_order[j]])
-            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

@@ -6,7 +6,7 @@ const MINITREE_DIR = Ref("/data/jiling/WVZ/v2.3")
 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/")
+const BDT_MODEL_PATH = Ref("/data/rjacobse/WVZ/kfolding/")
 
 function init_ONNX()
     model=ONNX.load(ONNX_MODEL_PATH[], zeros(Float32, 30, 1))
@@ -326,54 +326,12 @@ end
 
 
 """
-    hist_root(tag; output_dir, kw...)
-
-return a dictionary of all systematics histograms for a given tag, the `output_dir` is
-used to store Julia Serializatio files as a backup measure, and for later conversion to
-`.root` histograms via `PythonCall.jl` + `uproot`.
-"""
-function hist_root(tag; output_dir, kw...)
-    p = output_dir
-    if !isdir(p)
-        mkdir(p)
-    end
-    @show Threads.nthreads()
-    @info "-------------- $tag SF begin ------------ "
-    sf_tasks = prep_tasks(tag)
-    println("$(length(sf_tasks)) tasks in total")
-    p1 = Progress(length(sf_tasks))
-    sf_list = map(sf_tasks) do t
-        x = main_looper(t)
-        next!(p1)
-        x
-    end
-    sf_hist = reduce(mergewith(+), sf_list)
-
-    @info "-------------- $tag shapes begin ------------ "
-    shape_tasks = mapreduce(shape_variation -> prep_tasks(tag; shape_variation), vcat,
-                                 SHAPE_TREE_NAMES)
-    sort!(shape_tasks; by = x->x.path)
-    println("$(length(shape_tasks)) tasks in total")
-    p2 = Progress(length(shape_tasks))
-    shape_list = map(shape_tasks) do t
-        x = main_looper(t)
-        next!(p2)
-        x
-    end
-    shape_hist = reduce(mergewith(+), shape_list)
-
-    Hs = merge(sf_hist, shape_hist)
-    serialize(joinpath(p, "$(tag).jlserialize"), Hs)
-    Hs
-end
-
-"""
-    hist_root_pmap(tag; output_dir, kw...)
+    hist_root(tag; mapper=pmap, no_shape=false, output_dir, kw...)
 
 Similar to the one without `pmap`, except uses pmap for everything. Checkout `ClusterManager.jl`
 and be sure to have a handful of workers before running this.
 """
-function hist_root_pmap(tag; output_dir, kw...)
+function hist_root(tag; mapper=pmap, no_shape = false, output_dir, kw...)
     p = output_dir
     if !isdir(p)
         mkdir(p)
@@ -382,16 +340,25 @@ function hist_root_pmap(tag; output_dir, kw...)
 
     all_tasks = prep_tasks(tag; sfsys=true)
     if tag != "Data"
-        @info "-------------- $tag SF + shapes ------------ "
-        shape_tasks = mapreduce(shape_variation -> prep_tasks(tag; shape_variation, sfsys=false), vcat,
-                                SHAPE_TREE_NAMES)
-        sort!(shape_tasks; by = x->x.path)
-        append!(all_tasks, shape_tasks)
+        if no_shape
+            @info "-------------- $tag SF ------------ "
+        else
+            @info "-------------- $tag SF + shapes ------------ "
+            shape_tasks = mapreduce(shape_variation -> prep_tasks(tag; shape_variation, sfsys=false), vcat,
+                                    SHAPE_TREE_NAMES)
+            sort!(shape_tasks; by = x->x.path)
+            append!(all_tasks, shape_tasks)
+        end
     else
         @info "-------------- !!! $tag !!! ------------ "
     end
     println("$(length(all_tasks)) tasks in total")
-    all_list = @showprogress pmap(main_looper, all_tasks)
+    prog = Progress(length(all_tasks))
+    all_list = mapper(all_tasks) do task
+        _m = main_looper(task)
+        next!(prog)
+        return _m
+    end
     # TODO
     # process_theory_syst!.(all_list)
     Hs = reduce(mergewith(+), all_list)

src/hist_exfiltration.jl

@@ -26,6 +26,13 @@ function NN_hist_init(; sfsys, shape_variation)
     populate_hist!(_dict, shape_variation, (:SFinZ__Njet, :SFnoZ__Njet, :DF__Njet, :ZZCR__Njet, :ttZCR__Njet), bins, sfsys)
 
     _dict[:CutFlow] = Hist1D(Int; bins=1:20)
+    _dict[:CutFlowWgt] = Hist1D(Float64; bins=1:20)
+    _dict[:SFinZCutFlow] = Hist1D(Int; bins=1:20)
+    _dict[:SFinZCutFlowWgt] = Hist1D(Float64; bins=1:20)
+    _dict[:SFnoZCutFlow] = Hist1D(Int; bins=1:20)
+    _dict[:SFnoZCutFlowWgt] = Hist1D(Float64; bins=1:20)
+    _dict[:DFCutFlow] = Hist1D(Int; bins=1:20)
+    _dict[:DFCutFlowWgt] = Hist1D(Float64; bins=1:20)
 
     return _dict
 end