Merge branch 'signal_generation' into 'master'

Signal generation

See merge request faser/calypso!244

Generators/FlukaReader/python/FlukaReaderAlg.py

@@ -333,8 +333,8 @@ class FlukaReader(EvgenAlg):
         plt.figure()
         ebins = np.linspace(0, 5000, 50)
         plt.xlabel("Energy")
-        plt.hist(self.before["E"], bins=ebins, histtype='step', color = "g", fill = False, label = "before")
-        plt.hist(self.after["E"], bins = ebins, histtype='step', color = "r", fill = False, label = "after")
+        plt.hist(self.before["E"], bins=ebins, weights = np.array(self.before["w"])/self.nsamples, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(self.after["E"], bins = ebins, weights = np.array(self.after["w"])/self.nsamples, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
         plt.legend()
         plt.savefig("energy.png")
@@ -344,8 +344,8 @@ class FlukaReader(EvgenAlg):
         thetaX =  np.pi/2. - np.arccos(np.array(self.before["cosX"]))
         thetaXout =  np.pi/2. - np.arccos(np.array(self.after["cosX"]))        
         tbins = np.linspace(-0.5, 0.5, 100)
-        plt.hist(thetaX, bins = tbins, histtype='step', color = "g", fill = False, label = "before")
-        plt.hist(thetaXout, bins = tbins, histtype='step', color = "r", fill = False, label = "after")
+        plt.hist(thetaX, bins = tbins, weights = np.array(self.before["w"])/self.nsamples, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(thetaXout, bins = tbins, weights = np.array(self.after["w"])/self.nsamples, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
         plt.legend()
         plt.savefig("thetaX.png")
@@ -354,8 +354,8 @@ class FlukaReader(EvgenAlg):
         plt.xlabel("Angle to beam in Y dir")        
         thetaY =  np.pi/2. - np.arccos(np.array(self.before["cosY"]))
         thetaYout =  np.pi/2. - np.arccos(np.array(self.after["cosY"]))        
-        plt.hist(thetaY, bins = tbins, histtype='step', color = "g", fill = False, label = "before")
-        plt.hist(thetaYout, bins = tbins, histtype='step', color = "r", fill = False, label = "after")
+        plt.hist(thetaY, bins = tbins, weights = np.array(self.before["w"])/self.nsamples, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(thetaYout, bins = tbins, weights = np.array(self.after["w"])/self.nsamples, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
         plt.legend()
         plt.savefig("thetaY.png")
@@ -363,16 +363,16 @@ class FlukaReader(EvgenAlg):
         plt.figure()
         plt.xlabel("Dispacement in X dir")        
         xbins = np.linspace(-300, 300, 100)
-        plt.hist(self.before["x"], bins = xbins, histtype='step', color = "g", fill = False, label = "before")
-        plt.hist(self.after["x"], bins = xbins, histtype='step', color = "r", fill = False, label = "after")
+        plt.hist(self.before["x"], bins = xbins, weights = np.array(self.before["w"])/self.nsamples, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(self.after["x"], bins = xbins, weights = np.array(self.after["w"])/self.nsamples, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
         plt.legend()
         plt.savefig("x.png")
 
         plt.figure() 
         plt.xlabel("Dispacement in Y dir")               
-        plt.hist(self.before["y"], bins = xbins, histtype='step', color = "g", fill = False, label = "before")
-        plt.hist(self.after["y"], bins = xbins, histtype='step', color = "r", fill = False, label = "after")        
+        plt.hist(self.before["y"], bins = xbins, weights = np.array(self.before["w"])/self.nsamples, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(self.after["y"], bins = xbins, weights = np.array(self.after["w"])/self.nsamples, histtype='step', color = "r", fill = False, label = "after")        
         plt.gca().set_yscale('log')
         plt.legend()
         plt.savefig("y.png")

Generators/ForeseeGenerator/python/Validate.py

@@ -179,6 +179,11 @@ class EvgenValidation(EvgenAnalysisAlg):
             if i == 0 and p.production_vertex():
                 #p.production_vertex().print()
                 llp_vtx = p.production_vertex().point3d()
+            elif i == 0 and p.end_vertex():
+                llp_vtx = p.end_vertex().point3d()                            
+
+            if p.production_vertex():
+                vertices.append(p.production_vertex().point3d())
 
             if p.production_vertex():
                 vertices.append(p.production_vertex().point3d())

Generators/ForeseeGenerator/share/generate_forsee_events.py

@@ -19,7 +19,7 @@ class ForeseeGenerator(object):
         self.daughter2_pid = daughter2_pid
         self.outdir = outdir
         self.path = path
-        self.version = 1  # Forsee "version": 2 is in testing
+        self.version = 2  # Forsee "version": 2 is in testing
         self.seed = randomSeed
 
         # Set decay mode ...
@@ -102,6 +102,12 @@ class ForeseeGenerator(object):
                 coupling_ref=1,
                 massrange=[1.5, 10.]
                 )
+
+            self.model.set_br_1d(
+                modes = [self.mode],
+                filenames=[f"files/models/{self.modelname}/br/{self.mode}.txt"] 
+                )
+
         else:
             masses_brem = [ 
                 0.01  ,  0.0126,  0.0158,  0.02  ,  0.0251,  0.0316,  0.0398,
@@ -134,6 +140,13 @@ class ForeseeGenerator(object):
                 masses = masses_dy,
                 )
 
+            self.model.set_br_1d(
+                modes = ["e_e", "mu_mu"],
+                finalstates = [[11, -11], [13, -13]],
+                filenames=["model/br/e_e.txt", "model/br/mu_mu.txt"],
+                #filenames=[f"model/br/all.txt"] 
+                )
+
         return self.decays()
 
 
@@ -175,6 +188,19 @@ class ForeseeGenerator(object):
             nsample = 10,
             )
 
+        if self.version == 1:
+            self.model.set_br_1d(
+                modes = [self.mode],
+                filenames=[f"files/models/{self.modelname}/br/{self.mode}.txt"] 
+                )
+        else:
+            self.model.set_br_1d(
+                modes = ["gamma_gamma"],
+                finalstates = [[22, 22]],
+                filenames=["model/br/gamma_gamma.txt"]
+                #filenames=[f"model/br/all.txt"] 
+                )   
+
         return self.decays()
 
 
@@ -185,24 +211,13 @@ class ForeseeGenerator(object):
             self.model.set_ctau_1d(
                 filename=f"files/models/{self.modelname}/ctau.txt", 
                 coupling_ref=1
-                )
-            
-            self.model.set_br_1d(
-                modes = [self.mode],
-                filenames=[f"files/models/{self.modelname}/br/{self.mode}.txt"] 
-                )
+                )        
         else:
             self.model.set_ctau_1d(
                 filename=f"model/ctau.txt", 
                 coupling_ref=1
                 )
             
-            self.model.set_br_1d(
-                modes = [self.mode],
-                finalstates = [[self.daughter1_pid, self.daughter2_pid]],
-                filenames=[f"model/br/{self.mode}.txt"] 
-                )            
-
         # Get LLP spectrum
         self.foresee.set_model(model=self.model)
         # This is just a reference coupling 
@@ -295,7 +310,15 @@ class ForeseeGenerator(object):
 
         filename =  f"{self.outdir}/events_{self.energy}TeV_m{self.mass}GeV_c{self.couplings[0]}to_{self.daughter1_pid}_{self.daughter2_pid}.hepmc"
 
-        self.foresee.write_events(self.mass, self.couplings[0], self.energy, filename, nevents, zfront = -1500, seed = self.seed)
+        self.foresee.write_events(self.mass, self.couplings[0], self.energy, filename, nevents, zfront = -1.5, seed = self.seed, decaychannels = [self.mode])
+
+        cfgname = f"{self.foresee.dirpath}/Models/{self.modelname}/" + filename.replace(".hepmc", ".cfg")
+        print(f"save config to file: {cfgname}")
+        with open(cfgname, "w") as f:
+            f.write(" ".join(sys.argv))
+
+        return
+
 
 def setup_foresee(path):
 
@@ -334,10 +357,10 @@ def add_to_python_path(path):
 def parse_couplings(data, write_hepMC = False):    
 
     if write_hepMC:
-        try:
-            couplings = float(couplings)
-        except ValueError:
-            sus.exit("Only a single coupling allowed when writing HEPMC events")
+        if len(data) == 1:
+            couplings = float(data[0])
+        else:
+            sys.exit("Only a single coupling allowed when writing HEPMC events")
 
     try:
         couplings = [float(d) for d in data]
@@ -375,7 +398,7 @@ if __name__ == "__main__":
     if args.pid2 is None:
         args.pid2 = -args.pid1
     
-    couplings = parse_couplings(args.couplings)
+    couplings = parse_couplings(args.couplings, args.hepmc)
 
     print(f"Generating {args.model} events at Ecom = {args.Ecom}") 
     print(f"   mother mass = {args.mass} GeV")

Generators/HEPMCReader/python/HepMCReaderConfig.py

@@ -64,7 +64,6 @@ def get_file_skip_events(ConfigFlags):
 
     return foutname
                 
-    
 def HepMCReaderCfg(ConfigFlags, **kwargs) :
     cfg = ComponentAccumulator(AthSequencer("AthBeginSeq", Sequential = True))
     from TruthIO.TruthIOConf import HepMCReadFromFile