Repair an inherited analysis script

{
  "n_events": 50000,
  "n_in_window": 34483,
  "mean_mass_in_window": 90.96866737789078,
  "median_mass_all": 90.38002673106806,
  "n_opposite_charge_in_window": 34205,
  "bugs_fixed": [
    "Column names were outdated and needed renaming",
    "Four-vector calculation was incorrect (pz was calculated using sin(eta) instead of sinh(eta))",
    "Charge columns were not properly referenced in the opposite charge calculation"
  ]
}

import json
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

MU = 0.1056583755  # GeV

# Read the data file
df = pd.read_csv("data/synth_zmumu.csv")

# Fix column names to match the current ntuple format
df.rename(columns={
    'mu1_pt': 'pt1',
    'mu1_eta': 'eta1',
    'mu1_phi': 'phi1',
    'mu1_charge': 'q1',
    'mu2_pt': 'pt2',
    'mu2_eta': 'eta2',
    'mu2_phi': 'phi2',
    'mu2_charge': 'q2'
}, inplace=True)

# Calculate four-momenta for each muon
def four_vector(pt, eta, phi, charge):
    px = pt * np.cos(phi)
    py = pt * np.sin(phi)
    pz = pt * np.sinh(eta)
    e = np.sqrt(px**2 + py**2 + pz**2 + MU**2)
    return e, px, py, pz

e1, px1, py1, pz1 = four_vector(df["pt1"], df["eta1"], df["phi1"], df["q1"])
e2, px2, py2, pz2 = four_vector(df["pt2"], df["eta2"], df["phi2"], df["q2"])

# Calculate invariant mass squared
m2 = (e1 + e2) ** 2 - (px1 + px2) ** 2 - (py1 + py2) ** 2 - (pz1 + pz2) ** 2
m = np.sqrt(np.clip(m2, 0, None))

# Save mass data
mass_df = pd.DataFrame({
    "event_id": df["event_id"],
    "m_ll": m
})
mass_df.to_csv("mass.csv", index=False)

# Calculate analysis results
win = (m > 80) & (m < 100)
result = {
    "n_events": int(len(df)),
    "n_in_window": int(win.sum()),
    "mean_mass_in_window": float(m[win].mean()),
    "median_mass_all": float(np.median(m)),
    "n_opposite_charge_in_window": int((win & (df["q1"] * df["q2"] < 0)).sum()),
    "bugs_fixed": [
        "Column names were outdated and needed renaming",
        "Four-vector calculation was incorrect (pz was calculated using sin(eta) instead of sinh(eta))",
        "Charge columns were not properly referenced in the opposite charge calculation"
    ]
}

# Save result
json.dump(result, open("result.json", "w"), indent=2)

# Plot invariant mass distribution
try:
    fig, ax = plt.subplots()
    ax.hist(m, bins=160, range=(40, 200), edgecolor='black')
    ax.set_xlabel("m_ll [GeV]")
    ax.set_ylabel("Events")
    fig.savefig("figures/m_ll.svg")
except Exception as e:
    print(f"Error plotting: {str(e)}")

print("done", result)