Update foresee generation to allow testing with new version. Old version default for now

Generators/ForeseeGenerator/share/generate_forsee_events.py

@@ -4,23 +4,22 @@ import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib
 
-from foresee import Foresee, Model, Utility
-
 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, outdir = None):
+    def __init__(self, modelname, energy, mass, couplings, daughter1_pid, daughter2_pid, outdir = None, path = '.'):
 
         self.modelname = modelname
         self.energy = energy
         self.mass = mass
         self.couplings = [couplings] if isinstance(couplings, (str, int, float)) else couplings
-        self.mother_pid = mother_pid
         self.daughter1_pid = daughter1_pid
         self.daughter2_pid = daughter2_pid
         self.outdir = outdir
+        self.path = path
+        self.version = 1  # Forsee "version": 2 is in testing
 
         #if not os.path.exists("files"):
         #    os.symlink(os.path.expandvars("$Calypso_DIR/../calypso/Generators/foresee/files"), "files")
@@ -28,28 +27,36 @@ class ForeseeGenerator(object):
         # Set decay mode ...
         
         self.pid_map = { 
-            (11, 11) : "e_e",           
-            (13, 13) : "mu_mu",
+            (-11, 11) : "e_e",
+            (11, -11) : "e_e",                       
+            (-13, 13) : "mu_mu",
+            (13, -13) : "mu_mu",            
             (22, 22) : "gamma_gamma",
             }
 
-        self.mode = self.pid_map.get((abs(self.daughter1_pid), abs(self.daughter2_pid)), None)
+        self.mode = self.pid_map.get((self.daughter1_pid, self.daughter2_pid), None)
         if self.mode is None:
             sys.exit(f"Undefined decay to {self.daughter1_pid} + {self.daughter2_pid} for {self.modelname}")
 
         # Set detector ...
-        self.foresee = Foresee()
+        if self.version == 1:
+            self.foresee = Foresee()
+        else:
+            self.foresee = Foresee(path = self.path)
         self.foresee.set_detector(selection="np.sqrt(x.x**2 + x.y**2)< 0.1",
                                   channels=[self.mode], distance=480, length=1.5 ,
                                   luminosity=1/1000.) # 1 pb-1        
 
         # Set model ...
-        self.model = Model(self.modelname)
+        if self.version == 1:
+            self.model = Model(self.modelname)
+        else:
+            self.model = Model(self.modelname, path = f"{self.path}/Models/{self.modelname}/")
         
         if self.modelname == "DarkPhoton":
             self.data = self.darkphoton()
         elif self.modelname == "ALP-W":
-            self.data = self.alps()
+            self.data = self.alp_W()
         else:
             sys.exit(f"Unknown model {self.modelname}")
 
@@ -81,24 +88,56 @@ class ForeseeGenerator(object):
             energy = self.energy,
             )
 
-        self.model.add_production_direct(
-            label = "Brem",
-            energy = self.energy,
-            condition = "p.pt<1",
-            coupling_ref=1,
-            )
-        
-        self.model.add_production_direct(
-            label = "DY",
-            energy = self.energy,
-            coupling_ref=1,
-            massrange=[1.5, 10.]
-            )
+        if self.version == 1:
+            self.model.add_production_direct(
+                label = "Brem",
+                energy = self.energy,
+                condition = "p.pt<1",
+                coupling_ref=1,
+                )
+            
+            self.model.add_production_direct(
+                label = "DY",
+                energy = self.energy,
+                coupling_ref=1,
+                massrange=[1.5, 10.]
+                )
+        else:
+            masses_brem = [ 
+                0.01  ,  0.0126,  0.0158,  0.02  ,  0.0251,  0.0316,  0.0398,
+                0.0501,  0.0631,  0.0794,  0.1   ,  0.1122,  0.1259,  0.1413,
+                0.1585,  0.1778,  0.1995,  0.2239,  0.2512,  0.2818,  0.3162,
+                0.3548,  0.3981,  0.4467,  0.5012,  0.5623,  0.6026,  0.631 ,
+                0.6457,  0.6607,  0.6761,  0.6918,  0.7079,  0.7244,  0.7413,
+                0.7586,  0.7762,  0.7943,  0.8128,  0.8318,  0.8511,  0.871 ,
+                0.8913,  0.912 ,  0.9333,  0.955 ,  0.9772,  1.    ,  1.122 ,
+                1.2589,  1.4125,  1.5849,  1.7783,  1.9953,  2.2387,  2.5119,
+                2.8184,  3.1623,  3.9811,  5.0119,  6.3096,  7.9433, 10.    
+                ]
+            
+            self.model.add_production_direct(
+                label = "Brem",
+                energy = self.energy,
+                condition = "p.pt<1",
+                coupling_ref=1,
+                masses = masses_brem,
+                )
+            
+            masses_dy = [ 
+                1.5849,  1.7783,  1.9953,  2.2387,  2.5119, 2.8184,  3.1623,  3.9811,  5.0119,  6.3096,  7.9433, 10.    
+                ]
+            
+            self.model.add_production_direct(
+                label = "DY",
+                energy = self.energy,
+                coupling_ref=1,
+                masses = masses_dy,
+                )
 
         return self.decays()
 
 
-    def alps(self):
+    def alp_W(self):
 
         self.model.add_production_2bodydecay(
             pid0 = "5",
@@ -141,16 +180,28 @@ class ForeseeGenerator(object):
 
     def decays(self):
         # Set up liftime and BRs
-        
-        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"] 
-            )
+        if self.version == 1:
+            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)
@@ -211,7 +262,11 @@ class ForeseeGenerator(object):
         energies, thetas, weights = self.data
 
         if self.outdir is None:
-            self.outdir = f"files/models/{self.modelname}/events"
+            if self.version == 1:
+                self.outdir = f"files/models/{self.modelname}/events"
+            else:
+                self.outdir = f"{self.foresee.dirpath}/Models/{self.modelname}/model/events"
+
         if not os.path.exists(self.outdir):
             os.mkdir(self.outdir)
 
@@ -231,6 +286,18 @@ class ForeseeGenerator(object):
 
         return
 
+    def write_hepmc(self, nevents):
+
+        if self.outdir is None:
+            self.outdir = f"{self.foresee.dirpath}/Models/{self.modelname}/model/events"
+
+        if not os.path.exists(self.outdir):
+            os.mkdir(self.outdir)
+
+        filename =  "test.hepmc" #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)
+
 def setup_foresee(path):
 
     if path is None:
@@ -240,9 +307,11 @@ def setup_foresee(path):
     path = os.path.expandvars(os.path.expanduser(path))
     os.sys.path.append(f"{path}/FORESEE/src")
 
-    # Symlink foresee files dir to current dir
-    if not os.path.exists("files"):
-        os.symlink(os.path.expandvars(f"{path}/FORESEE/files"), "files")
+    # Symlink foresee files/Models dirs to current dir
+    #if not os.path.exists("files"):
+    #    os.symlink(os.path.expandvars(f"{path}/FORESEE/files"), "files")
+    #if not os.path.exists("Models"):
+    #    os.symlink(os.path.expandvars(f"{path}/FORESEE/Models"), "files")        
 
     # Install scikit-hep if needed.
 
@@ -257,8 +326,19 @@ def setup_foresee(path):
         
     return
 
+def add_to_python_path(path):
+    if path in sys.path: return
+    path = os.path.expandvars(os.path.expanduser(path))
+    os.sys.path.append(path)
+    return
+
+def parse_couplings(data, write_hepMC = False):    
 
-def parse_couplings(data):    
+    if write_hepMC:
+        try:
+            couplings = float(couplings)
+        except ValueError:
+            sus.exit("Only a single coupling allowed when writing HEPMC events")
 
     try:
         couplings = [float(d) for d in data]
@@ -283,10 +363,13 @@ if __name__ == "__main__":
     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")
+    parser.add_argument("--hepmc", action = "store_true", help = "Write HepMC events")
+    parser.add_argument("--nevents", "-n", default = 10, type = int, help = "Number of HepMC events ")    
     args = parser.parse_args()
 
-    setup_foresee(args.path)
-    
+    add_to_python_path(f"{args.path}/src")
+
+    from foresee import Foresee, Model, Utility
 
     # Create PIDs
     if args.pid2 is None:
@@ -297,10 +380,14 @@ if __name__ == "__main__":
     print(f"Generating {args.model} events at Ecom = {args.Ecom}") 
     print(f"   mother mass = {args.mass} GeV")
     print(f"   decay = {args.pid1} {args.pid2}")
-    print(f"   couplings = {couplings}")
+    print(f"   couplings = {couplings}")    
+
+    f = ForeseeGenerator(args.model, args.Ecom, args.mass, couplings, args.pid1, args.pid2, outdir = args.outdir, path = args.path)
 
-    f = ForeseeGenerator(args.model, args.Ecom, args.mass, couplings, args.pid1, args.pid2)
-    f.write()
+    if args.hepmc:
+        f.write_hepmc(args.nevents)
+    else:
+        f.write()