diff --git a/Generators/ForeseeGenerator/share/generate_forsee_events.py b/Generators/ForeseeGenerator/share/generate_forsee_events.py
index ba63087bf08bf1e2571865651f255cd9ad177084..bd44eb2faa0f97ae5d8278f0300028221fe190e4 100644
--- a/Generators/ForeseeGenerator/share/generate_forsee_events.py
+++ b/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()