Merge branch 'signal_generation' into 'master'

Update foresee gen and valid

See merge request faser/calypso!212

Generators/ForeseeGenerator/python/Validate.py

@@ -71,8 +71,8 @@ class EvgenValidation(EvgenAnalysisAlg):
 
     def binning(self):
         "binning for theta vs phi plot"
-        tmin, tmax, tnum = [-6, 0, 12]
-        pmin, pmax, pnum = [ 0, 5, 5]
+        tmin, tmax, tnum = [-6, 0, 24]
+        pmin, pmax, pnum = [ 0, 5, 10]
         t_edges = np.logspace(tmin, tmax, num=tnum+1)
         p_edges = np.logspace(pmin, pmax, num=pnum+1)
         return t_edges, p_edges
@@ -80,7 +80,6 @@ class EvgenValidation(EvgenAnalysisAlg):
     def initialize(self):
 
         # All daughters
-        self.hists.add("Masses", 100, 0, 0.01)
         self.hists.add("PIDs", 60, -30, 30)
 
         # Daughter i
@@ -88,22 +87,23 @@ class EvgenValidation(EvgenAnalysisAlg):
         for i in range(self.ndaughter):
             self.hists.add(f"P_d{i}", 100, 0, 10000)
             self.hists.add(f"Pt_d{i}", 100, 0, 1)
-            self.hists.add(f"Theta_d{i}", 10, 0, 0.001)
+            self.hists.add(f"Theta_d{i}", 20, 0, 0.001)
             self.hists.add(f"Phi_d{i}", 16, -3.2, 3.2)
             self.hists.add(f"ThetaVsP_d{i}", arrayX = tbins, arrayY = pbins)
+            self.hists.add(f"Mass_d{i}", 200, 0, 0.01)            
 
         # Mother
         self.hists.add("P_M", 100, 0, 10000) 
         self.hists.add("Pt_M", 100, 0, 1)       
-        self.hists.add("Theta_M", 10, 0, 0.001)
+        self.hists.add("Theta_M", 20, 0, 0.001)
         self.hists.add("Phi_M", 16, -3.2, 3.2)
-        self.hists.add("Mass_M", 100, 0, 1)
+        self.hists.add("Mass_M", 200, 0, 1)
         self.hists.add("ThetaVsP_M", arrayX = tbins, arrayY = pbins)
         
         return StatusCode.Success
 
 
-    def fillKin(self, label, p, mass = False, twoD = True):
+    def fillKin(self, label, p, mass = True, twoD = True):
 
         self.hists[f"P_{label}"].Fill(p.rho()/1000, 1)        
         self.hists[f"Pt_{label}"].Fill(p.perp()/1000, 1)
@@ -119,7 +119,6 @@ class EvgenValidation(EvgenAnalysisAlg):
         return
 
     def fillDaughter(self, p):
-        self.hists["Masses"].Fill(p.momentum().m()/1000, 1)
         self.hists["PIDs"].Fill(p.pdg_id())
         return
 
@@ -147,12 +146,13 @@ class EvgenValidation(EvgenAnalysisAlg):
     def finalize(self):
         self.hists.write(self.outname)
         return StatusCode.Success
+
         
 if __name__ == "__main__":
 
     import argparse, sys
     parser = argparse.ArgumentParser(description="Run gen-level validation")
-    parser.add_argument("file", help = "full path to imput file")
+    parser.add_argument("file", nargs="+", help = "full path to imput file")
     parser.add_argument("--ndaugthers", "-d", default = 2, type = float, help = "Number of daugthers to plot")
     parser.add_argument("--output", "-o",  default = "validation.root", help = "Name of output file")
     parser.add_argument("--mcEventKey", "-k",  default = "BeamTruthEvent", help = "Name of MC collection")
@@ -171,7 +171,7 @@ if __name__ == "__main__":
     ConfigFlags.IOVDb.GlobalTag = "OFLCOND-FASER-01"             # Always needed; must match FaserVersion
     ConfigFlags.GeoModel.FaserVersion     = "FASER-01"           # Default FASER geometry
     ConfigFlags.Detector.EnableFaserSCT = True
-    ConfigFlags.Input.Files = [args.file]
+    ConfigFlags.Input.Files = args.file
     ConfigFlags.lock()
 
     from CalypsoConfiguration.MainServicesConfig import MainServicesCfg

Generators/ForeseeGenerator/share/generate_forsee_events.py

 import os
 
 import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib
 
 from foresee import Foresee, Model, Utility
 
@@ -9,7 +11,7 @@ class ForeseeGenerator(object):
     Generate LLP particles within FASER acceptance from FORESEE
     """
     
-    def __init__(self, modelname, energy, mass, couplings, daughter1_pid, daughter2_pid, mother_pid = None):
+    def __init__(self, modelname, energy, mass, couplings, daughter1_pid, daughter2_pid, mother_pid = None, outdir = None):
 
         self.modelname = modelname
         self.energy = energy
@@ -18,6 +20,7 @@ class ForeseeGenerator(object):
         self.mother_pid = mother_pid
         self.daughter1_pid = daughter1_pid
         self.daughter2_pid = daughter2_pid
+        self.outdir = outdir
 
         #if not os.path.exists("files"):
         #    os.symlink(os.path.expandvars("$Calypso_DIR/../calypso/Generators/foresee/files"), "files")
@@ -152,30 +155,68 @@ class ForeseeGenerator(object):
         # Get LLP spectrum
         self.foresee.set_model(model=self.model)        
         plt = self.foresee.get_llp_spectrum(self.mass, coupling=1, do_plot=True)  # This is just a reference coupling 
-        #plt.savefig(f"{self.modelname}.png")
+        plt.savefig(f"{self.modelname}_m{self.mass}.png")
 
         def flatten(l):
             return [i for sublist in l for i in sublist]
 
         # Get list of events within detector
-        coups, ctaus, nsigs, energies, weights, thetas = self.foresee.get_events(mass=self.mass, energy=self.energy, couplings=self.couplings)
+        coups, ctaus, nsigs, energies, weights, thetas = self.foresee.get_events(mass=self.mass, energy=self.energy, couplings=self.couplings)        
+
+        self.plot(flatten(thetas), flatten(energies), flatten(weights))
 
         # Return energy (converting to MeV), theta and weights
         return [[e*1000 for e in flatten(energies)], flatten(thetas), flatten(weights)] 
 
+    def plot(self, thetas, energies, weights):
+        # Plot the results in Forsee format
+
+        t = np.array(thetas)
+        p = np.sqrt(np.array(energies)**2 - self.mass**2)
+
+        prange=[[-6, 0, 120],[ 0, 5, 50]]
+        tmin, tmax, tnum = prange[0]
+        pmin, pmax, pnum = prange[1]
+        t_edges = np.logspace(tmin, tmax, num=tnum+1)
+        p_edges = np.logspace(pmin, pmax, num=pnum+1)  
+
+        ticks = np.array([[np.linspace(10**(j),10**(j+1),9)] for j in range(-7,6)]).flatten()
+        ticks = [np.log10(x) for x in ticks]
+        ticklabels = np.array([[r"$10^{"+str(j)+"}$","","","","","","","",""] for j in range(-7,6)]).flatten()
+        matplotlib.rcParams.update({'font.size': 15})
+
+        fig = plt.figure(figsize=(8,5.5))
+        ax = plt.subplot(1,1,1)
+        h=ax.hist2d(x=np.log10(t),y=np.log10(p),weights=weights,
+                    bins=[tnum,pnum],range=[[tmin,tmax],[pmin,pmax]],
+                    norm=matplotlib.colors.LogNorm(), cmap="hsv",
+                    )
+        
+        fig.colorbar(h[3], ax=ax)
+        ax.set_xlabel(r"angle wrt. beam axis $\theta$ [rad]")
+        ax.set_ylabel(r"momentum $p$ [GeV]")
+        ax.set_xticks(ticks)
+        ax.set_xticklabels(ticklabels)
+        ax.set_yticks(ticks)
+        ax.set_yticklabels(ticklabels)
+        ax.set_xlim(tmin, tmax)
+        ax.set_ylim(pmin, pmax)
+        plt.savefig(f"{self.modelname}_m{self.mass}_acc.png")
+
     def write(self):
         # Write LLP results to a file
         
         energies, thetas, weights = self.data
 
-        dirname = f"files/models/{self.modelname}/events"
-        if not os.path.exists(dirname):
-            os.mkdir(dirname)
+        if self.outdir is None:
+            self.outdir = f"files/models/{self.modelname}/events"
+        if not os.path.exists(self.outdir):
+            os.mkdir(self.outdir)
 
         if len(self.couplings) == 1:
-            filename = f"{dirname}/events_{self.energy}TeV_m{self.mass}GeV_c{self.couplings[0]}to_{self.daughter1_pid}_{self.daughter2_pid}.npy"
+            filename = f"{self.outdir}/events_{self.energy}TeV_m{self.mass}GeV_c{self.couplings[0]}to_{self.daughter1_pid}_{self.daughter2_pid}.npy"
         else:
-            filename = f"{dirname}/events_{self.energy}TeV_m{self.mass}GeV_to_{self.daughter1_pid}_{self.daughter2_pid}.npy"
+            filename = f"{self.outdir}/events_{self.energy}TeV_m{self.mass}GeV_to_{self.daughter1_pid}_{self.daughter2_pid}.npy"
 
         print(f"Generated {len(thetas)} events")
         print(f"save data to file: {filename}")
@@ -238,6 +279,7 @@ if __name__ == "__main__":
     parser.add_argument("--pid1", required = True, type = int, help = "PID of daughter 1")
     parser.add_argument("--pid2", default = None, type = int, help = "PID of daughter 2 (if not set then will be -PID1)")
     parser.add_argument("--Ecom", default = "14", help = "Center of mass energy [TeV]")
+    parser.add_argument("--outdir", "-o", default = ".", help = "Output path")    
     parser.add_argument("--path", default = ".", help = "Path to foresee installation")
     args = parser.parse_args()
 

Generators/ForeseeGenerator/share/plot_validation.py

+import ROOT as R
+from collections import namedtuple
+
+Hist = namedtuple("Hist", "name, xtitle, ytitle, xlo, xhi, ylo, yhi, r, d, logx, logy, ndiv",
+                  defaults = [None, None, None, None, None, None, 1, "hist", False, False, None])
+
+def plot(f, name, xtitle, ytitle, xlo = None, xhi = None, ylo = None, yhi = None,
+         r = 1, d = "hist", logx = False, logy = False, ndiv = None):
+
+    h = f.Get(name)
+
+    if xlo is not None and xhi is not None:
+        h.SetAxisRange(xlo, xhi)
+
+    if ylo is not None and yhi is not None:
+        h.SetAxisRange(ylo, yhi, "Y")
+        
+    if r != 1:
+        h.Rebin(r)
+
+    if xtitle is not None:
+        h.GetXaxis().SetTitle(xtitle)
+
+    if ytitle is not None:
+        h.GetYaxis().SetTitle(ytitle)
+
+    if logx:
+        R.gPad.SetLogx()
+
+    if logy:
+        R.gPad.SetLogy()
+
+    if ndiv is not None:
+        h.SetNdivisions(ndiv)
+
+    h.SetLabelSize(0.05, "X")
+    h.SetTitleSize(0.05, "X")
+    h.SetLabelSize(0.05, "Y")
+    h.SetTitleSize(0.05, "Y")
+
+    h.GetXaxis().SetTitleOffset(1.2)
+
+    R.gPad.SetBottomMargin(0.15)
+    R.gPad.SetLeftMargin(0.12)
+    R.gPad.SetRightMargin(0.2)        
+
+    h.Draw(d)
+    return h
+
+def plotn(f, configs, x, y, outname = "valplot"):
+
+    c = R.TCanvas()
+    c.Divide(x, y)
+    c._objs = []
+
+    if isinstance(configs, tuple):
+        configs = [configs]
+        
+    for i, cfg in enumerate(configs):
+        c.cd(i+1)
+        c._objs.append(plot(f, *cfg))
+        
+    c.Print(f"{outname}.eps")
+
+    return
+
+if __name__ == "__main__":
+
+    R.gROOT.SetBatch(True)
+    R.gStyle.SetOptStat(0)
+
+    fname = "validation.root"
+    f = R.TFile.Open(fname)
+    
+    config = [Hist("P_d0", logy = True, xtitle = "p^{0} [GeV]", ndiv = 5, r = 5),
+              Hist("Theta_d0", xtitle = "#theta [rad]", ndiv = -4),
+              Hist("Mass_d0", xtitle = "m^{0} [GeV]", xlo = 0, xhi = 0.001, ndiv = 4),               
+              Hist("Pt_d0", logy = True, xtitle = "p_{T}^{0} [GeV]", ndiv = 10, r = 5),
+              Hist("Phi_d0", xtitle = "#phi [rad]"),
+              Hist("ThetaVsP_d0", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz")
+              ]
+
+    plotn(f, config, 3, 2, "daug0")
+
+    config = [Hist("P_d1", logy = True, xtitle = "p^{0} [GeV]", ndiv = 5, r = 5),
+              Hist("Theta_d1", xtitle = "#theta [rad]", ndiv = -4),
+              Hist("Mass_d1", xtitle = "m^{0} [GeV]", xlo = 0, xhi = 0.001, ndiv = 4),               
+              Hist("Pt_d1", logy = True, xtitle = "p_{T}^{0} [GeV]", ndiv = 10, r = 5),
+              Hist("Phi_d1", xtitle = "#phi [rad]"),
+              Hist("ThetaVsP_d1", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz")
+              ]
+
+    plotn(f, config, 3, 2, "daug1")    
+
+    config = [Hist("P_M", logy = True, xtitle = "p^{0} [GeV]", ndiv = 5, r = 5),
+              Hist("Theta_M", xtitle = "#theta [rad]", ndiv = -4),
+              Hist("Mass_M", xtitle = "m^{0} [GeV]", xlo = 0, xhi = 0.05),               
+              Hist("Pt_M", logy = True, xtitle = "p_{T}^{0} [GeV]", ndiv = 10, r = 5),
+              Hist("Phi_M", xtitle = "#phi [rad]"),
+              Hist("ThetaVsP_M", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz")
+              ]
+
+    plotn(f, config, 3, 2, "mother")
+
+    plotn(f, Hist("PIDs", xtitle="PDG Id"), 1, 1, "pid") 
+
+    config = [Hist("ThetaVsP_M", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz"),
+              Hist("ThetaVsP_d0", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz"),
+              Hist("ThetaVsP_d1", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz")
+              ]                   
+
+    plotn(f, config, 2, 2, "twod")