Pickle_plot.py

#!/usr/bin/env python
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
import mplhep as hep

# Set up directory
pickle_dir = "/home/jtong/lbyl/cutflow_jar/"

# Define the load_or_compute function (here, we only load; if file doesn't exist, we raise an error)
def load_or_compute(filename, data):
    if os.path.exists(filename):
        with open(filename, 'rb') as f:
            loaded_data = pickle.load(f)
    else:
        # If no file exists, use the provided default data and save it
        loaded_data = data
        with open(filename, 'wb') as f:
            pickle.dump(loaded_data, f)
    return loaded_data

# Use ATLAS style
hep.style.use(hep.style.ATLASAlt)

fig_size = (10, 10)

# Histogram for diphoton masses
diphoton_masses = load_or_compute(os.path.join(pickle_dir, "m_gg.pkl"), None)
m_gg = np.histogram(diphoton_masses, bins=50)
plt.figure(figsize=fig_size)
hep.histplot(m_gg, histtype='step', linewidth=2, label="Signal")
hep.atlas.label(year=2025, data = False, label="Work in Progress", loc=4, com = "?")
plt.margins(x=0.05)  # 5% margin on the left and right of the x-axis
plt.ylim(top = 2200)
plt.xlabel("Diphoton Mass (GeV)")
plt.ylabel("Counts")
plt.title("Diphoton Mass Distribution")
plt.legend()
plt.savefig("Plots/diphoton_mass_distribution.png")
plt.close()

# Histogram for diphoton pT
diphoton_pts = load_or_compute(os.path.join(pickle_dir, "pt_gg.pkl"), None)
pt_gg = np.histogram(diphoton_pts, bins=50)
plt.figure(figsize=fig_size)
hep.histplot(pt_gg, histtype='step', linewidth=2, label="Signal")
hep.atlas.label(year=2025, data = False, label="Work in Progress", loc=4, com = "?")
plt.margins(x=0.05)  # 5% margin on the left and right of the x-axis
plt.ylim(top = 1700)
plt.xlabel("Diphoton $p_{T}$ (GeV)")
plt.ylabel("Counts")
plt.title("Diphoton $p_{T}$ Distribution")
plt.legend()
plt.savefig("Plots/diphoton_pt_distribution.png")
plt.close()

# Histogram for diphoton acoplanarity
diphoton_acoplanarity = load_or_compute(os.path.join(pickle_dir, "a_gg.pkl"), None)
a_gg = np.histogram(diphoton_acoplanarity, bins=50)
plt.figure(figsize=fig_size)
hep.histplot(a_gg, histtype='step', linewidth=2, label="Signal")
hep.atlas.label(year=2025, label="Work in Progress", loc=4, com = "?")
plt.margins(x=0.05)  # 5% margin on the left and right of the x-axis
plt.xscale("linear")
plt.ylim(top = 350)
plt.xlabel("Diphoton $A_{\phi}$")
plt.ylabel("Counts")
plt.title("Diphoton $A_{\phi}$ Distribution")
plt.legend()
plt.savefig("Plots/Diphoton_acoplanarity.png")
plt.close()