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
1 merge request!215Update foresee generation to allow testing with new version. Old version default for now
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): ...@@ -11,16 +11,17 @@ class ForeseeGenerator(object):
Generate LLP particles within FASER acceptance from FORESEE 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.modelname = modelname
self.energy = energy self.energy = energy
self.mass = mass self.mass = mass
self.couplings = [couplings] if isinstance(couplings, (str, int, float)) else couplings self.couplings = [couplings] if isinstance(couplings, (str, int, float)) else couplings
self.mother_pid = mother_pid
self.daughter1_pid = daughter1_pid self.daughter1_pid = daughter1_pid
self.daughter2_pid = daughter2_pid self.daughter2_pid = daughter2_pid
self.outdir = outdir self.outdir = outdir
self.path = path
self.version = 1 # Forsee "version": 2 is in testing
#if not os.path.exists("files"): #if not os.path.exists("files"):
# os.symlink(os.path.expandvars("$Calypso_DIR/../calypso/Generators/foresee/files"), "files") # os.symlink(os.path.expandvars("$Calypso_DIR/../calypso/Generators/foresee/files"), "files")
...@@ -28,28 +29,36 @@ class ForeseeGenerator(object): ...@@ -28,28 +29,36 @@ class ForeseeGenerator(object):
# Set decay mode ... # Set decay mode ...
self.pid_map = { self.pid_map = {
(11, 11) : "e_e", (-11, 11) : "e_e",
(13, 13) : "mu_mu", (11, -11) : "e_e",
(-13, 13) : "mu_mu",
(13, -13) : "mu_mu",
(22, 22) : "gamma_gamma", (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: if self.mode is None:
sys.exit(f"Undefined decay to {self.daughter1_pid} + {self.daughter2_pid} for {self.modelname}") sys.exit(f"Undefined decay to {self.daughter1_pid} + {self.daughter2_pid} for {self.modelname}")
# Set detector ... # 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", self.foresee.set_detector(selection="np.sqrt(x.x**2 + x.y**2)< 0.1",
channels=[self.mode], distance=480, length=1.5 , channels=[self.mode], distance=480, length=1.5 ,
luminosity=1/1000.) # 1 pb-1 luminosity=1/1000.) # 1 pb-1
# Set model ... # 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": if self.modelname == "DarkPhoton":
self.data = self.darkphoton() self.data = self.darkphoton()
elif self.modelname == "ALP-W": elif self.modelname == "ALP-W":
self.data = self.alps() self.data = self.alp_W()
else: else:
sys.exit(f"Unknown model {self.modelname}") sys.exit(f"Unknown model {self.modelname}")
...@@ -81,24 +90,56 @@ class ForeseeGenerator(object): ...@@ -81,24 +90,56 @@ class ForeseeGenerator(object):
energy = self.energy, energy = self.energy,
) )
self.model.add_production_direct( if self.version == 1:
label = "Brem", self.model.add_production_direct(
energy = self.energy, label = "Brem",
condition = "p.pt<1", energy = self.energy,
coupling_ref=1, condition = "p.pt<1",
) coupling_ref=1,
)
self.model.add_production_direct(
label = "DY", self.model.add_production_direct(
energy = self.energy, label = "DY",
coupling_ref=1, energy = self.energy,
massrange=[1.5, 10.] 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() return self.decays()
def alps(self): def alp_W(self):
self.model.add_production_2bodydecay( self.model.add_production_2bodydecay(
pid0 = "5", pid0 = "5",
...@@ -141,16 +182,28 @@ class ForeseeGenerator(object): ...@@ -141,16 +182,28 @@ class ForeseeGenerator(object):
def decays(self): def decays(self):
# Set up liftime and BRs # 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( if self.version == 1:
modes = [self.mode], self.model.set_ctau_1d(
filenames=[f"files/models/{self.modelname}/br/{self.mode}.txt"] 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 # Get LLP spectrum
self.foresee.set_model(model=self.model) self.foresee.set_model(model=self.model)
...@@ -211,7 +264,11 @@ class ForeseeGenerator(object): ...@@ -211,7 +264,11 @@ class ForeseeGenerator(object):
energies, thetas, weights = self.data energies, thetas, weights = self.data
if self.outdir is None: 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): if not os.path.exists(self.outdir):
os.mkdir(self.outdir) os.mkdir(self.outdir)
...@@ -231,6 +288,18 @@ class ForeseeGenerator(object): ...@@ -231,6 +288,18 @@ class ForeseeGenerator(object):
return 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): def setup_foresee(path):
if path is None: if path is None:
...@@ -240,9 +309,11 @@ def setup_foresee(path): ...@@ -240,9 +309,11 @@ def setup_foresee(path):
path = os.path.expandvars(os.path.expanduser(path)) path = os.path.expandvars(os.path.expanduser(path))
os.sys.path.append(f"{path}/FORESEE/src") os.sys.path.append(f"{path}/FORESEE/src")
# Symlink foresee files dir to current dir # Symlink foresee files/Models dirs to current dir
if not os.path.exists("files"): #if not os.path.exists("files"):
os.symlink(os.path.expandvars(f"{path}/FORESEE/files"), "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. # Install scikit-hep if needed.
...@@ -258,7 +329,13 @@ def setup_foresee(path): ...@@ -258,7 +329,13 @@ def setup_foresee(path):
return 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: try:
couplings = [float(d) for d in data] couplings = [float(d) for d in data]
...@@ -283,11 +360,10 @@ if __name__ == "__main__": ...@@ -283,11 +360,10 @@ if __name__ == "__main__":
parser.add_argument("--Ecom", default = "14", help = "Center of mass energy [TeV]") parser.add_argument("--Ecom", default = "14", help = "Center of mass energy [TeV]")
parser.add_argument("--outdir", "-o", default = ".", help = "Output path") parser.add_argument("--outdir", "-o", default = ".", help = "Output path")
parser.add_argument("--path", default = ".", help = "Path to foresee installation") 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() args = parser.parse_args()
setup_foresee(args.path)
# Create PIDs # Create PIDs
if args.pid2 is None: if args.pid2 is None:
args.pid2 = -args.pid1 args.pid2 = -args.pid1
...@@ -299,8 +375,12 @@ if __name__ == "__main__": ...@@ -299,8 +375,12 @@ if __name__ == "__main__":
print(f" decay = {args.pid1} {args.pid2}") 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) f = ForeseeGenerator(args.model, args.Ecom, args.mass, couplings, args.pid1, args.pid2, outdir = args.outdir, path = args.path)
f.write()
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