Calculating cosTheta from TLotentzVector

I used MadGraph to simulate Z->e+e-, and then delphes to get a ROOT file. I am now using uproot in Python to plot dsigma/dcosTheta where the cosTheta is the angle, in the correspondent Z rest frame, between the electron direction and the Z direction in the lab frame. In order to do this, I am using ROOT’s TLorentzVector module. My issue is that my cosTheta calculation is not giving me the correct values, and I do not understand what is going wrong. Some perspective would be greatly appreciated.

import uproot
import numpy as np
import ROOT

# Open the ROOT file
file = uproot.open("delphes_ZZunpol.root")
tree = file["Delphes"]

# Access the 'Particle' data
particle_data = tree.arrays(["Particle.PID", "Particle.Px", "Particle.Py", "Particle.Pz", "Particle.E", "Event.Weight"], library="np")

PID = np.array(particle_data["Particle.PID"])
Px = np.array(particle_data["Particle.Px"])
Py = np.array(particle_data["Particle.Py"])
Pz = np.array(particle_data["Particle.Pz"])
E = np.array(particle_data["Particle.E"])
event_weights = np.array(particle_data["Event.Weight"])

# Define particle PIDs
Z_PID = 23
e_minus_PID = 11
e_plus_PID = -11

# Define mass of the electron
m_e = 0.00051099895

cos_theta = []
weighted_dcs = []

N_mad = len(event_weights)

for i in range(len(PID)):
    weight = event_weights[i][0]

    Z = PID[i] == Z_PID
    
    # Construct 4-momentum for the two Z bosons
    Z1_4mom = ROOT.TLorentzVector(Px[i][Z][0], Py[i][Z][0], Pz[i][Z][0], E[i][Z][0])
    Z2_4mom = ROOT.TLorentzVector(Px[i][Z][1], Py[i][Z][1], Pz[i][Z][1], E[i][Z][1])
    
    # Compute invariant mass of ZZ system
    M_ZZ = (Z1_4mom + Z2_4mom).M()
    if M_ZZ < 200:
        continue

    # Transverse momentum cut
    Z_px, Z_py, Z_pz, Z_E = Px[i][Z][0], Py[i][Z][0], Pz[i][Z][0], E[i][Z][0]
    pT_Ze = np.sqrt(Z_px**2 + Z_py**2)
    if not (200 < pT_Ze < 400):
        continue
        
    # Get electron four-momentum
    e_minus = PID[i] == e_minus_PID
    e_px, e_py, e_pz = Px[i][e_minus][0], Py[i][e_minus][0], Pz[i][e_minus][0]
    
    # Compute energy of electron with nonzero electron mass
    e_E = np.sqrt(e_px**2 + e_py**2 + e_pz**2 + m_e**2)

    # Define electron 4-momentum
    electron = ROOT.TLorentzVector(e_px, e_py, e_pz, e_E)
    
    # Define Z boson 4-momentum
    Z_boson = ROOT.TLorentzVector(Z_px, Z_py, Z_pz, Z_E)

    # Boost electron to Z-boson rest frame
    electron_Z_rest = ROOT.TLorentzVector(electron)
    electron_Z_rest.Boost(-Z_boson.BoostVector())

    # Compute cos(theta) in the Z rest frame
    e_px_rest = electron_Z_rest.Px()
    e_py_rest = electron_Z_rest.Py()
    e_pz_rest = electron_Z_rest.Pz()
    
    electron_3mom = np.array([e_px_rest, e_py_rest, e_pz_rest])
    Z_3mom = np.array([Z_px, Z_py, Z_pz])
    cos_theta_val = np.dot(electron_3mom, Z_3mom) / (np.linalg.norm(electron_3mom) * np.linalg.norm(Z_3mom))
    cos_theta.append(cos_theta_val)

ROOT Version: 6.30/02
Python Version: 3.9.12
Platform: Mac

Welcome to the ROOT Forum!
Maybe @moneta can help with this