code for making efficiency plots

efficiency.py

pt_true = []
matched_pt_true = []
pred_pt = []

eta_true = []
matched_eta_true = []
pred_eta = []

phi_true = []
matched_phi_true = []
pred_phi = []

nhits_true = []
matched_nhits_true = []

test_offset = 200 # should be the event number at which the test set starts
n_test = 10 # how many test events
n_tracks = 420 # should be mean number of tracks + maximum with variation.

for evt in range(n_test):
    idata = test_offset + evt
    d_gpu = data[idata].to('cuda')

    d_gpu.x = d_gpu.x[(d_gpu.y < n_tracks)] # just take the first n tracks
    d_gpu.x_scaled = (d_gpu.x - torch.min(d_gpu.x, axis=0).values)/(torch.max(d_gpu.x, axis=0).values - torch.min(d_gpu.x, axis=0).values) # Normalise
    d_gpu.y_particle_barcodes = d_gpu.y_particle_barcodes[d_gpu.y < n_tracks]
    d_gpu.y = d_gpu.y[(d_gpu.y < n_tracks)]

    coords, edge_scores, edges, cluster_map, cluster_props, cluster_batch = model(d_gpu.x_scaled)

    pred_cluster_match, y_properties = match_cluster_targets(cluster_map, d_gpu.y, d_gpu)

    first_n_tracks = data[idata].y < n_tracks

    n_clusters = data[idata].y[first_n_tracks].max().item() + 1

    color_cycle = plt.cm.coolwarm(np.linspace(0.1,0.9,n_clusters))

    print(cluster_props.size())
    print(n_clusters)
    
    track_lengths = []
    true_lengths = []

    for i in range(n_clusters):
        mapped_i = pred_cluster_match[i].item()
        r = d_gpu.x[cluster_map == mapped_i,0].cpu().detach().numpy()
        r_true = d_gpu.x[d_gpu.y == i,0].cpu().detach().numpy()
        phi = d_gpu.x[cluster_map == mapped_i,1].cpu().detach().numpy()
        z = d_gpu.x[cluster_map == mapped_i,2].cpu().detach().numpy()
    
        track_lengths.append(r.shape[0])
        true_lengths.append(r_true.shape[0])
    
        if r_true.shape[0] > 1:
            pt_true.append(1/y_properties[i,0].item())
            eta_true.append(y_properties[i,1].item())
            phi_true.append(y_properties[i,2].item())
            nhits_true.append(r_true.shape[0])
            if r.shape[0] > 1:
                matched_pt_true.append(1/y_properties[i,0].item())
                pred_pt.append(1./F.softplus(cluster_props[mapped_i,0]).item())
                matched_eta_true.append(y_properties[i,1].item())
                pred_eta.append(5.0*(2*torch.sigmoid(cluster_props[mapped_i,1]) - 1))
                matched_phi_true.append(y_properties[i,2].item())
                pred_phi.append(math.pi*(2*torch.sigmoid(cluster_props[mapped_i,2]) - 1))
                matched_nhits_true.append(r_true.shape[0])

# to get below pip install coffea
from coffea import hist
from coffea.hist import Bin,Hist

pt_bins = np.concatenate( (np.linspace(0,1,num=7)[:-1], np.linspace(1,2,num=3)[:-1], np.linspace(2,10,num=4)) )
eta_bins = np.linspace(-5,5,num=10)
phi_bins = np.linspace(-math.pi,math.pi,num=10)
nhits_bins = np.linspace(0, 17, num=17)

pt_truth = Hist("Counts", Bin("pt", r"$p_{T}$ (GeV)", pt_bins))
matched_pt_truth = Hist("Counts", Bin("pt", r"$p_{T}$ (GeV)", pt_bins))

eta_truth = Hist("Counts", Bin("eta", r"$\eta$", eta_bins))
matched_eta_truth = Hist("Counts", Bin("eta", r"$\eta$", eta_bins))

phi_truth = Hist("Counts", Bin("phi", r"$\varphi$", phi_bins))
matched_phi_truth = Hist("Counts", Bin("phi", r"$\varphi$", phi_bins))

nhits_truth = Hist("Counts", Bin("nhits", r"$N_{hits}$", nhits_bins))
matched_nhits_truth = Hist("Counts", Bin("nhits", r"$N_{hits}$", nhits_bins))

pt_true = np.array(pt_true)
matched_pt_true = np.array(matched_pt_true)

eta_true = np.array(eta_true)
matched_eta_true = np.array(matched_eta_true)

phi_true = np.array(phi_true)
matched_phi_true = np.array(matched_phi_true)

nhits_true = np.array(nhits_true)
matched_nhits_true = np.array(matched_nhits_true)

pt_truth.fill(pt=pt_true)
matched_pt_truth.fill(pt=matched_pt_true)

eta_truth.fill(eta=eta_true)
matched_eta_truth.fill(eta=matched_eta_true)

phi_truth.fill(phi=phi_true)
matched_phi_truth.fill(phi=matched_phi_true)

nhits_truth.fill(nhits=nhits_true)
matched_nhits_truth.fill(nhits=matched_nhits_true)

plt.rcParams.update({
    'font.size': 14,
    'axes.titlesize': 18,
    'axes.labelsize': 18,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12
})

fig, ax = plt.subplots(2, 2, figsize=(20, 20))

matched_pt_truth.label = r'$\epsilon$'
hist.plotratio(num=matched_pt_truth, denom=pt_truth, 
               error_opts={'marker': '.'},
               unc='clopper-pearson',
               ax=ax[0,0]
              )
ax[0,0].set_xlim(0.1, 10)
ax[0,0].set_xscale('log')
ax[0,0].set_ylim(0, 1.1)


matched_eta_truth.label = r'$\epsilon$'
hist.plotratio(num=matched_eta_truth, denom=eta_truth, 
               error_opts={'marker': '.'},
               unc='clopper-pearson',
               ax=ax[0,1]
              )
ax[0,1].set_xlim(-5, 5)
ax[0,1].set_ylim(0, 1.1)


matched_phi_truth.label = r'$\epsilon$'
hist.plotratio(num=matched_phi_truth, denom=phi_truth, 
               error_opts={'marker': '.'},
               unc='clopper-pearson',
               ax=ax[1,0]
              )
ax[1,0].set_xlim(-math.pi, math.pi)
ax[1,0].set_ylim(0, 1.1)

matched_nhits_truth.label = r'$\epsilon$'
hist.plotratio(num=matched_nhits_truth, denom=nhits_truth, 
               error_opts={'marker': '.'},
               unc='clopper-pearson',
               ax=ax[1,1]
              )
ax[1,1].set_xlim(0, 18)
ax[1,1].set_ylim(0, 1.1)