From c88b5542cedd08479240bb4a6d71b92d57a83484 Mon Sep 17 00:00:00 2001
From: Dmitry Kalinkin <dmitry.kalinkin@gmail.com>
Date: Mon, 3 Mar 2025 02:44:45 -0500
Subject: [PATCH] zdc_{lambda,photon}: fix logic, run more code conditionally
 (#141)

---
 .../zdc_lambda/analysis/lambda_plots.py       | 242 +++++++++---------
 .../zdc_photon/analysis/zdc_photon_plots.py   |   2 +-
 2 files changed, 122 insertions(+), 122 deletions(-)

diff --git a/benchmarks/zdc_lambda/analysis/lambda_plots.py b/benchmarks/zdc_lambda/analysis/lambda_plots.py
index a083cc8..f81de0b 100644
--- a/benchmarks/zdc_lambda/analysis/lambda_plots.py
+++ b/benchmarks/zdc_lambda/analysis/lambda_plots.py
@@ -138,138 +138,138 @@ else:
     plt.savefig(outdir+"thetastar_recon.pdf")
     #plt.show()
 
-#vtx z
-fig,axs=plt.subplots(1,3, figsize=(24, 8))
-plt.sca(axs[0])
-plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
-x=[]
-y=[]
-for p in momenta:
-    x+=list(ak.flatten(arrays_sim[p]['MCParticles.vertex.z'][:,3]+0*zvtx_recon[p])/1000)
-    y+=list(ak.flatten(zvtx_recon[p])/1000)
-plt.scatter(x,y)
-#print(x,y)
-plt.xlabel("$z^{\\rm truth}_{\\rm vtx}$ [m]")
-plt.ylabel("$z^{\\rm recon}_{\\rm vtx}$  [m]")
-plt.xlim(0,40)
-plt.ylim(0,40)
-
-plt.sca(axs[1])
-plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
-y,x,_=plt.hist(y-np.array(x), bins=50, range=(-10,10))
-bc=(x[1:]+x[:-1])/2
-
-from scipy.optimize import curve_fit
-slc=abs(bc)<5
-fnc=gauss
-p0=[100, 0, 1]
-coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
-                                 sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
-x=np.linspace(-5, 5)
-plt.plot(x, gauss(x, *coeff), color='tab:orange')
-print(coeff[2], np.sqrt(var_matrix[2][2]))
-plt.xlabel("$z^{*\\rm recon}_{\\rm vtx}-z^{*\\rm truth}_{\\rm vtx}$ [m]")
-plt.ylabel("events")
-
-plt.sca(axs[2])
-sigmas=[]
-dsigmas=[]
-xvals=[]
-for p in momenta:
-    
+    #vtx z
+    fig,axs=plt.subplots(1,3, figsize=(24, 8))
+    plt.sca(axs[0])
+    plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
+    x=[]
+    y=[]
+    for p in momenta:
+        x+=list(ak.flatten(arrays_sim[p]['MCParticles.vertex.z'][:,3]+0*zvtx_recon[p])/1000)
+        y+=list(ak.flatten(zvtx_recon[p])/1000)
+    plt.scatter(x,y)
+    #print(x,y)
+    plt.xlabel("$z^{\\rm truth}_{\\rm vtx}$ [m]")
+    plt.ylabel("$z^{\\rm recon}_{\\rm vtx}$  [m]")
+    plt.xlim(0,40)
+    plt.ylim(0,40)
 
-    a=ak.flatten((zvtx_recon[p]-arrays_sim[p]['MCParticles.vertex.z'][:,3])/1000)
-    y,x=np.histogram(a, bins=100, range=(-10,10))
+    plt.sca(axs[1])
+    plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
+    y,x,_=plt.hist(y-np.array(x), bins=50, range=(-10,10))
     bc=(x[1:]+x[:-1])/2
 
     from scipy.optimize import curve_fit
     slc=abs(bc)<5
     fnc=gauss
-    p0=(100, 0, 1)
-    #print(bc[slc],y[slc])
-    sigma=np.sqrt(y[slc])+(y[slc]==0)
-    try:
-        coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
-        sigmas.append(coeff[2])
-        dsigmas.append(np.sqrt(var_matrix[2][2]))
-        xvals.append(p)
-    except:
-        print("fit failed")
-plt.ylim(0, 2)
-
-plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
-x=np.linspace(100, 275, 100)
-
-avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
-plt.axhline(avg, color='tab:orange')
-plt.text(150, 1.25,f"$\\sigma\\approx${avg:.1f} m")
-
-plt.xlabel("$E_{\\Lambda}$ [GeV]")
-plt.ylabel("$\\sigma[z_{\\rm vtx}]$ [m]")
-plt.tight_layout()
-plt.savefig(outdir+"zvtx_recon.pdf")
-#plt.show()
-
-p=100
-fig,axs=plt.subplots(1,2, figsize=(16, 8))
-plt.sca(axs[0])
-lambda_mass=1.115683
-vals=[]
-for p in momenta:
-    vals+=list(ak.flatten(mass_recon[p]))
-
-y,x,_= plt.hist(vals, bins=100, range=(1.0, 1.25))
-bc=(x[1:]+x[:-1])/2
-plt.axvline(lambda_mass, ls='--', color='tab:green', lw=3)
-plt.text(lambda_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
-plt.xlabel("$m_{\\Lambda}^{\\rm recon}$ [GeV]")
-plt.ylim(0, np.max(y)*1.2)
-plt.xlim(1.0, 1.25)
-
-from scipy.optimize import curve_fit
-slc=abs(bc-lambda_mass)<0.05
-fnc=gauss
-p0=[100, lambda_mass, 0.04]
-coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
-                                 sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
-x=np.linspace(0.8, 1.3, 200)
-plt.plot(x, gauss(x, *coeff), color='tab:orange')
-print(coeff[2], np.sqrt(var_matrix[2][2]))
-plt.xlabel("$m^{\\rm recon}_{\\Lambda}$ [GeV]")
-plt.ylabel("events")
-plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
-
-plt.sca(axs[1])
-xvals=[]
-sigmas=[]
-dsigmas=[]
-for p in momenta:
-    y,x= np.histogram(ak.flatten(mass_recon[p]), bins=100, range=(0.6,1.4))
+    p0=[100, 0, 1]
+    coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
+                                     sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
+    x=np.linspace(-5, 5)
+    plt.plot(x, gauss(x, *coeff), color='tab:orange')
+    print(coeff[2], np.sqrt(var_matrix[2][2]))
+    plt.xlabel("$z^{*\\rm recon}_{\\rm vtx}-z^{*\\rm truth}_{\\rm vtx}$ [m]")
+    plt.ylabel("events")
+
+    plt.sca(axs[2])
+    sigmas=[]
+    dsigmas=[]
+    xvals=[]
+    for p in momenta:
+
+
+        a=ak.flatten((zvtx_recon[p]-arrays_sim[p]['MCParticles.vertex.z'][:,3])/1000)
+        y,x=np.histogram(a, bins=100, range=(-10,10))
+        bc=(x[1:]+x[:-1])/2
+
+        from scipy.optimize import curve_fit
+        slc=abs(bc)<5
+        fnc=gauss
+        p0=(100, 0, 1)
+        #print(bc[slc],y[slc])
+        sigma=np.sqrt(y[slc])+(y[slc]==0)
+        try:
+            coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
+            sigmas.append(coeff[2])
+            dsigmas.append(np.sqrt(var_matrix[2][2]))
+            xvals.append(p)
+        except:
+            print("fit failed")
+    plt.ylim(0, 2)
+
+    plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
+    x=np.linspace(100, 275, 100)
+
+    avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
+    plt.axhline(avg, color='tab:orange')
+    plt.text(150, 1.25,f"$\\sigma\\approx${avg:.1f} m")
+
+    plt.xlabel("$E_{\\Lambda}$ [GeV]")
+    plt.ylabel("$\\sigma[z_{\\rm vtx}]$ [m]")
+    plt.tight_layout()
+    plt.savefig(outdir+"zvtx_recon.pdf")
+    #plt.show()
+
+    p=100
+    fig,axs=plt.subplots(1,2, figsize=(16, 8))
+    plt.sca(axs[0])
+    lambda_mass=1.115683
+    vals=[]
+    for p in momenta:
+        vals+=list(ak.flatten(mass_recon[p]))
+
+    y,x,_= plt.hist(vals, bins=100, range=(1.0, 1.25))
     bc=(x[1:]+x[:-1])/2
+    plt.axvline(lambda_mass, ls='--', color='tab:green', lw=3)
+    plt.text(lambda_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
+    plt.xlabel("$m_{\\Lambda}^{\\rm recon}$ [GeV]")
+    plt.ylim(0, np.max(y)*1.2)
+    plt.xlim(1.0, 1.25)
 
     from scipy.optimize import curve_fit
     slc=abs(bc-lambda_mass)<0.05
     fnc=gauss
-    p0=[100, lambda_mass, 0.05]
-    try:
-        coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
-                                       sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
-        x=np.linspace(0.8, 1.3, 200)
-        sigmas.append(coeff[2])
-        dsigmas.append(np.sqrt(var_matrix[2][2]))
-        xvals.append(p)
-    except:
-        print("fit failed")
-    
-plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
-avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
-plt.axhline(avg, color='tab:orange')
-plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
-plt.xlabel("$E_{\\Lambda}$ [GeV]")
-plt.ylabel("$\\sigma[m_{\\Lambda}]$ [GeV]")
-plt.ylim(0, 0.02)
-plt.tight_layout()
-plt.savefig(outdir+"lambda_mass_rec.pdf")
+    p0=[100, lambda_mass, 0.04]
+    coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
+                                     sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
+    x=np.linspace(0.8, 1.3, 200)
+    plt.plot(x, gauss(x, *coeff), color='tab:orange')
+    print(coeff[2], np.sqrt(var_matrix[2][2]))
+    plt.xlabel("$m^{\\rm recon}_{\\Lambda}$ [GeV]")
+    plt.ylabel("events")
+    plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
+
+    plt.sca(axs[1])
+    xvals=[]
+    sigmas=[]
+    dsigmas=[]
+    for p in momenta:
+        y,x= np.histogram(ak.flatten(mass_recon[p]), bins=100, range=(0.6,1.4))
+        bc=(x[1:]+x[:-1])/2
+
+        from scipy.optimize import curve_fit
+        slc=abs(bc-lambda_mass)<0.05
+        fnc=gauss
+        p0=[100, lambda_mass, 0.05]
+        try:
+            coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
+                                           sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
+            x=np.linspace(0.8, 1.3, 200)
+            sigmas.append(coeff[2])
+            dsigmas.append(np.sqrt(var_matrix[2][2]))
+            xvals.append(p)
+        except:
+            print("fit failed")
+
+    plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
+    avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
+    plt.axhline(avg, color='tab:orange')
+    plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
+    plt.xlabel("$E_{\\Lambda}$ [GeV]")
+    plt.ylabel("$\\sigma[m_{\\Lambda}]$ [GeV]")
+    plt.ylim(0, 0.02)
+    plt.tight_layout()
+    plt.savefig(outdir+"lambda_mass_rec.pdf")
 
 
 #now for the CM stuff:
diff --git a/benchmarks/zdc_photon/analysis/zdc_photon_plots.py b/benchmarks/zdc_photon/analysis/zdc_photon_plots.py
index bd6abf3..a89f795 100644
--- a/benchmarks/zdc_photon/analysis/zdc_photon_plots.py
+++ b/benchmarks/zdc_photon/analysis/zdc_photon_plots.py
@@ -28,7 +28,7 @@ for p in momenta:
         for index in range(5)
     })
 
-if "ReconstructedFarForwardZDCNeutrals.PDG" in arrays_sim[p][momenta[0]].fields:
+if "ReconstructedFarForwardZDCNeutrals.PDG" not in arrays_sim[p][momenta[0]].fields:
     print("ReconstructedFarForwardZDCNeutrals collection is not available (needs EICrecon 1.23)")
     import sys
     sys.exit(0)
-- 
GitLab