Skip to content
Snippets Groups Projects
Commit 91798887 authored by Carl Gwilliam's avatar Carl Gwilliam
Browse files

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

parent 06edd777
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Generators/ForeseeGenerator/share/generate_forsee_events.py
+ 121
− 41
View file @ 91798887
......@@ -11,16 +11,17 @@ 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 +29,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 +90,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 +182,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 +264,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 +288,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 +309,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.
......@@ -258,7 +329,13 @@ def setup_foresee(path):
return
def parse_couplings(data):
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")
try:
couplings = [float(d) for d in data]
......@@ -283,11 +360,10 @@ 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)
# Create PIDs
if args.pid2 is None:
args.pid2 = -args.pid1
......@@ -299,8 +375,12 @@ if __name__ == "__main__":
print(f" decay = {args.pid1} {args.pid2}")
print(f" couplings = {couplings}")
f = ForeseeGenerator(args.model, args.Ecom, args.mass, couplings, args.pid1, args.pid2)
f.write()
f = ForeseeGenerator(args.model, args.Ecom, args.mass, couplings, args.pid1, args.pid2, outdir = args.outdir, path = args.path)
if args.hepmc:
f.write_hepmc(args.nevents)
else:
f.write()
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment