Skip to content
Snippets Groups Projects
Commit a0cdf048 authored by barkleya's avatar barkleya Committed by Alan Basil Barkley Yeung
Browse files

Added Decay in Flight generator

parent 8c65b741
No related branches found
No related tags found
No related merge requests found
################################################################################
# Package: DIFGenerator
################################################################################
# Declare the package name:
atlas_subdir( DIFGenerator )
# Install files from the package:
atlas_install_python_modules( python/*.py POST_BUILD_CMD ${ATLAS_FLAKE8} )
# Decay In Flight Sampler
The DIFSampler models a two-body decay.
## Usage
The `runG4DecayInFlight.py` file shows a default configuration.
Options available through the `DIFSampler` constructor include:
* daughter1_pid - The pid of the first decay particle. The default valid pids are listed [here](https://gitlab.cern.ch/atlas/athena/-/blob/master/Generators/ParticleGun/python/samplers.py#L825). (default: 11)
* daughter2_pid - The pid of the second decay particle. The default valid pids are listed [here](https://gitlab.cern.ch/atlas/athena/-/blob/master/Generators/ParticleGun/python/samplers.py#L825). (default: -11)
* mother_pid - The pid of the mother particle. Only set this if you need to get the mother's pid from the mother_sampler property (default: None)
* mother_mom - A momentum sampler representing the mother particle's momentum. (default: PG.EThetaMPhiSampler(sqrt((1 * TeV)**2 + (10 * MeV)**2),0,10*MeV,0))
* mother_pos - A position sampler representing the mother particle's position. (default: PG.PosSampler(0,0,[-1500,0],0))
The pid, momentum sampler, and position sampler of the mother particle can be modified after initialization by setting the `mother_sampler` property of the `DIFSampler`. The `mother_sampler` property should be a [ParticleSampler](https://gitlab.cern.ch/atlas/athena/-/blob/master/Generators/ParticleGun/python/samplers.py#L858).
#!/usr/bin/env python
#TODO: exponential decay lifetime?
import ParticleGun as PG
from math import sqrt, sin, cos, acos
from random import uniform
from AthenaCommon.SystemOfUnits import TeV, MeV
from AthenaCommon.PhysicalConstants import pi
class DIFSampler(PG.ParticleSampler):
"""
A particle gun which models 2 particle decay.
"""
class particle:
"""
A helper class to hold information about a particle.
"""
def __init__(self,pid = None,mass = None,mom=None,theta=None,phi=None,pos=None):
self.pid = pid
if mass is None:
self.mass = PG.MASSES[abs(pid)]
else:
self.mass = mass
self.mom_mag = mom
self.theta = theta
self.phi = phi
self.pos = pos
if self.mass is not None and self.mom_mag is not None:
self.E = sqrt(self.mom_mag**2 + self.mass**2)
else:
self.E = None
if self.theta is not None and self.phi is not None and self.mass is not None and self.E is not None:
self.mom = PG.EThetaMPhiSampler(self.E,self.theta,self.mass,self.phi).shoot()
else:
self.mom = None
def __init__(self, daughter1_pid = 11, daughter2_pid = -11, mother_pid = None, mother_mom = PG.EThetaMPhiSampler(sqrt((1*TeV)**2 + (10*MeV)**2),0,10*MeV,0),mother_pos = PG.PosSampler(0,0,[-1500,0],0)):
self._mother_sampler = PG.ParticleSampler(mother_pid,mother_mom,mother_pos)
self.daughter1 = self.particle(daughter1_pid)
self.daughter2 = self.particle(daughter2_pid)
@property
def mother_sampler(self):
return self._mother_sampler
@mother_sampler.setter
def mother_sampler(self,sampler):
if isinstance(sampler,PG.ParticleSampler) or issubclass(type(sampler),PG.ParticleSampler):
self._mother_sampler = sampler
else:
raise AttributeError("DIFSampler: Trying to set mother_sampler to something other than a ParticleSampler.")
def calculate_decay_p(self,m0,m1,m2):
return sqrt(m0**4 - (2*m0**2 * m1**2) + (m1**4) - (2*m0**2 * m2**2) - (2*m1**2 *m2**2) + (m2**4)) / (2*m0)
def lorentz_transformation(self):
self.mother_mom = self.mother.mom.shoot()
Bx = self.mother_mom.Px() / self.mother_mom.E()
By = self.mother_mom.Py() / self.mother_mom.E()
Bz = self.mother_mom.Pz() / self.mother_mom.E()
self.daughter1.mom.Boost(Bx,By,Bz)
self.daughter2.mom.Boost(Bx,By,Bz)
def shoot(self):
"Mother particle decays into two daughter particles"
## Shoot daughter 1 in a random direction in CM frame
self.daughter1.phi = uniform(0,2 * pi)
## ensure cos(theta) is uniform from -1 to 1
cos_theta = uniform(-1,1)
self.daughter1.theta = acos(cos_theta)
self.daughter2.phi = self.daughter1.phi + pi
self.daughter2.theta = acos(-cos_theta)
## The magnitude of the momentum will be equal and opposite
self.mother = self.mother_sampler
p = self.calculate_decay_p(self.mother.mom.mass(),self.daughter1.mass,self.daughter2.mass)
self.daughter1.E = sqrt(self.daughter1.mass**2 + p**2)
self.daughter2.E = sqrt(self.daughter2.mass**2 + p**2)
self.daughter1.mom = PG.EThetaMPhiSampler(self.daughter1.E,self.daughter1.theta,self.daughter1.mass,self.daughter1.phi).shoot()
self.daughter2.mom = PG.EThetaMPhiSampler(self.daughter2.E,self.daughter2.theta,self.daughter2.mass,self.daughter2.phi).shoot()
## boost into lab frame
self.lorentz_transformation()
self.mother_pos = self.mother.pos.shoot()
d1 = PG.SampledParticle(self.daughter1.pid,self.daughter1.mom,self.mother_pos)
d2 = PG.SampledParticle(self.daughter2.pid,self.daughter2.mom,self.mother_pos)
return [d1,d2]
if __name__ == "__main__":
import matplotlib.pyplot as plt
import sys
show_graphs = False
if len(sys.argv) > 1 and sys.argv[1] == "-g":
show_graphs = True
epsilon = 10**-6 #account for round off error
## test motionless mother particle
mother_pid = 321 #K+
mother_mass = PG.MASSES[mother_pid]
daughter0_pid = 211 #pi+
daughter1_pid = 111 #pi0
## check uniformifty of angles
if show_graphs:
cos_thetas = [[],[]]
thetas = [[],[]]
phis = [[],[]]
n = 100
for i in range(n):
try:
DIFS = DIFSampler(daughter0_pid,daughter1_pid, mother_pid) # K+ -> pi+ and pi0
DIFS.mother_sampler.mom = PG.EThetaMPhiSampler(mother_mass,0,mother_mass,0)
daughters = DIFS.shoot()
mother_mom = DIFS.mother_sampler.mom.shoot()
assert daughters[0].mom.E() + daughters[1].mom.E() == mother_mass
assert (daughters[0].mom.E() + daughters[1].mom.E())**2 - (daughters[0].mom.P() - daughters[1].mom.P())**2 == (PG.MASSES[mother_pid])**2
assert abs(daughters[0].mom.P() - daughters[1].mom.P()) <= epsilon
assert daughters[0].mom.Theta() - epsilon <= abs(pi - daughters[1].mom.Theta()) <= daughters[0].mom.Theta() + epsilon
assert abs(daughters[0].mom.Px() + daughters[1].mom.Px()) <= epsilon
assert abs(daughters[0].mom.Py() + daughters[1].mom.Py()) <= epsilon
assert abs(daughters[0].mom.Pz() + daughters[1].mom.Pz()) <= epsilon
assert daughters[0].pos.X() == daughters[1].pos.X() == DIFS.mother_pos.X()
assert daughters[0].pos.Y() == daughters[1].pos.Y() == DIFS.mother_pos.Y()
assert daughters[0].pos.Z() == daughters[1].pos.Z() == DIFS.mother_pos.Z()
if show_graphs:
cos_thetas[0].append(cos(daughters[0].mom.Theta()))
cos_thetas[1].append(cos(daughters[1].mom.Theta()))
thetas[0].append(daughters[0].mom.Theta())
thetas[1].append(daughters[1].mom.Theta())
phis[0].append(daughters[0].mom.Phi())
phis[1].append(daughters[1].mom.Phi())
except AssertionError:
print("Error on run " + str(i))
print("mother particle:")
print(" E = " + str(mother_mom.E()))
print(" M = " + str(mother_mom.M()))
print(" P = " + str(mother_mom.P()))
print(" Px = " + str(mother_mom.Px()))
print(" Py = " + str(mother_mom.Py()))
print(" Pz = " + str(mother_mom.Pz()))
print(" theta = " + str(mother_mom.Theta()))
print(" phi = " + str(mother_mom.Phi()))
print(" x = " + str(DIFS.mother_pos.X()))
print(" y = " + str(DIFS.mother_pos.Y()))
print(" z = " + str(DIFS.mother_pos.Z()))
print("daughter 0 particle:")
print(" E = " + str(daughters[0].mom.E()))
print(" M = " + str(daughters[0].mom.M()))
print(" P = " + str(daughters[0].mom.P()))
print(" Px = " + str(daughters[0].mom.Px()))
print(" Py = " + str(daughters[0].mom.Py()))
print(" Pz = " + str(daughters[0].mom.Pz()))
print(" theta = " + str(daughters[0].mom.Theta()))
print(" phi = " + str(daughters[0].mom.Phi()))
print(" x = " + str(daughters[0].pos.X()))
print(" y = " + str(daughters[0].pos.Y()))
print(" z = " + str(daughters[0].pos.Z()))
print("daughter 1 particle:")
print(" E = " + str(daughters[1].mom.E()))
print(" M = " + str(daughters[1].mom.M()))
print(" P = " + str(daughters[1].mom.P()))
print(" Px = " + str(daughters[1].mom.Px()))
print(" Py = " + str(daughters[1].mom.Py()))
print(" Pz = " + str(daughters[1].mom.Pz()))
print(" theta = " + str(daughters[1].mom.Theta()))
print(" phi = " + str(daughters[1].mom.Phi()))
print(" x = " + str(daughters[1].pos.X()))
print(" y = " + str(daughters[1].pos.Y()))
print(" z = " + str(daughters[1].pos.Z()))
raise
if show_graphs:
fig,axs = plt.subplots(2,3)
axs[0,0].hist(thetas[0])
axs[0,0].set_title("daughter 0: theta")
axs[0,1].hist(cos_thetas[1])
axs[0,1].set_title("daughter 0: cos(theta)")
axs[0,2].hist(phis[1])
axs[0,2].set_title("daughter 0: phi")
axs[1,0].hist(thetas[1])
axs[1,0].set_title("daughter 1: theta")
axs[1,1].hist(cos_thetas[0])
axs[1,1].set_title("daughter 1, cos(theta)")
axs[1,2].hist(phis[0])
axs[1,2].set_title("daughter 1, phi")
plt.show()
# test moving mother particle
mother_pid = 15 #tau
daughter0_pid = 16 #nu_tau
daughter1_pid = 211 #pi+-
n = 100
for i in range(n):
M = PG.MASSES[mother_pid]
theta = acos(uniform(-1,1))
mother_mom = PG.EThetaMPhiSampler([M,3*M],theta,M,[0,2*pi])
mother_pos = PG.PosSampler(x=0,y=0,z=[-1500,0],t=0)
DIFS = DIFSampler(daughter0_pid,daughter1_pid,mother_pid) # tau -> nu_tau and pi+-
DIFS.mother_sampler = PG.ParticleSampler(pid=mother_pid,mom=mother_mom,n=1,pos=mother_pos)
daughters = DIFS.shoot()
try:
mother_mom = DIFS.mother_mom
s = daughters[0].mom+daughters[1].mom
assert mother_mom.E() - epsilon <= s.E() <= mother_mom.E() + epsilon
assert mother_mom.P() - epsilon <= s.P()<= mother_mom.P() + epsilon
assert mother_mom.Px() - epsilon <= s.Px() <= mother_mom.Px() + epsilon
assert mother_mom.Py() - epsilon <= s.Py() <= mother_mom.Py() + epsilon
assert mother_mom.Pz() - epsilon <= s.Pz() <= mother_mom.Pz() + epsilon
assert daughters[0].pos.X() == daughters[1].pos.X() == DIFS.mother_pos.X()
assert daughters[0].pos.Y() == daughters[1].pos.Y() == DIFS.mother_pos.Y()
assert daughters[0].pos.Z() == daughters[1].pos.Z() == DIFS.mother_pos.Z()
except AssertionError:
print("Error on run " + str(i))
print("mother particle:")
print(" E = " + str(mother_mom.E()))
print(" M = " + str(mother_mom.M()))
print(" P = " + str(mother_mom.P()))
print(" Px = " + str(mother_mom.Px()))
print(" Py = " + str(mother_mom.Py()))
print(" Pz = " + str(mother_mom.Pz()))
print(" theta = " + str(mother_mom.Theta()))
print(" phi = " + str(mother_mom.Phi()))
print(" x = " + str(DIFS.mother_pos.X()))
print(" y = " + str(DIFS.mother_pos.Y()))
print(" z = " + str(DIFS.mother_pos.Z()))
print("daughter 0 particle:")
print(" E = " + str(daughters[0].mom.E()))
print(" M = " + str(daughters[0].mom.M()))
print(" P = " + str(daughters[0].mom.P()))
print(" Px = " + str(daughters[0].mom.Px()))
print(" Py = " + str(daughters[0].mom.Py()))
print(" Pz = " + str(daughters[0].mom.Pz()))
print(" theta = " + str(daughters[0].mom.Theta()))
print(" phi = " + str(daughters[0].mom.Phi()))
print(" x = " + str(daughters[0].pos.X()))
print(" y = " + str(daughters[0].pos.Y()))
print(" z = " + str(daughters[0].pos.Z()))
print("daughter 1 particle:")
print(" E = " + str(daughters[1].mom.E()))
print(" M = " + str(daughters[1].mom.M()))
print(" P = " + str(daughters[1].mom.P()))
print(" Px = " + str(daughters[1].mom.Px()))
print(" Py = " + str(daughters[1].mom.Py()))
print(" Pz = " + str(daughters[1].mom.Pz()))
print(" theta = " + str(daughters[1].mom.Theta()))
print(" phi = " + str(daughters[1].mom.Phi()))
print(" x = " + str(daughters[1].pos.X()))
print(" y = " + str(daughters[1].pos.Y()))
print(" z = " + str(daughters[1].pos.Z()))
raise
# Test default arguments
epsilon = 2 # increase round off error because the numbers are much larger
n = 100
for i in range(n):
DIFS = DIFSampler()
daughters = DIFS.shoot()
mother_mom = DIFS.mother_mom
s = daughters[0].mom+daughters[1].mom
try:
assert DIFS.mother_sampler.pid() == None
assert mother_mom.E()**2 - mother_mom.P()**2 - epsilon <= s.E()**2 - s.P()**2 <= mother_mom.E()**2 - mother_mom.P()**2 + epsilon
assert mother_mom.Px() - epsilon <= s.Px() <= mother_mom.Px() + epsilon
assert mother_mom.Py() - epsilon <= s.Py() <= mother_mom.Py() + epsilon
assert mother_mom.Pz() - epsilon <= s.Pz() <= mother_mom.Pz() + epsilon
assert daughters[0].pos.X() == daughters[1].pos.X() == DIFS.mother_pos.X()
assert daughters[0].pos.Y() == daughters[1].pos.Y() == DIFS.mother_pos.Y()
assert daughters[0].pos.Z() == daughters[1].pos.Z() == DIFS.mother_pos.Z()
except AssertionError:
print("Error on run " + str(i))
print("mother particle:")
print(" E = " + str(mother_mom.E()))
print(" M = " + str(mother_mom.M()))
print(" P = " + str(mother_mom.P()))
print(" Px = " + str(mother_mom.Px()))
print(" Py = " + str(mother_mom.Py()))
print(" Pz = " + str(mother_mom.Pz()))
print(" theta = " + str(mother_mom.Theta()))
print(" phi = " + str(mother_mom.Phi()))
print(" x = " + str(DIFS.mother_pos.X()))
print(" y = " + str(DIFS.mother_pos.Y()))
print(" z = " + str(DIFS.mother_pos.Z()))
print("daughter 0 particle:")
print(" E = " + str(daughters[0].mom.E()))
print(" M = " + str(daughters[0].mom.M()))
print(" P = " + str(daughters[0].mom.P()))
print(" Px = " + str(daughters[0].mom.Px()))
print(" Py = " + str(daughters[0].mom.Py()))
print(" Pz = " + str(daughters[0].mom.Pz()))
print(" theta = " + str(daughters[0].mom.Theta()))
print(" phi = " + str(daughters[0].mom.Phi()))
print(" x = " + str(daughters[0].pos.X()))
print(" y = " + str(daughters[0].pos.Y()))
print(" z = " + str(daughters[0].pos.Z()))
print("daughter 1 particle:")
print(" E = " + str(daughters[1].mom.E()))
print(" M = " + str(daughters[1].mom.M()))
print(" P = " + str(daughters[1].mom.P()))
print(" Px = " + str(daughters[1].mom.Px()))
print(" Py = " + str(daughters[1].mom.Py()))
print(" Pz = " + str(daughters[1].mom.Pz()))
print(" theta = " + str(daughters[1].mom.Theta()))
print(" phi = " + str(daughters[1].mom.Phi()))
print(" x = " + str(daughters[1].pos.X()))
print(" y = " + str(daughters[1].pos.Y()))
print(" z = " + str(daughters[1].pos.Z()))
raise
# pass invalid type to mother_sampler
try:
DIFS = DIFSampler()
DIFS.mother_sampler = PG.SampledParticle()
raise AssertionError("mother_sampler set to improper type")
except AttributeError:
pass
# pass derived type to mother_sampler
class SP(PG.ParticleSampler):
pass
DIFS = DIFSampler()
DIFS.mother_sampler = SP()
print("DIFSampler: All unit tests passed")
\ No newline at end of file
# Copyright (C) 2021 CERN for the benefit of the FASER collaboration
from DIFGenerator.DIFSampler import DIFSampler
\ No newline at end of file
#!/usr/bin/env python
if __name__ == "__main__":
import os
import sys
import GaudiPython
import ParticleGun as PG
from DIFGenerator import DIFSampler
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
from AthenaCommon.AppMgr import *
from AthenaCommon.Logging import log, logging
from AthenaCommon.SystemOfUnits import TeV
from AthenaCommon.PhysicalConstants import pi
from AthenaCommon.Constants import VERBOSE, INFO
from AthenaCommon.Configurable import Configurable
from CalypsoConfiguration.AllConfigFlags import ConfigFlags
from CalypsoConfiguration.MainServicesConfig import MainServicesCfg
from AthenaPoolCnvSvc.PoolReadConfig import PoolReadCfg
from McEventSelector.McEventSelectorConfig import McEventSelectorCfg
from OutputStreamAthenaPool.OutputStreamConfig import OutputStreamCfg
from FaserGeoModel.FaserGeoModelConfig import FaserGeometryCfg
from G4FaserAlg.G4FaserAlgConfigNew import G4FaserAlgCfg
from G4FaserServices.G4FaserServicesConfigNew import G4GeometryNotifierSvcCfg
#
# Set up logging and new style config
#
log.setLevel(VERBOSE)
Configurable.configurableRun3Behavior = True
#
# Input settings (Generator file)
#
# from AthenaConfiguration.TestDefaults import defaultTestFiles
# ConfigFlags.Input.Files = defaultTestFiles.EVNT
#
# Alternatively, these must ALL be explicitly set to run without an input file
# (if missing, it will try to read metadata from a non-existent file and crash)
#
ConfigFlags.Input.Files = [""]
ConfigFlags.Input.isMC = True
ConfigFlags.Input.RunNumber = 12345
ConfigFlags.Input.Collections = [""]
ConfigFlags.Input.ProjectName = "mc19"
ConfigFlags.Common.isOnline = False
ConfigFlags.Beam.Type = "collisions"
ConfigFlags.Beam.Energy = 7*TeV # Informational, does not affect simulation
ConfigFlags.GeoModel.FaserVersion = "FASER-01" # Always needed
ConfigFlags.IOVDb.GlobalTag = "OFLCOND-XXXX-XXX-XX" # Always needed; only the OFLCOND part matters
# Workaround for bug/missing flag; unimportant otherwise
ConfigFlags.addFlag("Input.InitialTimeStamp", 0)
# Workaround to avoid problematic ISF code
ConfigFlags.GeoModel.Layout = "Development"
#
# Output settings
#
ConfigFlags.Output.HITSFileName = "DecayInFlight.HITS.pool.root"
ConfigFlags.GeoModel.GeoExportFile = "faserGeo.db" # Optional dump of geometry for browsing in vp1light
#
# Geometry-related settings
# Do not change!
#
ConfigFlags.Detector.SimulateVeto = True
ConfigFlags.Detector.GeometryVeto = True
ConfigFlags.Detector.SimulateTrigger= True
ConfigFlags.Detector.GeometryTrigger= True
ConfigFlags.Detector.SimulatePreshower = True
ConfigFlags.Detector.GeometryPreshower = True
ConfigFlags.Detector.SimulateFaserSCT = True
ConfigFlags.Detector.GeometryFaserSCT = True
ConfigFlags.Detector.SimulateUpstreamDipole = True
ConfigFlags.Detector.SimulateCentralDipole = True
ConfigFlags.Detector.SimulateDownstreamDipole = True
ConfigFlags.Detector.GeometryUpstreamDipole = True
ConfigFlags.Detector.GeometryCentralDipole = True
ConfigFlags.Detector.GeometryDownstreamDipole = True
ConfigFlags.GeoModel.Align.Dynamic = False
ConfigFlags.Sim.ReleaseGeoModel = False
#
# All flags should be set before calling lock
#
ConfigFlags.lock()
#
# Construct ComponentAccumulator
#
acc = MainServicesCfg(ConfigFlags)
#
# Particle Gun generator (comment out to read generator file)
# Raw energies (without units given) are interpreted as MeV
#
pg = PG.ParticleGun()
pg.McEventKey = "GEN_EVENT"
pg.randomSeed = 123456
pg.sampler = DIFSampler()
acc.addEventAlgo(pg, "AthBeginSeq") # to run *before* G4
#
# Only one of these two should be used in a given job
# (MCEventSelectorCfg for generating events with no input file,
# PoolReadCfg when reading generator data from an input file)
#
acc.merge(McEventSelectorCfg(ConfigFlags))
# acc.merge(PoolReadCfg(ConfigFlags))
#
# Output stream configuration
#
acc.merge(OutputStreamCfg(ConfigFlags,
"HITS",
["EventInfo#*",
"McEventCollection#TruthEvent",
"McEventCollection#GEN_EVENT",
"ScintHitCollection#*",
"FaserSiHitCollection#*"
], disableEventTag=True))
acc.getEventAlgo("OutputStreamHITS").AcceptAlgs = ["G4FaserAlg"] # optional
acc.getEventAlgo("OutputStreamHITS").WritingTool.ProcessingTag = "StreamHITS" # required
#
# Here is the configuration of the Geant4 pieces
#
acc.merge(FaserGeometryCfg(ConfigFlags))
acc.merge(G4FaserAlgCfg(ConfigFlags))
acc.addService(G4GeometryNotifierSvcCfg(ConfigFlags, ActivateLVNotifier=True))
#
# Verbosity
#
# ConfigFlags.dump()
# logging.getLogger('forcomps').setLevel(VERBOSE)
# acc.foreach_component("*").OutputLevel = VERBOSE
# acc.foreach_component("*ClassID*").OutputLevel = INFO
# acc.getService("StoreGateSvc").Dump=True
# acc.getService("ConditionStore").Dump=True
# acc.printConfig()
f=open('FaserG4AppCfg_EVNT.pkl','wb')
acc.store(f)
f.close()
#
# Execute and finish
#
sys.exit(int(acc.run(maxEvents=500).isFailure()))
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment