DisplacedTauAnalysis/fast_plotting.py at main · brandiskip/DisplacedTauAnalysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# fast_plotting.py
# use this file to only call one hist function to avoid making all histograms from all functions
import os
import awkward as ak
import numpy as np
import matplotlib.pyplot as plt
import hist
from hist import Hist, axis, intervals
from coffea.nanoevents import NanoEventsFactory, PFNanoAODSchema
import json

# Import functions from jet_selection.py
from jet_selection import (
    process_events,
    select_and_define_leading_jets,
    match_gen_taus,
    flatten_gen_tau_vars
    #delta_r_mask
)

# Import functions from jet_plotting.py
from jet_plotting import (
    plot_efficiency_from_json,
    overlay_dxy_match_hists
)

# Load the file
filenames = {
    'Stau_100_1mm'    : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_100_1mm/*.root',
    'Stau_100_10mm'   : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_100_10mm/*.root',
    'Stau_100_100mm'  : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_100_100mm/*.root',
    'Stau_100_1000mm' : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_100_1000mm/*.root',
    'Stau_200_1mm'    : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_200_1mm/*.root',
    'Stau_200_10mm'   : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_200_10mm/*.root',
    'Stau_200_100mm'  : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_200_100mm/*.root',
    'Stau_200_1000mm' : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_200_1000mm/*.root',
    'Stau_300_1mm'    : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_300_1mm/*.root',
    'Stau_300_10mm'   : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_300_10mm/*.root',
    'Stau_300_100mm'  : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_300_100mm/*.root',
    'Stau_300_1000mm' : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_300_1000mm/*.root',
    'Stau_500_1mm'    : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_500_1mm/*.root',
    'Stau_500_10mm'   : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_500_10mm/*.root',
    'Stau_500_100mm'  : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_500_100mm/*.root',
    'Stau_500_1000mm' : 'root://cmseos.fnal.gov///store/user/dally/first_skim_muon_root_GenVtx/Stau_500_1000mm/*.root',
}

#PFNanoAODSchema.mixins["DisMuon"] = "Muon"

samples = {}
for sample_name, files in filenames.items():
    samples[sample_name] = NanoEventsFactory.from_root(
        {files: "Events"},
        schemaclass=PFNanoAODSchema,
        metadata={"dataset": "MC"}
    ).events()

# Create output folder for histograms
output_dir = "histograms"
os.makedirs(output_dir, exist_ok=True)

# ----------------------------------------------------------------------
# Main loop: Process each sample and produce histograms.
# ----------------------------------------------------------------------
if __name__ == '__main__':
    for sample_name, events in samples.items():
        print(f"Processing sample: {sample_name}")

        # Process events: select staus, taus, and apply event filters
        cut_filtered_events = process_events(events)

        # Select jets and define leading jets (using the filtered events)
        jets, leading_pt_jets, leading_score_jets, leading_dxy_jets, total_nJets = select_and_define_leading_jets(cut_filtered_events)

        # Match gen taus to jets using both pt-based and leading-score methods
        (gen_taus, gen_taus_matched_by_dxy,
            gen_taus_matched_by_pt, jet_matched_gen_taus_pt,
            gen_vis_taus_matched_by_score, jet_matched_gen_taus_score,
            matched_leading_jets_flat, all_unmatched_jets_pt) = match_gen_taus(
            cut_filtered_events,
            leading_pt_jets, leading_dxy_jets,
            leading_score_jets,
            jets, total_nJets, sample_name
            )

        (gen_taus_flat_dxy, gen_taus_matched_by_flat_dxy, leading_dxy_jets_flat_dxy, gen_taus_flat_pt,
            gen_taus_matched_by_pt_flat_dxy, gen_taus_matched_by_pt_flat_pt) = flatten_gen_tau_vars(
            gen_taus, gen_taus_matched_by_pt, leading_dxy_jets, gen_taus_matched_by_dxy)

        #overlay_dxy_match_hists(gen_taus_matched_by_flat_dxy, leading_dxy_jets_flat_dxy, output_dir, sample_name)


    # Now call plot_sample_grid using stored efficiency results
    plot_efficiency_from_json("dxy_efficiency_results.json", "efficiency_grid.pdf")