UPCSample_AOD_Analyzer_Python_Junda.py

#!/usr/bin/env python

# Set up ROOT and other necessary libraries
import ROOT
import uproot
import numpy as np
import ctypes
import matplotlib.pyplot as plt
import mplhep as hep
hep.style.use(hep.style.ROOT)

# Define GeV conversion
GeV = 1000

# Function to calculate invariant mass
def invariant_mass(obj1, obj2):
    p1 = ROOT.TLorentzVector()
    p2 = ROOT.TLorentzVector()
    p1.SetPtEtaPhiM(obj1.pt(), obj1.eta(), obj1.phi(), 0)  # Massless assumption for photons
    p2.SetPtEtaPhiM(obj2.pt(), obj2.eta(), obj2.phi(), 0)
    return (p1 + p2).M()

# Load data (example, needs adapting for actual input format)
fileName = "mc23_5p36TeV.860223.SuperChic_4p2_gammagamma2gammagamma_2Mgg500.merge.AOD.e8558_e8528_s4357_r15716_r14663_tid40981863_00/AOD.40981863._000001.pool.root.1"  # Replace with file, signal, yyee, CEP etc that we have in the folder

treeName = 'CollectionTree'
f = ROOT.TFile.Open(fileName)

# Make the "transient tree" and load any iterators we may need
import xAODRootAccess.GenerateDVIterators
t = ROOT.xAOD.MakeTransientTree(f, treeName)

# Cut flow dictionary
cut_flow = {
    "Total Events": 0,
    "Pass Track Veto": 0,
    "Pass Pixel Track Veto": 0,
    "Pass Photon Selection": 0,
    "Pass Diphoton Selection": 0,
}

ev_n_chargedtracks = 0
ev_pix_track_eta = []
ev_pix_track_phi = []
ev_pix_track_pt = []

diphoton_masses = []
diphoton_pts = []

# Loop over events, stopping before event 5000 for now
for i, event in enumerate(t):
    if i >= 4999:
        print("Stopping at event 5000 to prevent I/O errors.")
        break
    
    cut_flow["Total Events"] += 1
    
    n_electrons = 0
    for el in event.Electrons:
        el_pt = el.pt()
        el_quality = el.isGoodOQ(ROOT.xAOD.EgammaParameters.BADCLUSELECTRON)
        el_eta = abs(el.caloCluster().etaBE(2))
        if el_quality and el_pt> 2 * GeV and el_eta < 2.47 and (el_eta > 1.52 or el_eta < 1.37):
            n_electrons += 1


    # Track veto
    n_charged_tracks = 0
    for track in event.InDetTrackParticles:
        eta = track.eta()
        pt = track.pt()
        passesTrackKinematicPreSelection = (pt > 0.1 * GeV and abs(eta) < 2.5)

        npix = ctypes.c_ubyte(0)
        track.summaryValue(npix, ROOT.xAOD.SummaryType.numberOfPixelHits)
        nsct = ctypes.c_ubyte(0)
        track.summaryValue(nsct, ROOT.xAOD.numberOfSCTHits)
        nBL = ctypes.c_ubyte(0)
        track.summaryValue(nBL, ROOT.xAOD.numberOfBLayerHits)
        nIBL = ctypes.c_ubyte(0)
        track.summaryValue(nIBL, ROOT.xAOD.numberOfInnermostPixelLayerHits)
        nPixHole = ctypes.c_ubyte(0)
        track.summaryValue(nPixHole, ROOT.xAOD.numberOfPixelHoles)
        npix = npix.value
        nsct = nsct.value
        nBL = nBL.value
        nIBL = nIBL.value
        nPixHole = nPixHole.value

        passesTrackPreSelection = (npix >= 1 and (npix + nsct) >= 6)
        passesTightTrackPreSelection = (((npix + nsct) >= 9 and abs(track.eta()) <= 1.65) or ((npix + nsct) >= 11 and abs(track.eta()) > 1.65)) and ((nIBL + nBL) > 0 and nPixHole == 0)

        if passesTrackKinematicPreSelection and passesTrackPreSelection:
            n_charged_tracks += 1
    
    if n_charged_tracks > 0:
        continue
    cut_flow["Pass Track Veto"] += 1
    
    # Pixel track veto
    n_pixel_tracks = 0
    for track in event.InDetPixelTrackParticles:
        eta = track.eta()
        phi = track.phi()
        pt = track.pt()
        
        npix = ctypes.c_ubyte(0)
        track.summaryValue(npix, ROOT.xAOD.SummaryType.numberOfPixelHits)
        npix = npix.value
        
        if pt > 0.05 * GeV and abs(eta) < 2.5 and npix >= 3:
            for ph in event.Photons:
                ph_phi = ph.phi()
                ph_eta = ph.eta()
                ph_aeta = abs(ph_eta)
                ph_pt = ph.pt()/GeV  # Convert to GeV
                
                if ph_aeta < 2.37 and not (1.37 < ph_aeta < 1.52) and ph_pt > 0.5  * GeV:
                    if abs(ph_eta - eta) < 0.5:
                        ev_pix_track_eta.append(eta)
                        ev_pix_track_phi.append(phi)
                        ev_pix_track_pt.append(pt)
                        n_pixel_tracks += 1
                        break  # Save only one match per track
    
    if n_pixel_tracks > 0:
        continue
    cut_flow["Pass Pixel Track Veto"] += 1
    
    selected_photons = []
    # Photon selection
    for ph in event.Photons:
        ph_pt = ph.pt()
        ph_eta = abs(ph.caloCluster().etaBE(2))                    
        ph_Author = ph.author()
        ph_isGoodPhoton = ph.isGoodOQ(ROOT.xAOD.EgammaParameters.BADCLUSPHOTON)

        if ph_isGoodPhoton and ph_pt > 2 * GeV and ph_eta < 2.47 and (ph_eta > 1.52 or ph_eta < 1.37):
            selected_photons.append(ph)

  
    if len(selected_photons) < 2:
        continue
    cut_flow["Pass Photon Selection"] += 1
    
    # Diphoton selection (leading two photons)
    selected_photons.sort(key=lambda x: x.pt(), reverse=True)
    leading_photon, subleading_photon = selected_photons[:2]
    
    diphoton_mass = invariant_mass(leading_photon, subleading_photon)  # Convert to GeV
    diphoton_pt = (leading_photon.pt() + subleading_photon.pt())  # Convert to GeV
    
    if diphoton_mass < 5 * GeV:
        continue
    cut_flow["Pass Diphoton Selection"] += 1
    
    diphoton_masses.append(diphoton_mass * GeV)
    diphoton_pts.append(diphoton_pt * GeV)

# Print cut flow summary
for cut, count in cut_flow.items():
    print(f"{cut}: {count}")

# Plot histograms
plt.figure()
plt.hist(diphoton_masses, bins=50, histtype='step', linewidth=2)
plt.xlabel("Diphoton Mass (GeV)")
plt.ylabel("Counts")
plt.title("Diphoton Mass Distribution")
plt.savefig("diphoton_mass_distribution.png")
plt.close()

plt.figure()
plt.hist(diphoton_pts, bins=50, histtype='step', linewidth=2)
plt.xlabel("Diphoton pT (GeV)")
plt.ylabel("Counts")
plt.title("Diphoton pT Distribution")
plt.savefig("diphoton_pt_distribution.png")
plt.close()