End-to-end Z-peak measurement on CMS open data

{"n_total": 100000, "n_selected": 5752, "efficiency": 0.05752, "n_in_range": 4656, "fit": {"mu": 90.79632906047802, "mu_err": 344.12147036281044, "sigma": 2.0314895353387765, "sigma_err": 421.46099473220096, "n_sig": 1876.0141975531728, "n_sig_err": 0.5269934736613361, "n_bkg": -31.175324893126287, "lam": -307769.4174918953, "lam_err": 7.289127555701092e-08, "chi2": 2225.5697635218603, "ndf": 35}}

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import least_squares
from scipy.stats import norm
import json
import os

# Read the data
data = pd.read_csv('data/cms_dimuon_2011.csv')

# Step 1: Quality selection
selected = data[
    (data['type1'] == 'G') & 
    (data['type2'] == 'G') &
    (np.abs(data['eta1']) < 2.4) &
    (np.abs(data['eta2']) < 2.4) &
    (data['pt1'] > 20) &
    (data['pt2'] > 20) &
    (data['Q1'] * data['Q2'] < 0)
]

n_total = len(data)
n_selected = len(selected)

# Calculate efficiency
efficiency = n_selected / n_total if n_total > 0 else 0.0

# Step 2: Histogram M
M = selected['M']
hist, bin_edges = np.histogram(M, bins=40, range=(80, 100), density=False)
bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2

# Step 3: Fit the histogram
def fit_function(params, x, hist, bin_width):
    mu, sigma, n_sig, n_bkg, lam = params
    # Normalized Gaussian
    gauss = norm.pdf(x, loc=mu, scale=sigma)
    # Normalized Exponential
    exp_term = np.exp(-x / lam)
    exp_norm = (exp_term - np.exp(-80 / lam)) / (np.exp(-80 / lam) - np.exp(-100 / lam))
    exp = exp_norm
    # Model
    model = n_sig * gauss + n_bkg * exp
    return model

def objective(params, x, hist, bin_width):
    model = fit_function(params, x, hist, bin_width)
    return (model - hist)

# Initial parameters
initial_mu = 91
initial_sigma = 2
initial_n_sig = 0.9 * len(M[(M >= 80) & (M <= 100)])
initial_n_bkg = 0.1 * len(M[(M >= 80) & (M <= 100)])
initial_lam = 30
initial_params = [initial_mu, initial_sigma, initial_n_sig, initial_n_bkg, initial_lam]

# Fit using least squares
result = least_squares(objective, initial_params, args=(bin_centers, hist, 0.5))

# Extract fitted parameters
mu_fit = result.x[0]
sigma_fit = result.x[1]
n_sig_fit = result.x[2]
n_bkg_fit = result.x[3]
lam_fit = result.x[4]

# Compute chi2 and ndf
chi2 = np.sum((hist - fit_function(result.x, bin_centers, hist, 0.5)) ** 2 / np.sqrt(hist))
ndf = len(hist) - len(result.x)

# Compute uncertainties from covariance matrix
cov_matrix = result.jac.T @ result.jac
std_devs = np.sqrt(np.diag(cov_matrix))
mu_err = std_devs[0]
sigma_err = std_devs[1]
n_sig_err = std_devs[2]
n_bkg_err = std_devs[3]
lam_err = std_devs[4]

# Step 4: Save results
result_data = {
    "n_total": int(n_total),
    "n_selected": int(n_selected),
    "efficiency": float(efficiency),
    "n_in_range": int(len(M[(M >= 80) & (M <= 100)])),
    "fit": {
        "mu": float(mu_fit),
        "mu_err": float(mu_err),
        "sigma": float(sigma_fit),
        "sigma_err": float(sigma_err),
        "n_sig": float(n_sig_fit),
        "n_sig_err": float(n_sig_err),
        "n_bkg": float(n_bkg_fit),
        "lam": float(lam_fit),
        "lam_err": float(lam_err),
        "chi2": float(chi2),
        "ndf": int(ndf)
    }
}

with open('result.json', 'w') as f:
    json.dump(result_data, f)

# Plot the data and fit
plt.figure(figsize=(10, 6))
plt.errorbar(bin_centers, hist, yerr=np.sqrt(hist), fmt='o', label='Data', color='blue')
x_fit = np.linspace(80, 100, 1000)
y_fit = fit_function(result.x, x_fit, hist, 0.5)
plt.plot(x_fit, y_fit, label='Fit', color='red')
plt.xlabel('Invariant Mass (GeV)')
plt.ylabel('Counts')
plt.legend()
plt.title('Z-peak Measurement')
plt.savefig('figures/zfit.svg')

# Write REPORT.md
report_content = """
## Selection
We applied the following quality cuts to select events:
- Both muons are global (`type1 == "G"` and `type2 == "G"`)
- |eta1| < 2.4 and |eta2| < 2.4
- pt1 > 20 GeV and pt2 > 20 GeV
- Opposite charge (Q1*Q2 < 0)

The total number of events is {n_total}, and after applying these cuts, we have {n_selected} selected events with an efficiency of {efficiency:.4f}.

## Fit
We fitted the invariant mass distribution in the range [80, 100] GeV using a combination of a Gaussian for signal and an exponential for background. The fit parameters are:
- Mean (mu): {mu_fit:.2f} ± {mu_err:.2f} GeV
- Sigma (sigma): {sigma_fit:.2f} ± {sigma_err:.2f} GeV
- Signal yield (n_sig): {n_sig_fit:.2f} ± {n_sig_err:.2f}
- Background yield (n_bkg): {n_bkg_fit:.2f} ± {n_bkg_err:.2f}
- Background parameter (lam): {lam_fit:.2f} ± {lam_err:.2f}
- Chi-squared (chi2): {chi2:.2f} with {ndf} degrees of freedom

## Assumptions
- The signal is modeled as a Gaussian distribution.
- The background is modeled as an exponential distribution.
- The fit uses a least squares minimization with uncertainties calculated from the covariance matrix.
- The initial parameters for the fit were chosen based on physical expectations.
""".format(
    n_total=n_total,
    n_selected=n_selected,
    efficiency=efficiency,
    mu_fit=mu_fit,
    mu_err=mu_err,
    sigma_fit=sigma_fit,
    sigma_err=sigma_err,
    n_sig_fit=n_sig_fit,
    n_sig_err=n_sig_err,
    n_bkg_fit=n_bkg_fit,
    n_bkg_err=n_bkg_err,
    lam_fit=lam_fit,
    lam_err=lam_err,
    chi2=chi2,
    ndf=ndf
)

with open('REPORT.md', 'w') as f:
    f.write(report_content)

# Save solution.py
with open('solution.py', 'w') as f:
    f.write('''import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import least_squares
from scipy.stats import norm
import json
import os

# Read the data
data = pd.read_csv('data/cms_dimuon_2011.csv')

# Step 1: Quality selection
selected = data[
    (data['type1'] == 'G') & 
    (data['type2'] == 'G') &
    (np.abs(data['eta1']) < 2.4) &
    (np.abs(data['eta2']) < 2.4) &
    (data['pt1'] > 20) &
    (data['pt2'] > 20) &
    (data['Q1'] * data['Q2'] < 0)
]

n_total = len(data)
n_selected = len(selected)

# Calculate efficiency
efficiency = n_selected / n_total if n_total > 0 else 0.0

# Step 2: Histogram M
M = selected['M']
hist, bin_edges = np.histogram(M, bins=40, range=(80, 100), density=False)
bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2

# Step 3: Fit the histogram
def fit_function(params, x, hist, bin_width):
    mu, sigma, n_sig, n_bkg, lam = params
    # Normalized Gaussian
    gauss = norm.pdf(x, loc=mu, scale=sigma)
    # Normalized Exponential
    exp_term = np.exp(-x / lam)
    exp_norm = (exp_term - np.exp(-80 / lam)) / (np.exp(-80 / lam) - np.exp(-100 / lam))
    exp = exp_norm
    # Model
    model = n_sig * gauss + n_bkg * exp
    return model

def objective(params, x, hist, bin_width):
    model = fit_function(params, x, hist, bin_width)
    return (model - hist)

# Initial parameters
initial_mu = 91
initial_sigma = 2
initial_n_sig = 0.9 * len(M[(M >= 80) & (M <= 100)])
initial_n_bkg = 0.1 * len(M[(M >= 80) & (M <= 100)])
initial_lam = 30
initial_params = [initial_mu, initial_sigma, initial_n_sig, initial_n_bkg, initial_lam]

# Fit using least squares
result = least_squares(objective, initial_params, args=(bin_centers, hist, 0.5))

# Extract fitted parameters
mu_fit = result.x[0]
sigma_fit = result.x[1]
n_sig_fit = result.x[2]
n_bkg_fit = result.x[3]
lam_fit = result.x[4]

# Compute chi2 and ndf
chi2 = np.sum((hist - fit_function(result.x, bin_centers, hist, 0.5)) ** 2 / np.sqrt(hist))
ndf = len(hist) - len(result.x)

# Compute uncertainties from covariance matrix
cov_matrix = result.jac.T @ result.jac
std_devs = np.sqrt(np.diag(cov_matrix))
mu_err = std_devs[0]
sigma_err = std_devs[1]
n_sig_err = std_devs[2]
n_bkg_err = std_devs[3]
lam_err = std_devs[4]

# Step 4: Save results
result_data = {
    "n_total": int(n_total),
    "n_selected": int(n_selected),
    "efficiency": float(efficiency),
    "n_in_range": int(len(M[(M >= 80) & (M <= 100)])),
    "fit": {
        "mu": float(mu_fit),
        "mu_err": float(mu_err),
        "sigma": float(sigma_fit),
        "sigma_err": float(sigma_err),
        "n_sig": float(n_sig_fit),
        "n_sig_err": float(n_sig_err),
        "n_bkg": float(n_bkg_fit),
        "lam": float(lam_fit),
        "lam_err": float(lam_err),
        "chi2": float(chi2),
        "ndf": int(ndf)
    }
}

with open('result.json', 'w') as f:
    json.dump(result_data, f)

# Plot the data and fit
plt.figure(figsize=(10, 6))
plt.errorbar(bin_centers, hist, yerr=np.sqrt(hist), fmt='o', label='Data', color='blue')
x_fit = np.linspace(80, 100, 1000)
y_fit = fit_function(result.x, x_fit, hist, 0.5)
plt.plot(x_fit, y_fit, label='Fit', color='red')
plt.xlabel('Invariant Mass (GeV)')
plt.ylabel('Counts')
plt.legend()
plt.title('Z-peak Measurement')
plt.savefig('figures/zfit.svg')

# Write REPORT.md
report_content = """
## Selection
We applied the following quality cuts to select events:
- Both muons are global (`type1 == "G"` and `type2 == "G"`)
- |eta1| < 2.4 and |eta2| < 2.4
- pt1 > 20 GeV and pt2 > 20 GeV
- Opposite charge (Q1*Q2 < 0)

The total number of events is {n_total}, and after applying these cuts, we have {n_selected} selected events with an efficiency of {efficiency:.4f}.

## Fit
We fitted the invariant mass distribution in the range [80, 100] GeV using a combination of a Gaussian for signal and an exponential for background. The fit parameters are:
- Mean (mu): {mu_fit:.2f} ± {mu_err:.2f} GeV
- Sigma (sigma): {sigma_fit:.2f} ± {sigma_err:.2f} GeV
- Signal yield (n_sig): {n_sig_fit:.2f} ± {n_sig_err:.2f}
- Background yield (n_bkg): {n_bkg_fit:.2f} ± {n_bkg_err:.2f}
- Background parameter (lam): {lam_fit:.2f} ± {lam_err:.2f}
- Chi-squared (chi2): {chi2:.2f} with {ndf} degrees of freedom

## Assumptions
- The signal is modeled as a Gaussian distribution.
- The background is modeled as an exponential distribution.
- The fit uses a least squares minimization with uncertainties calculated from the covariance matrix.
- The initial parameters for the fit were chosen based on physical expectations.
""".format(
    n_total=n_total,
    n_selected=n_selected,
    efficiency=efficiency,
    mu_fit=mu_fit,
    mu_err=mu_err,
    sigma_fit=sigma_fit,
    sigma_err=sigma_err,
    n_sig_fit=n_sig_fit,
    n_sig_err=n_sig_err,
    n_bkg_fit=n_bkg_fit,
    n_bkg_err=n_bkg_err,
    lam_fit=lam_fit,
    lam_err=lam_err,
    chi2=chi2,
    ndf=ndf
)

with open('REPORT.md', 'w') as f:
    f.write(report_content)
''')