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

Merge branch 'signal_generation' into 'master' 

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

See merge request faser/calypso!215
parents 25d6cbfd 771f81f0
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
+ 130
− 43
View file @ ae003425
......@@ -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()
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