From 3e6a6efb07a893a068b33473fddbd071732d9b24 Mon Sep 17 00:00:00 2001
From: Sebouh Paul <sebouh.paul@gmail.com>
Date: Mon, 26 Aug 2024 14:51:40 -0400
Subject: [PATCH] Neutron in insert (#24)

* updated Snakefile

* added benchmark for neutron recon in ZDC

* added neutron_plots.py, updated Snakefile

* Update .gitlab-ci.yml

include the neutron benchmark in the ones running in gitlab

* Update config.yml

use correct generated momenta in parallelized jobs

* Update config.yml

used the correct configuration (craterlake) in the config.yml

* Update neutron_plots.py

* Update config.yml

* Update config.yml

used "touch" to force snakemake to not update files in the neutron:analyze phase
---
 .gitlab-ci.yml                                |   2 +
 Snakefile                                     |   1 +
 benchmarks/neutron/Snakefile                  |  58 ++++++
 benchmarks/neutron/analysis/gen_particles.cxx | 127 +++++++++++++
 benchmarks/neutron/analysis/neutron_plots.py  | 175 ++++++++++++++++++
 benchmarks/neutron/config.yml                 |  33 ++++
 6 files changed, 396 insertions(+)
 create mode 100644 benchmarks/neutron/Snakefile
 create mode 100644 benchmarks/neutron/analysis/gen_particles.cxx
 create mode 100644 benchmarks/neutron/analysis/neutron_plots.py
 create mode 100644 benchmarks/neutron/config.yml

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index ae8c71fc..3544c549 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -120,6 +120,7 @@ include:
     #- local: 'benchmarks/dvmp/config.yml'
   - local: 'benchmarks/dvcs/config.yml'
   - local: 'benchmarks/lambda/config.yml'
+  - local: 'benchmarks/neutron/config.yml'
   - local: 'benchmarks/sigma/config.yml'
   - local: 'benchmarks/tcs/config.yml'
   - local: 'benchmarks/u_omega/config.yml'
@@ -134,6 +135,7 @@ summary:
     - "dis:results"
     - "dvcs:results"
     - "lambda:results"
+    - "neutron:results"
     - "sigma:results"
     - "tcs:results"
     - "u_omega:results"
diff --git a/Snakefile b/Snakefile
index c11b5879..10998fe9 100644
--- a/Snakefile
+++ b/Snakefile
@@ -43,5 +43,6 @@ ddsim \
 include: "benchmarks/diffractive_vm/Snakefile"
 include: "benchmarks/dis/Snakefile"
 include: "benchmarks/lambda/Snakefile"
+include: "benchmarks/neutron/Snakefile"
 include: "benchmarks/demp/Snakefile"
 include: "benchmarks/sigma/Snakefile"
\ No newline at end of file
diff --git a/benchmarks/neutron/Snakefile b/benchmarks/neutron/Snakefile
new file mode 100644
index 00000000..3f24d738
--- /dev/null
+++ b/benchmarks/neutron/Snakefile
@@ -0,0 +1,58 @@
+rule neutron_generate:
+	input:
+                script="benchmarks/neutron/analysis/gen_particles.cxx",
+	params:
+		NEVENTS_GEN=1000,
+		th_max=5.7,
+		th_min=2.0
+	output:
+		GEN_FILE="sim_output/neutron/neutron_{P}GeV.hepmc"
+	shell:
+		"""
+mkdir -p sim_output/neutron
+root -l -b -q '{input.script}({params.NEVENTS_GEN},"{output.GEN_FILE}", "neutron", {params.th_min}, {params.th_max}, 0., 360., {wildcards.P})'
+"""
+
+rule neutron_simulate:
+	input:
+		GEN_FILE="sim_output/neutron/neutron_{P}GeV.hepmc"
+	params:
+		PHYSICS_LIST="FTFP_BERT"
+	output:
+		SIM_FILE="sim_output/neutron/{DETECTOR_CONFIG}_sim_neutron_{P}GeV.edm4hep.root"
+	shell:
+		"""
+NEVENTS_SIM=1000
+# Running simulation
+npsim \
+   --compactFile $DETECTOR_PATH/{wildcards.DETECTOR_CONFIG}.xml \
+   --numberOfEvents $NEVENTS_SIM \
+   --physicsList {params.PHYSICS_LIST} \
+   --inputFiles {input.GEN_FILE} \
+   --outputFile {output.SIM_FILE}
+"""
+
+rule neutron_recon:
+        input:
+                SIM_FILE="sim_output/neutron/{DETECTOR_CONFIG}_sim_neutron_{P}GeV.edm4hep.root"
+        output:
+                REC_FILE="sim_output/neutron/{DETECTOR_CONFIG}_rec_neutron_{P}GeV.edm4hep.root"
+        shell:
+                """
+NEVENTS_REC=1000
+eicrecon {input.SIM_FILE} -Ppodio:output_file={output.REC_FILE} -Pdd4hep:xml_files=$DETECTOR_PATH/{wildcards.DETECTOR_CONFIG}.xml -Ppodio:output_collections=MCParticles,HcalEndcapPInsertRecHits,HcalEndcapPInsertClusters,HcalEndcapPInsertSubcellHits,EcalEndcapPInsertRecHits,EcalEndcapPInsertClusters  -Pjana:nevents=$NEVENTS_REC
+"""
+
+rule neutron_analysis:
+	input:
+                expand("sim_output/neutron/{DETECTOR_CONFIG}_rec_neutron_{P}GeV.edm4hep.root",
+		    P=[20, 30, 40, 50, 60, 70, 80],
+		    DETECTOR_CONFIG=["{DETECTOR_CONFIG}"]),
+                script="benchmarks/neutron/analysis/neutron_plots.py",
+	output:
+		results_dir=directory("results/{DETECTOR_CONFIG}/neutron"),
+	shell:
+		"""
+mkdir -p {output.results_dir}
+python {input.script} {output.results_dir}
+"""
diff --git a/benchmarks/neutron/analysis/gen_particles.cxx b/benchmarks/neutron/analysis/gen_particles.cxx
new file mode 100644
index 00000000..0877a852
--- /dev/null
+++ b/benchmarks/neutron/analysis/gen_particles.cxx
@@ -0,0 +1,127 @@
+#include "HepMC3/GenEvent.h"
+#include "HepMC3/ReaderAscii.h"
+#include "HepMC3/WriterAscii.h"
+#include "HepMC3/Print.h"
+
+#include "TRandom3.h"
+#include "TVector3.h"
+
+#include <iostream>
+#include <random>
+#include <cmath>
+#include <math.h>
+#include <TMath.h>
+#include <TDatabasePDG.h>
+#include <TParticlePDG.h>
+
+using namespace HepMC3;
+
+// Generate single electron as input to the Insert simulation.
+// --
+// We generate events with a constant polar angle with respect to
+// the proton direction and then rotate so that the events are given
+// in normal lab coordinate system
+// --
+void gen_particles(
+                    int n_events = 1000, 
+                    const char* out_fname = "gen_particles.hepmc", 
+                    TString particle_name = "e-",
+                    double th_min = 3., // Minimum polar angle, in degrees
+		    double th_max = 3., // Maximum polar angle, in degrees
+		    double phi_min = 0., // Minimum azimuthal angle, in degrees
+                    double phi_max = 360., // Maximum azimuthal angle, in degrees
+                    double p = 10.,  // Momentum in GeV/c
+		    int dist = 0,  //Momentum distribution: 0=fixed, 1=uniform, 2=Gaussian
+		    int useCrossingAngle=1  // 0= no rotation, 1 = -25 mrad
+                  )
+{ 
+  WriterAscii hepmc_output(out_fname);
+  int events_parsed = 0;
+  GenEvent evt(Units::GEV, Units::MM);
+
+  // Random number generator
+  TRandom3 *r1 = new TRandom3(0); //Use time as random seed
+  
+  // Getting generated particle information
+  TDatabasePDG *pdg = new TDatabasePDG();
+  TParticlePDG *particle = pdg->GetParticle(particle_name);
+  const double mass = particle->Mass();
+  const int pdgID = particle->PdgCode();
+
+  for (events_parsed = 0; events_parsed < n_events; events_parsed++) {
+
+    //Set the event number
+    evt.set_event_number(events_parsed);
+
+    // FourVector(px,py,pz,e,pdgid,status)
+    // type 4 is beam
+    // pdgid 11 - electron
+    // pdgid 111 - pi0
+    // pdgid 2212 - proton
+    GenParticlePtr p1 =
+        std::make_shared<GenParticle>(FourVector(0.0, 0.0, 10.0, 10.0), 11, 4);
+    GenParticlePtr p2 = std::make_shared<GenParticle>(
+        FourVector(0.0, 0.0, 0.0, 0.938), 2212, 4);
+
+    // Define momentum with respect to proton direction
+    double phi   = r1->Uniform(phi_min*TMath::DegToRad(),phi_max*TMath::DegToRad());
+    double th    = r1->Uniform(th_min*TMath::DegToRad(),th_max*TMath::DegToRad());
+
+    //Total momentum distribution
+    double pevent = -1;
+    if(dist==0){ //fixed
+	pevent = p;
+    }
+    else if(dist==1){ //Uniform: +-50% variation
+	pevent = p*(1. + r1->Uniform(-0.5,0.5) );
+    }
+    else if(dist==2){  //Gaussian: Sigma = 0.1*mean
+	while(pevent<0) //Avoid negative values
+		pevent = r1->Gaus(p,0.1*p);
+    }
+
+    double px    = pevent * std::cos(phi) * std::sin(th);
+    double py    = pevent * std::sin(phi) * std::sin(th);
+    double pz    = pevent * std::cos(th);
+    TVector3 pvec(px,py,pz); 
+
+    //Rotate to lab coordinate system
+    double cross_angle = -25./1000.*useCrossingAngle; //in Rad
+    TVector3 pbeam_dir(sin(cross_angle),0,cos(cross_angle)); //proton beam direction
+    pvec.RotateY(-pbeam_dir.Theta()); // Theta is returned positive, beam in negative X
+    // type 1 is final state
+    // pdgid 11 - electron 0.510 MeV/c^2
+    GenParticlePtr p3 = std::make_shared<GenParticle>(
+        FourVector(
+            pvec.X(), pvec.Y(), pvec.Z(),
+            sqrt(pevent*pevent + (mass * mass))),
+        pdgID, 1);
+
+    //If wanted, set non-zero vertex
+    double vx = 0.;
+    double vy = 0.;
+    double vz = 0.;
+    double vt = 0.;
+
+    GenVertexPtr v1 = std::make_shared<GenVertex>();
+    evt.shift_position_by(FourVector(vx, vy, vz, vt));
+    v1->add_particle_in(p1);
+    v1->add_particle_in(p2);
+
+    v1->add_particle_out(p3);
+    evt.add_vertex(v1);
+
+    if (events_parsed == 0) {
+      std::cout << "First event: " << std::endl;
+      Print::listing(evt);
+    }
+
+    hepmc_output.write_event(evt);
+    if (events_parsed % 100 == 0) {
+      std::cout << "Event: " << events_parsed << std::endl;
+    }
+    evt.clear();
+  }
+  hepmc_output.close();
+  std::cout << "Events parsed and written: " << events_parsed << std::endl;
+}
diff --git a/benchmarks/neutron/analysis/neutron_plots.py b/benchmarks/neutron/analysis/neutron_plots.py
new file mode 100644
index 00000000..a7b919b6
--- /dev/null
+++ b/benchmarks/neutron/analysis/neutron_plots.py
@@ -0,0 +1,175 @@
+import numpy as np, pandas as pd, matplotlib.pyplot as plt, matplotlib as mpl, awkward as ak, sys
+import mplhep as hep
+hep.style.use("CMS")
+
+plt.rcParams['figure.facecolor']='white'
+plt.rcParams['savefig.facecolor']='white'
+plt.rcParams['savefig.bbox']='tight'
+
+plt.rcParams["figure.figsize"] = (7, 7)
+config=sys.argv[1].split("/")[1]  #results/{config}/neutron
+outdir=sys.argv[1]+"/"
+try:
+    import os
+    os.mkdir(outdir[:-1])
+except:
+    pass
+
+#read files
+arrays_sim={}
+for p in 20,30,40,50,60,70,80:
+    arrays_sim[p] = ur.open(f'sim_output/neutron/{config}_rec_neutron_{p}GeV.edm4hep.root:events')\
+                    .arrays()
+
+#get reconstructed theta, eta, and E
+def gauss(x, A,mu, sigma):
+    return A * np.exp(-(x-mu)**2/(2*sigma**2))
+
+#get the truth pseudorapidity and energy
+for array in arrays_sim.values():
+    tilt=-0.025
+    px=array['MCParticles.momentum.x'][:,2]
+    py=array['MCParticles.momentum.y'][:,2]
+    pz=array['MCParticles.momentum.z'][:,2]
+    p=np.sqrt(px**2+py**2+pz**2)
+    
+    pxp=px*np.cos(tilt)-pz*np.sin(tilt)
+    pyp=py
+    pzp=pz*np.cos(tilt)+px*np.sin(tilt)
+    
+    array['E_truth']=np.hypot(p, 0.9406)
+    array['eta_truth']=1/2*np.log((p+pzp)/(p-pzp))
+    array['theta_truth']=np.arccos(pzp/p)
+
+#weighted avg of theta of cluster centers, by energy
+for array in arrays_sim.values():
+    tilt=-0.025
+    x=array['HcalEndcapPInsertClusters.position.x']
+    y=array['HcalEndcapPInsertClusters.position.y']
+    z=array['HcalEndcapPInsertClusters.position.z']
+    E=array['HcalEndcapPInsertClusters.energy']
+    r=np.sqrt(x**2+y**2+z**2)
+    
+    xp=x*np.cos(tilt)-z*np.sin(tilt)
+    yp=y
+    zp=z*np.cos(tilt)+x*np.sin(tilt)
+    
+    w=E
+    
+    array['theta_recon']=np.sum(np.arccos(zp/r)*w, axis=-1)/np.sum(w, axis=-1)
+    array['eta_recon']=-np.log(np.tan(['theta_recon']/2))
+    
+    
+    array['E_Hcal']=np.sum(array['HcalEndcapPInsertClusters.energy'], axis=-1)#*20/12.5
+    array['E_Ecal']=np.sum(array['EcalEndcapPInsertClusters.energy'], axis=-1)#*27/20
+    
+    array['E_recon']=array['E_Hcal']/(0.70-.30/np.sqrt(array['E_Hcal']))\
+                        +array['E_Ecal']/(0.82-0.35/np.sqrt(array['E_Ecal']))
+
+#plot theta residuals:
+from scipy.optimize import curve_fit
+
+fig, axs=plt.subplots(1,2, figsize=(16,8))
+plt.sca(axs[0])
+p=40
+eta_min=3.4; eta_max=3.6
+y,x,_=plt.hist(1000*(arrays_sim[p]['theta_recon']-arrays_sim[p]['theta_truth'])\
+               [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)], bins=50,
+                    range=(-10,10), histtype='step')
+bc=(x[1:]+x[:-1])/2
+slc=abs(bc)<3
+# try:
+fnc=gauss
+sigma=np.sqrt(y[slc])+(y[slc]==0)
+p0=(100, 0, 5)
+coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
+xx=np.linspace(-5,5,100)
+plt.plot(xx,fnc(xx,*coeff))
+# except:
+#     pass
+plt.xlabel("$\\theta_{rec}-\\theta_{truth}$ [mrad]")
+plt.ylabel("events")
+plt.title(f"$p={p}$ GeV, ${eta_min}<\\eta<{eta_max}$")
+
+r=[3.0,3.2,3.4,3.6,3.8]
+for eta_min, eta_max in zip(r[:-1],r[1:]):
+    xvals=[]
+    sigmas=[]
+    dsigmas=[]
+    for p in 20,30,40, 50, 60, 70, 80:
+        y,x=np.histogram(1000*(arrays_sim[p]['theta_recon']-arrays_sim[p]['theta_truth'])\
+                         [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)],
+                         bins=50, range=(-10,10))
+        bc=(x[1:]+x[:-1])/2
+        slc=abs(bc)<3
+        fnc=gauss
+        p0=(100, 0, 5)
+        #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))
+            sigmas.append(np.abs(coeff[2]))
+            dsigmas.append(np.sqrt(var_matrix[2][2]))
+            xvals.append(p)
+        except:
+            pass
+    plt.sca(axs[1])
+    plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', label=f"${eta_min}<\\eta<{eta_max}$")
+plt.xlabel("$p_{n}$ [GeV]")
+plt.ylabel("$\\sigma[\\theta]$ [mrad]")
+plt.ylim(0)
+plt.legend()
+plt.tight_layout()
+plt.savefig(outdir+"neutron_theta_recon.pdf")
+
+fig, axs=plt.subplots(1,2, figsize=(16,8))
+plt.sca(axs[0])
+p=50
+eta_min=3.4; eta_max=3.6
+y,x,_=plt.hist(((arrays_sim[p]['E_recon']-arrays_sim[p]['E_truth'])/arrays_sim[p]['E_truth'])\
+               [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)&(arrays_sim[p]['E_Hcal']>0)], bins=50,
+                    range=(-.5,.5), histtype='step')
+bc=(x[1:]+x[:-1])/2
+slc=abs(bc)<0.4
+# try:
+p0=(100, 0, 0.3)
+fnc=gauss
+sigma=np.sqrt(y[slc])+(y[slc]==0)
+coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
+xx=np.linspace(-.5,.5,100)
+plt.plot(xx,fnc(xx,*coeff))
+# except:
+#     pass
+plt.xlabel("$(E_{rec}-E_{truth})/E_{truth}$")
+plt.ylabel("events")
+plt.title(f"$p={p}$ GeV, ${eta_min}<\\eta<{eta_max}$")
+
+r=[3.0,3.2,3.4,3.6,3.8]
+for eta_min, eta_max in zip(r[:-1],r[1:]):
+    xvals=[]
+    sigmas=[]
+    dsigmas=[]
+    for p in 20,30,40, 50, 60, 70, 80:
+        y,x=np.histogram(((arrays_sim[p]['E_recon']-arrays_sim[p]['E_truth'])/arrays_sim[p]['E_truth'])\
+                       [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)], bins=50,  range=(-.5,.5))
+        bc=(x[1:]+x[:-1])/2
+        slc=abs(bc)<0.5
+        fnc=gauss
+        p0=(100, 0, 0.3)
+        #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))
+            sigmas.append(np.abs(coeff[2]))
+            dsigmas.append(np.sqrt(var_matrix[2][2]))
+            xvals.append(p)
+        except:
+            pass
+    plt.sca(axs[1])
+    plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', label=f"${eta_min}<\\eta<{eta_max}$")
+plt.xlabel("$p_{n}$ [GeV]")
+plt.ylabel("$\\sigma[E]/E$")
+plt.ylim(0)
+plt.legend()
+plt.tight_layout()
+plt.savefig(outdir+"neutron_energy_recon.pdf")
diff --git a/benchmarks/neutron/config.yml b/benchmarks/neutron/config.yml
new file mode 100644
index 00000000..be42df42
--- /dev/null
+++ b/benchmarks/neutron/config.yml
@@ -0,0 +1,33 @@
+neutron:simulate:
+  stage: simulate
+  extends: .phy_benchmark
+  needs: ["common:detector"]
+  parallel:
+    matrix:
+      - P: 20
+      - P: 30
+      - P: 40
+      - P: 50
+      - P: 60
+      - P: 70
+      - P: 80
+  timeout: 1 hours
+  script:
+    - snakemake --cores 1 sim_output/neutron/epic_craterlake_rec_neutron_${P}GeV.edm4hep.root 
+  retry:
+    max: 2
+    when:
+      - runner_system_failure
+
+neutron:analyze:
+  stage: analyze
+  extends: .phy_benchmark
+  needs: ["neutron:simulate"]
+  script:
+    - snakemake --cores 1 --touch results/epic_craterlake/neutron 
+
+neutron:results:
+  stage: collect
+  needs: ["neutron:analyze"]
+  script:
+    - ls -al
-- 
GitLab