From f4c9eb40df3d45590ab36fafb07735147d461ea6 Mon Sep 17 00:00:00 2001
From: Andy Buckley <a.g.buckley@gmail.com>
Date: Thu, 29 Jan 2015 19:47:32 +0100
Subject: [PATCH] '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)
---
...bOption.ParticleGun_constenergy_flateta.py | 3 +-
.../jobOption.ParticleGun_corrhist.py | 2 +-
Generators/ParticleGun/python/__init__.py | 49 ++++++++++++-------
Generators/ParticleGun/python/histsampling.py | 30 +++++++++---
Generators/ParticleGun/python/samplers.py | 21 +++++---
5 files changed, 69 insertions(+), 36 deletions(-)
diff --git a/Generators/ParticleGun/examples/jobOption.ParticleGun_constenergy_flateta.py b/Generators/ParticleGun/examples/jobOption.ParticleGun_constenergy_flateta.py
index 7b1d14e11f1..3c3cb0321b3 100644
--- a/Generators/ParticleGun/examples/jobOption.ParticleGun_constenergy_flateta.py
+++ b/Generators/ParticleGun/examples/jobOption.ParticleGun_constenergy_flateta.py
@@ -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
diff --git a/Generators/ParticleGun/examples/jobOption.ParticleGun_corrhist.py b/Generators/ParticleGun/examples/jobOption.ParticleGun_corrhist.py
index 33ad1d9a118..f5539702d5d 100644
--- a/Generators/ParticleGun/examples/jobOption.ParticleGun_corrhist.py
+++ b/Generators/ParticleGun/examples/jobOption.ParticleGun_corrhist.py
@@ -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)
diff --git a/Generators/ParticleGun/python/__init__.py b/Generators/ParticleGun/python/__init__.py
index a5bc6a1564c..e6aaa7feac2 100644
--- a/Generators/ParticleGun/python/__init__.py
+++ b/Generators/ParticleGun/python/__init__.py
@@ -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)
diff --git a/Generators/ParticleGun/python/histsampling.py b/Generators/ParticleGun/python/histsampling.py
index e54e6ef0052..6d97ca2a929 100644
--- a/Generators/ParticleGun/python/histsampling.py
+++ b/Generators/ParticleGun/python/histsampling.py
@@ -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):
diff --git a/Generators/ParticleGun/python/samplers.py b/Generators/ParticleGun/python/samplers.py
index 94a38fa577c..d29078e82d1 100644
--- a/Generators/ParticleGun/python/samplers.py
+++ b/Generators/ParticleGun/python/samplers.py
@@ -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
--
GitLab