Small fixes to Forsee generation

Generators/FaserParticleGun/python/FaserParticleGunConfig.py

@@ -177,6 +177,7 @@ def FaserParticleGunForeseeCfg(ConfigFlags, **kwargs) :
         mother_mass = kwargs.get("mother_mass", 0.01),
         daughter1_pid = kwargs.get("daughter1_pid", 11),
         daughter2_pid = kwargs.get("daughter2_pid", -11),
+        randomSeed = kwargs.get("randomSeed", None)        
         )
     
     pg.sampler.mother_sampler = PG.ParticleSampler(pid=kwargs.get("mother_pid", None),mom=mother_mom,n=1,pos=pg.sampler.mother_sampler.pos)

Generators/ForeseeGenerator/python/ForeseeSampler.py

+import os, sys
+import random
 import numpy as np
 
 import ParticleGun as PG
 
 class ForeseeNumpySampler(PG.MomSampler):
     """ Sample from the output of Foresee generation in numpy format with columns E, theta, weight"""
-    def __init__(self, model_path = ".", model_name = "DarkPhoton",  com_energy = "14", mother_mass = 0.01,  coupling = None, daughter1_pid = 11, daughter2_pid = -11):
+    def __init__(self, model_path = ".", model_name = "DarkPhoton",  com_energy = "14", mother_mass = 0.01,
+                 coupling = None, daughter1_pid = 11, daughter2_pid = -11, randomSeed = None):
 
-        self.path = model_path
+        self.path = os.path.expanduser(os.path.expandvars(model_path))
         self.modelname = model_name
         self.energy = com_energy
         self._mass = mother_mass
@@ -14,7 +17,12 @@ class ForeseeNumpySampler(PG.MomSampler):
         self.daughter1_pid = daughter1_pid
         self.daughter2_pid = daughter2_pid
 
-        self.rng = np.random.default_rng()
+        if randomSeed is not None:
+            print(f"Setting seed to {randomSeed}")
+            self.rng = np.random.default_rng(randomSeed)
+        else:
+            self.rng = np.random.default_rng()        
+
         self.xs = 0
 
         self.read()
@@ -101,6 +109,7 @@ class ForeseeSampler(PG.MomSampler):
         if self.mode is None:
             sys.exit(f"Undefined decay to {self.daughter1_pid} + {self.daughter2_pid} for {self.modelname}")
 
+        from foresee import Foresee, Model, Utility
         self.foresee = Foresee()
         self.foresee.set_detector(selection="np.sqrt(x.x**2 + x.y**2)< 0.1", channels=[self.mode], distance=480, length=1.5 , luminosity=150)        
 
@@ -250,7 +259,6 @@ if __name__ == "__main__":
 
     # Testing ...
 
-    import os
     from math import sqrt, log10
     import matplotlib.pyplot as plt
     import matplotlib

Generators/ForeseeGenerator/share/generate_forsee_events.py

@@ -177,9 +177,15 @@ class ForeseeGenerator(object):
         else:
             filename = f"{dirname}/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}")
         np.save(filename,[energies,thetas, weights])
 
+        cfgname = filename.replace(".npy", ".cfg")
+        print(f"save config to file: {cfgname}")
+        with open(cfgname, "w") as f:
+            f.write(" ".join(sys.argv))
+
         return
 
 def setup_foresee(path):
@@ -230,7 +236,7 @@ if __name__ == "__main__":
     parser.add_argument("--mass", "-m", required = True, type = float, help = "Mass of mother [GeV]")
     parser.add_argument("--couplings", "-c", required = True, nargs = "+", help = "Couplings of mother (either single/mulitple values or tuple to pass to np.logspace)")
     parser.add_argument("--pid1", required = True, type = int, help = "PID of daughter 1")
-    parser.add_argument("--pid2", default = None, help = "PID of daughter 2 (if not set then will be -PID1)")
+    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("--path", default = ".", help = "Path to foresee installation")
     args = parser.parse_args()