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 not updated to match the current ntuple format",
    "Fixed calculation of pz using sinh(eta) instead of sin(eta)",
    "Added charge information for both muons"
  ]
}

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 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 not updated to match the current ntuple format",
        "Fixed calculation of pz using sinh(eta) instead of sin(eta)",
        "Added charge information for both muons"
    ]
}

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

# Plot histogram
try:
    fig, ax = plt.subplots()
    ax.hist(m, bins=160, range=(40, 200))
    ax.set_xlabel("m_ll [GeV]")
    ax.set_ylabel("Events")
    fig.savefig("figures/m_ll.svg")
except Exception:
    pass  # plotting is optional

print("done", result)