Commit f4c9eb40 authored by Andy Buckley's avatar Andy Buckley Committed by Graeme Stewart
Browse files

'Re-fix mass setting ;-) Previous version is mostly right but could produce...

'Re-fix mass setting ;-) Previous version is mostly right but could produce out-of-sync vector and gen masses.' (ParticleGun-02-00-06)

	* Tagging: ParticleGun-02-00-06

	* Re-fix mass setting ;-) Previous version is mostly right but could produce out-of-sync vector and gen masses.

2015-01-29  Ewelina Lobodzinska  <ewelina@mail.desy.de>

        * Tagging: ParticleGun-02-00-05

        * implement fix provided by Zach for mass setting

2015-01-28  Andy Buckley  <andy.buckley@cern.ch>

	* Tagging: ParticleGun-02-00-04

	* Improve generated_mass setting and remove accidental print.

2015-01-25  Andy Buckley  <andy.buckley@cern.ch>

	* Tagging: ParticleGun-02-00-03
...
(Long ChangeLog diff - truncated)
parent 4c619ee8
......@@ -7,7 +7,8 @@ theApp.EvtMax = 100
import ParticleGun as PG
pg = PG.ParticleGun()
pg.sampler.pid = 11
pg.randomSeed = 123456
pg.sampler.pid = {11,-11,211,111}
pg.sampler.mom = PG.EEtaMPhiSampler(energy=10000, eta=[-2,2])
topSeq += pg
......
......@@ -26,7 +26,7 @@ class PtEtaHistParticleSampler(PG.ParticleSampler):
ptrand *= 1000 # NB. This _particular_ histogram is in GeV, but Athena needs MeV!
# TODO: Provide 4-mom construction functions to avoid building this one-time sampler
pid = self.pid()
mom = PG.PtEtaMPhiSampler(pt=ptrand, eta=etarand, mass=PG.MASSES[pid])
mom = PG.PtEtaMPhiSampler(pt=ptrand, eta=etarand, mass=PG.MASSES[abs(pid)])
p = PG.SampledParticle(pid, mom())
#print p.mom.Pt(), "\t", p.mom.Eta(), "\t", p.mom.Phi(), "\t", p.mom.M()
particles.append(p)
......
......@@ -15,12 +15,12 @@ class ParticleGun(EvgenAlg):
A simple but flexible algorithm for generating events from simple distributions.
"""
def __init__(self, name="ParticleGun", randomSvcName="AtRndmGenSvc", randomStream="ParticleGun"):
def __init__(self, name="ParticleGun", randomSvcName="AtRndmGenSvc", randomStream="ParticleGun", randomSeed=None):
super(ParticleGun, self).__init__(name=name)
self.samplers = [ParticleSampler()]
self.randomStream = randomStream
self.randomSvcName = randomSvcName
self.randomSeed = randomSeed
@property
def sampler(self):
......@@ -35,20 +35,29 @@ class ParticleGun(EvgenAlg):
def initialize(self):
"""
Pass the AtRndmGenSvc seed to Python's random module (or fall back to a fixed value for reproducibility).
TODO: Also allow directly specifying the seed as a constructor/Athena property -> self.randomSeed=None?
"""
seed = "123456"
randomSvc = getattr(svcMgr, self.randomSvcName, None)
if randomSvc is not None:
for seedstr in randomSvc.Seeds:
if seedstr.startswith(self.randomStream):
seed = seedstr
self.msg.info("ParticleGun: Using random seed '%s'" % seed)
seed = None
## Use self.randomSeed directly, or if it's None find a seed string from the ATLAS random number service
if self.randomSeed is not None:
seed = self.randomSeed
else:
self.msg.warning("ParticleGun: Failed to find random number service called '%s'" % self.randomSvcName)
random.seed(seed)
return StatusCode.Success
randomSvc = getattr(svcMgr, self.randomSvcName, None)
if randomSvc is not None:
for seedstr in randomSvc.Seeds:
if seedstr.startswith(self.randomStream):
seed = seedstr
self.msg.info("ParticleGun: Using random seed '%s'" % seed)
break
if seed is None:
self.msg.warning("ParticleGun: Failed to find a seed for the random stream named '%s'" % self.randomStream)
else:
self.msg.warning("ParticleGun: Failed to find random number service called '%s'" % self.randomSvcName)
## Apply the seed
if seed is not None:
random.seed(seed)
return StatusCode.Success
else:
return StatusCode.Failure
def fillEvent(self, evt):
......@@ -64,10 +73,10 @@ class ParticleGun(EvgenAlg):
particles = s.shoot()
for p in particles:
## Debug printout of particle properties
#print p.pid, p.mom.E(), p.mom.M()
#print "(px,py,pz,E) = (%0.2e, %0.2e, %0.2e, %0.2e)" % (p.mom.Px(), p.mom.Py(), p.mom.Pz(), p.mom.E())
#print "(eta,phi,pt,m) = (%0.2e, %0.2e, %0.2e, %0.2e)" % (p.mom.Eta(), p.mom.Phi(), p.mom.Pt(), p.mom.M())
#print "(x,y,z,t) = (%0.2e, %0.2e, %0.2e, %0.2e)" % (p.pos.X(), p.pos.Y(), p.pos.Z(), p.pos.T())
#print DEBUG0, p.pid, p.mom.E(), p.mom.Pt(), p.mom.M()
#print "DEBUG1 (px,py,pz,E) = (%0.2e, %0.2e, %0.2e, %0.2e)" % (p.mom.Px(), p.mom.Py(), p.mom.Pz(), p.mom.E())
#print "DEBUG2 (eta,phi,pt,m) = (%0.2e, %0.2e, %0.2e, %0.2e)" % (p.mom.Eta(), p.mom.Phi(), p.mom.Pt(), p.mom.M())
#print "DEBUG3 (x,y,z,t) = (%0.2e, %0.2e, %0.2e, %0.2e)" % (p.pos.X(), p.pos.Y(), p.pos.Z(), p.pos.T())
## Make particle-creation vertex
# TODO: do something cleverer than one vertex per particle?
......@@ -79,9 +88,11 @@ class ParticleGun(EvgenAlg):
## Make particle with status == 1
mom = HepMC.FourVector(p.mom.Px(), p.mom.Py(), p.mom.Pz(), p.mom.E())
gp = HepMC.GenParticle()
gp.set_status(1)
gp.set_pdg_id(p.pid)
gp.set_momentum(mom)
gp.set_status(1)
if p.mass is not None:
gp.set_generated_mass(p.mass)
ROOT.SetOwnership(gp, False)
gv.add_particle_out(gp)
......
......@@ -41,12 +41,19 @@ def get_sampling_vars(h):
globalbin_to_axisbin = {} # for reverse axis bin lookup to get edges
globalbins = [] # because they aren't easily predicted, nor contiguous
cheights = [0] # cumulative "histogram" from which to uniformly sample
for ix in xrange(1, h.GetNbinsX()+1):
for iy in xrange(1, h.GetNbinsY()+1):
iglobal = h.GetBin(ix, iy)
if issubclass(type(h), ROOT.TH1):
for ix in xrange(1, h.GetNbinsX()+1):
iglobal = h.GetBin(ix)
globalbins.append(iglobal)
globalbin_to_axisbin[iglobal] = (ix, iy) if issubclass(type(h), ROOT.TH2) else tuple(ix,)
globalbin_to_axisbin[iglobal] = (ix,)
cheights.append(cheights[-1] + h.GetBinContent(iglobal))
elif issubclass(type(h), ROOT.TH2):
for ix in xrange(1, h.GetNbinsX()+1):
for iy in xrange(1, h.GetNbinsY()+1):
iglobal = h.GetBin(ix, iy)
globalbins.append(iglobal)
globalbin_to_axisbin[iglobal] = (ix, iy)
cheights.append(cheights[-1] + h.GetBinContent(iglobal))
return globalbins, globalbin_to_axisbin, cheights
......@@ -74,7 +81,7 @@ def get_random_x(h, globalbins, cheights, globalbin_to_axisbin):
irand = get_random_bin(globalbins, cheights)
axisids = globalbin_to_axisbin.get(irand)
assert axisids is not None
xrand = random.uniform(h2.GetXaxis().GetBinLowEdge(axisids[0]), h2.GetXaxis().GetBinUpEdge(axisids[0]))
xrand = random.uniform(h.GetXaxis().GetBinLowEdge(axisids[0]), h.GetXaxis().GetBinUpEdge(axisids[0]))
return xrand
......@@ -92,10 +99,17 @@ def get_random_xy(h2, globalbins, cheights, globalbin_to_axisbin):
class TH1(object):
"Minimal wrapper for ROOT TH1, for consistency and easy loading"
"Minimal wrapper for ROOT TH1, for sampling consistency and easy loading"
def __init__(self, *args):
self.th1 = load_hist(*args)
self.globalbins, self.globalbin_to_axisbin, self.cheights = None, None, None
def GetRandom(self):
"A GetRandom that works for TH1s and uses Python random numbers"
if self.globalbins is None or self.globalbin_to_axisbin is None or self.cheights is None:
self.globalbins, self.globalbin_to_axisbin, self.cheights = get_sampling_vars(self.th1)
return get_random_x(self.th1, self.globalbins, self.cheights, self.globalbin_to_axisbin)
def __getattr__(self, attr):
"Forward all attributes to the contained TH1"
......@@ -107,10 +121,12 @@ class TH2(object):
def __init__(self, *args):
self.th2 = load_hist(*args)
self.globalbins, self.globalbin_to_axisbin, self.cheights = get_sampling_vars(self.th2)
self.globalbins, self.globalbin_to_axisbin, self.cheights = None, None, None
def GetRandom(self):
"A GetRandom that works for TH2s"
if self.globalbins is None or self.globalbin_to_axisbin is None or self.cheights is None:
self.globalbins, self.globalbin_to_axisbin, self.cheights = get_sampling_vars(self.th2)
return get_random_xy(self.th2, self.globalbins, self.cheights, self.globalbin_to_axisbin)
def __getattr__(self, attr):
......
......@@ -204,7 +204,7 @@ class RandomSeqSampler(DiscreteSampler):
self.sequence = args
def shoot(self):
return self.sequence.random()
return random.choice(self.sequence)
# Alias:
RndmSeq = RandomSeqSampler
......@@ -841,6 +841,7 @@ class SampledParticle(object):
self.pid = pid
self.mom = mom
self.pos = pos
self.mass = None
class ParticleSampler(Sampler):
......@@ -877,16 +878,20 @@ class ParticleSampler(Sampler):
def shoot(self):
"Return a vector of sampled particles"
numparticles = self.n()
rtn = []
for i in xrange(numparticles):
## Sample the particle ID and pass mass info to the v4 sampler
## Sample the particle ID and create a particle
pid = self.pid()
p = SampledParticle(pid)
## Pass mass info to the v4 sampler and set same generated mass
if self.mass_override and self.massdict.has_key(abs(pid)):
self.mom.mass = self.massdict[abs(pid)]
## Sample momentum and vertex positions
mom = self.mom()
pos = self.pos()
rtn.append( SampledParticle(pid, mom, pos) )
m = self.massdict[abs(pid)]
self.mom.mass = m
p.mass = m
## Sample momentum and vertex positions into the particle
p.mom = self.mom()
p.pos = self.pos()
## Add particle to output list
rtn.append(p)
return rtn
Supports Markdown
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