'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)

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
 

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)

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)
 

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):

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