Reconstruct the dimuon invariant mass from muon kinematics

{"n_events": 50000, "n_in_window": 34483, "mean_mass_in_window": 90.96866737789078, "median_mass_all": 90.38002673106806, "n_opposite_charge_in_window": 34205}

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os

# Read the data card
data_path = 'data/synth_zmumu.csv'
df = pd.read_csv(data_path)

# Constants from the data card
muon_mass = 0.1056583755  # GeV

# Function to compute invariant mass
def compute_invariant_mass(pt1, eta1, phi1, pt2, eta2, phi2):
    # Convert to momentum components
    px1 = pt1 * np.cos(phi1)
    py1 = pt1 * np.sin(phi1)
    pz1 = pt1 * np.sinh(eta1)
    E1 = np.sqrt(px1**2 + py1**2 + pz1**2 + muon_mass**2)
    
    px2 = pt2 * np.cos(phi2)
    py2 = pt2 * np.sin(phi2)
    pz2 = pt2 * np.sinh(eta2)
    E2 = np.sqrt(px2**2 + py2**2 + pz2**2 + muon_mass**2)
    
    # Invariant mass squared
    m_ll_squared = (E1 + E2)**2 - (px1 + px2)**2 - (py1 + py2)**2 - (pz1 + pz2)**2
    m_ll = np.sqrt(m_ll_squared)
    return m_ll

# Compute invariant mass for each event
df['m_ll'] = df.apply(
    lambda row: compute_invariant_mass(
        row['mu1_pt'], row['mu1_eta'], row['mu1_phi'],
        row['mu2_pt'], row['mu2_eta'], row['mu2_phi']
    ),
    axis=1
)

# Save mass.csv
mass_csv_path = 'mass.csv'
df[['event_id', 'm_ll']].to_csv(mass_csv_path, index=False)

# Calculate required statistics
n_events = len(df)
n_in_window = ((df['m_ll'] > 80) & (df['m_ll'] < 100)).sum()
mean_mass_in_window = df[(df['m_ll'] > 80) & (df['m_ll'] < 100)]['m_ll'].mean()
median_mass_all = df['m_ll'].median()
n_opposite_charge_in_window = (
    (df['mu1_charge'] * df['mu2_charge'] < 0) &
    (df['m_ll'] > 80) &
    (df['m_ll'] < 100)
).sum()

# Save result.json
result_json_path = 'result.json'
result = {
    'n_events': int(n_events),
    'n_in_window': int(n_in_window),
    'mean_mass_in_window': float(mean_mass_in_window),
    'median_mass_all': float(median_mass_all),
    'n_opposite_charge_in_window': int(n_opposite_charge_in_window)
}
with open(result_json_path, 'w') as f:
    import json
    json.dump(result, f)

# Plot histogram of m_ll
plt.figure(figsize=(10, 6))
plt.hist(df['m_ll'], bins=100, range=(40, 200), color='blue', alpha=0.7, edgecolor='black')
plt.xlabel('Invariant Mass (GeV)')
plt.ylabel('Number of Events')
plt.title('Distribution of Dimuon Invariant Mass')
plt.grid(axis='y', alpha=0.5)
plt.savefig('figures/m_ll.svg')

# Ensure directories exist
os.makedirs('figures', exist_ok=True)

# Save solution.py
solution_py_path = 'solution.py'
with open(solution_py_path, 'w') as f:
    f.write('''import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os

# Read the data card
data_path = 'data/synth_zmumu.csv'
df = pd.read_csv(data_path)

# Constants from the data card
muon_mass = 0.1056583755  # GeV

# Function to compute invariant mass
def compute_invariant_mass(pt1, eta1, phi1, pt2, eta2, phi2):
    # Convert to momentum components
    px1 = pt1 * np.cos(phi1)
    py1 = pt1 * np.sin(phi1)
    pz1 = pt1 * np.sinh(eta1)
    E1 = np.sqrt(px1**2 + py1**2 + pz1**2 + muon_mass**2)
    
    px2 = pt2 * np.cos(phi2)
    py2 = pt2 * np.sin(phi2)
    pz2 = pt2 * np.sinh(eta2)
    E2 = np.sqrt(px2**2 + py2**2 + pz2**2 + muon_mass**2)
    
    # Invariant mass squared
    m_ll_squared = (E1 + E2)**2 - (px1 + px2)**2 - (py1 + py2)**2 - (pz1 + pz2)**2
    m_ll = np.sqrt(m_ll_squared)
    return m_ll

# Compute invariant mass for each event
df['m_ll'] = df.apply(
    lambda row: compute_invariant_mass(
        row['mu1_pt'], row['mu1_eta'], row['mu1_phi'],
        row['mu2_pt'], row['mu2_eta'], row['mu2_phi']
    ),
    axis=1
)

# Save mass.csv
mass_csv_path = 'mass.csv'
df[['event_id', 'm_ll']].to_csv(mass_csv_path, index=False)

# Calculate required statistics
n_events = len(df)
n_in_window = ((df['m_ll'] > 80) & (df['m_ll'] < 100)).sum()
mean_mass_in_window = df[(df['m_ll'] > 80) & (df['m_ll'] < 100)]['m_ll'].mean()
median_mass_all = df['m_ll'].median()
n_opposite_charge_in_window = (
    (df['mu1_charge'] * df['mu2_charge'] < 0) &
    (df['m_ll'] > 80) &
    (df['m_ll'] < 100)
).sum()

# Save result.json
result_json_path = 'result.json'
result = {
    'n_events': int(n_events),
    'n_in_window': int(n_in_window),
    'mean_mass_in_window': float(mean_mass_in_window),
    'median_mass_all': float(median_mass_all),
    'n_opposite_charge_in_window': int(n_opposite_charge_in_window)
}
with open(result_json_path, 'w') as f:
    import json
    json.dump(result, f)

# Plot histogram of m_ll
plt.figure(figsize=(10, 6))
plt.hist(df['m_ll'], bins=100, range=(40, 200), color='blue', alpha=0.7, edgecolor='black')
plt.xlabel('Invariant Mass (GeV)')
plt.ylabel('Number of Events')
plt.title('Distribution of Dimuon Invariant Mass')
plt.grid(axis='y', alpha=0.5)
plt.savefig('figures/m_ll.svg')

# Ensure directories exist
os.makedirs('figures', exist_ok=True)
''')

# Submit the results
print("Deliverables generated successfully.")
print("Assumptions made:")
print("- The muon mass is correctly given as 0.1056583755 GeV.")
print("- The data file contains no missing values and is properly formatted.")
print("- The invariant mass calculation uses the correct kinematic formulas.")