'Adding a 2-particle example JO for the new multiple samplers feature.' (ParticleGun-02-00-01)

Generators/ParticleGun/README

+ParticleGun documentation
+-------------------------
+See https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/ParticleGunForAtlas
+for some coherent documentation that should be kept up to date.
+

Generators/ParticleGun/TODO

+ParticleGun TODO list
+---------------------
+
+NOTHING!

Generators/ParticleGun/cmt/requirements

+package ParticleGun
+
+author Andy Buckley <andy.buckley@cern.ch>
+
+use AtlasPolicy           AtlasPolicy-*
+
+## Note: although these packages are required, the dependencies are Python only
+#use GeneratorModules      GeneratorModules-*      Generators
+#use AthenaKernel          AthenaKernel-*          Control
+#use GaudiInterface        GaudiInterface-*        External
+#use AthenaBaseComps       AthenaBaseComps-*       Control
+#use AtlasCLHEP            AtlasCLHEP-*            External
+#use StoreGate             StoreGate-*             Control
+#use McParticleEvent       McParticleEvent-*       PhysicsAnalysis/TruthParticleID
+#use NavFourMom            NavFourMom-*            Event
+
+apply_pattern declare_joboptions files="*.py"
+apply_pattern declare_python_modules files="*.py"

Generators/ParticleGun/examples/jobOption.ParticleGun_constenergy_flateta.py

+#! -*- python -*-
+
+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+include("GeneratorUtils/StdEvgenSetup.py")
+theApp.EvtMax = 100
+
+import ParticleGun as PG
+pg = PG.ParticleGun()
+pg.sampler.pid = 11
+pg.sampler.mom = PG.EEtaMPhiSampler(energy=10000, eta=[-2,2])
+topSeq += pg
+
+include("GeneratorUtils/postJO.CopyWeights.py")
+include("GeneratorUtils/postJO.PoolOutput.py")
+include("GeneratorUtils/postJO.DumpMC.py")

Generators/ParticleGun/examples/jobOption.ParticleGun_correlated.py

+#! -*- python -*-
+
+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+include("GeneratorUtils/StdEvgenSetup.py")
+theApp.EvtMax = 100
+
+import ParticleGun as PG
+
+class MyParticleSampler(PG.ParticleSampler):
+    "A special sampler with two _correlated_ particles."
+
+    def __init__(self):
+        self.mom1 = PG.PtEtaMPhiSampler(pt=25000, eta=[-2,2])
+
+    def shoot(self):
+        "Return a vector of sampled particles"
+        p1 = PG.SampledParticle(11, self.mom1.shoot())
+        eta1 = p1.mom.Eta()
+        phi1 = p1.mom.Phi()
+        # TODO: will phi be properly wrapped into range?
+        mom2 = PG.PtEtaMPhiSampler(pt=25000,
+                                   eta=[eta1-0.5, eta1+0.5],
+                                   phi=[phi1-0.5, phi1+0.5])
+        p2 = PG.SampledParticle(11, mom2.shoot())
+        return [p1, p2]
+
+topSeq += PG.ParticleGun()
+topSeq.ParticleGun.sampler = MyParticleSampler()
+
+include("GeneratorUtils/postJO.CopyWeights.py")
+include("GeneratorUtils/postJO.PoolOutput.py")
+include("GeneratorUtils/postJO.DumpMC.py")

Generators/ParticleGun/examples/jobOption.ParticleGun_corrhist.py

+#! -*- python -*-
+
+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+## ROOT 2D histogram sampling alg (in ParticleGun.histsampling) by Andy Buckley
+## Thanks to Alejandro Alonso for the initial Athena example on which this is based.
+
+include("GeneratorUtils/StdEvgenSetup.py")
+theApp.EvtMax = 100
+
+import ParticleGun as PG
+
+class PtEtaHistParticleSampler(PG.ParticleSampler):
+    "Particle sampler with correlated pT and eta from a 2D histogram."
+
+    def __init__(self, pid, histfile, num=100):
+        self.pid = PG.mksampler(pid)
+        self.hist = PG.TH2(histfile, "h_pt_eta")
+        self.numparticles = num
+
+    def shoot(self):
+        "Return a vector of sampled particles from the provided pT--eta histogram"
+        particles = []
+        for i in xrange(self.numparticles):
+            ptrand, etarand = self.hist.GetRandom()
+            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])
+            p = PG.SampledParticle(pid, mom())
+            #print p.mom.Pt(), "\t", p.mom.Eta(), "\t", p.mom.Phi(), "\t", p.mom.M()
+            particles.append(p)
+        return particles
+
+topSeq += PG.ParticleGun()
+topSeq.ParticleGun.sampler = PtEtaHistParticleSampler(11, "data_histos_el_1470pt.root")
+
+include("GeneratorUtils/postJO.CopyWeights.py")
+include("GeneratorUtils/postJO.PoolOutput.py")
+include("GeneratorUtils/postJO.DumpMC.py")

Generators/ParticleGun/examples/jobOption.ParticleGun_flatpt_2particle.py

+#! -*- python -*-
+
+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+include("GeneratorUtils/StdEvgenSetup.py")
+theApp.EvtMax = 100
+
+import ParticleGun as PG
+pg = PG.ParticleGun()
+pg.samplers.append(PG.ParticleSampler()) # add a second sampler
+pg.samplers[0].pid = (-13, 13) # cycle mu+-
+pg.samplers[0].mom = PG.PtEtaMPhiSampler(pt=[4000, 100000], eta=[1.0, 3.2]) # flat in pt and +ve eta
+pg.samplers[1].pid = (13, -13) # cycle mu-+
+pg.samplers[1].mom = PG.PtEtaMPhiSampler(pt=[4000, 100000], eta=[-3.2, -1.0]) # flat in pt and -ve eta
+topSeq += pg
+
+include("GeneratorUtils/postJO.CopyWeights.py")
+include("GeneratorUtils/postJO.PoolOutput.py")
+include("GeneratorUtils/postJO.DumpMC.py")

Generators/ParticleGun/examples/jobOption.ParticleGun_fwd_sequence.py

+#! -*- python -*-
+
+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+include("GeneratorUtils/StdEvgenSetup.py")
+theApp.EvtMax = 100
+
+import ParticleGun as PG
+pg = PG.ParticleGun()
+pg.sampler.pid = (2112, 22, 2112, 22)
+pg.sampler.mom = PG.EThetaMPhiSampler(energy=(1360000, 500000, 1360000, 500000),
+                                      theta=(0, 0, PG.PI, PG.PI))
+pg.sampler.pos = PG.PosSampler(x=[-120,-100], y=[-10,10], z=203950)
+topSeq += pg
+
+include("GeneratorUtils/postJO.CopyWeights.py")
+include("GeneratorUtils/postJO.PoolOutput.py")
+include("GeneratorUtils/postJO.DumpMC.py")

Generators/ParticleGun/examples/jobOption.ParticleGun_vtx.py

+#! -*- python -*-
+
+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+include("GeneratorUtils/StdEvgenSetup.py")
+theApp.EvtMax = 100
+
+import ParticleGun as PG
+pg = PG.ParticleGun()
+pg.sampler.pid = 13
+pg.sampler.pos = PG.PosSampler(x=3140.0, y=[-154.134,154.134], z=[4938.76,5121.29], t=5929.7)
+pg.sampler.mom = PG.EEtaMPhiSampler(energy=100000, eta=1.25, phi=0.0)
+topSeq += pg
+
+include("GeneratorUtils/postJO.CopyWeights.py")
+include("GeneratorUtils/postJO.PoolOutput.py")
+include("GeneratorUtils/postJO.DumpMC.py")

Generators/ParticleGun/python/__init__.py

+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+from AthenaCommon.AppMgr import ServiceMgr as svcMgr
+from GeneratorModules.EvgenAlg import EvgenAlg
+from ParticleGun.samplers import *
+from ParticleGun.histsampling import TH1, TH2
+from AthenaPython.PyAthena import HepMC, StatusCode
+import McParticleEvent.Pythonizations
+
+__author__ = "Andy Buckley <andy.buckley@cern.ch>"
+
+
+class ParticleGun(EvgenAlg):
+    """
+    A simple but flexible algorithm for generating events from simple distributions.
+    """
+
+    def __init__(self, name="ParticleGun", randomSvcName="AtRndmGenSvc", randomStream="ParticleGun"):
+        super(ParticleGun, self).__init__(name=name)
+        self.samplers = [ParticleSampler()]
+        self.randomStream = randomStream
+        self.randomSvcName = randomSvcName
+
+
+    @property
+    def sampler(self):
+        "Get the first (and presumed only) sampler"
+        return self.samplers[0] if self.samplers else None
+    @sampler.setter
+    def sampler(self, s):
+        "Set the samplers list to include only a single sampler, s"
+        self.samplers = [s]
+
+
+    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)
+        else:
+            self.msg.warning("ParticleGun: Failed to find random number service called '%s'" % self.randomSvcName)
+        random.seed(seed)
+        return StatusCode.Success
+
+
+    def fillEvent(self, evt):
+        """
+        Sample a list of particle properties, which are then used to create a new GenEvent in StoreGate.
+        """
+        ## Set event weight(s)
+        # TODO: allow weighted sampling?
+        evt.weights().push_back(1.0)
+
+        ## Make and fill particles
+        for s in self.samplers:
+            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())
+
+                ## Make particle-creation vertex
+                # TODO: do something cleverer than one vertex per particle?
+                pos = HepMC.FourVector(p.pos.X(), p.pos.Y(), p.pos.Z(), p.pos.T())
+                gv = HepMC.GenVertex(pos)
+                ROOT.SetOwnership(gv, False)
+                evt.add_vertex(gv)
+
+                ## 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_pdg_id(p.pid)
+                gp.set_momentum(mom)
+                gp.set_status(1)
+                ROOT.SetOwnership(gp, False)
+                gv.add_particle_out(gp)
+
+        return StatusCode.Success
+
+
+## PyAthena HepMC notes
+#
+## evt.print() isn't valid syntax in Python2 due to reserved word
+# TODO: Add a Pythonisation, e.g. evt.py_print()?
+#getattr(evt, 'print')()
+#
+## How to check that the StoreGate key exists and is an McEventCollection
+# if self.sg.contains(McEventCollection, self.sgkey):
+#     print self.sgkey + " found!"
+#
+## Modifying an event other than that supplied as an arg
+# mcevts = self.sg[self.sgkey]
+# for vtx in mcevts[0].vertices: # only way to get the first vtx?!
+#     gp2 = HepMC.GenParticle()
+#     gp2.set_momentum(HepMC.FourVector(1,2,3,4))
+#     gp2.set_status(1)
+#     vtx.add_particle_out(gp2)
+#     break

Generators/ParticleGun/python/histsampling.py

+# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+
+"""
+Tools for histogram sampling, in particular inverse transform sampling which is
+missing from ROOT's TH2 classes.
+"""
+
+__author__ = "Andy Buckley <andy.buckley@cern.ch>"
+
+import random, ROOT
+
+
+def load_hist(*args):
+    """
+    Load a histogram from a filename/TFile and histo name. If a single arg is
+    provided, it has to be a histo object and will be cloned before return.
+    """
+    h = None
+    if len(args) == 1 and issubclass(type(args[0]), ROOT.TH1):
+        h = args[0].Clone()
+    elif len(args) == 2:
+        if type(args[0]) is type(args[1]) is str:
+            f = ROOT.TFile.Open(args[0])
+            htmp = f.Get(args[1])
+            h = f.Get(args[1]).Clone()
+            #f.Close()
+        elif type(args[0]) is ROOT.TFile and type(args[1]) is str:
+            h = args[0].Get(args[1]).Clone()
+    if h is None:
+        raise Exception("Error in histogram loading from " + args)
+    return h
+
+
+def get_sampling_vars(h):
+    """
+    Get the following from a histogram h, since the ROOT API sucks:
+    * list of global bin IDs (not even contiguous for 2D, gee thanks ROOT)
+    * dict mapping global bin IDs to a tuple of axis bin IDs
+    * list of nbins+1 cumulative bin values, in the same order as globalbins
+    """
+    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)
+            globalbins.append(iglobal)
+            globalbin_to_axisbin[iglobal] = (ix, iy) if issubclass(type(h), ROOT.TH2) else tuple(ix,)
+            cheights.append(cheights[-1] + h.GetBinContent(iglobal))
+    return globalbins, globalbin_to_axisbin, cheights
+
+
+def get_random_bin(globalbins, cheights):
+    """
+    Choose a random bin from the cumulative distribution list of nbins+1 entries.
+
+    TODO: Search more efficiently (lin and log guesses, then lin search or
+    binary split depending on vector size).
+    """
+    assert len(cheights) == len(globalbins)+1
+    randomheight = random.uniform(0, cheights[-1])
+    randombin = None
+    for i, iglobal in enumerate(globalbins):
+        if randomheight >= cheights[i] and randomheight < cheights[i+1]:
+            return iglobal
+    raise Exception("Sample fell outside range of cumulative distribution?!?!")
+
+
+def get_random_x(h, globalbins, cheights, globalbin_to_axisbin):
+    """
+    Choose a random bin via get_random_bin, then pick a uniform random x
+    point in that bin (without any attempt at estimating the in-bin distribution).
+    """
+    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]))
+    return xrand
+
+
+def get_random_xy(h2, globalbins, cheights, globalbin_to_axisbin):
+    """
+    Choose a random bin via get_random_bin, then pick a uniform random x,y
+    point in that bin (without any attempt at estimating the in-bin distribution).
+    """
+    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]))
+    yrand = random.uniform(h2.GetYaxis().GetBinLowEdge(axisids[1]), h2.GetYaxis().GetBinUpEdge(axisids[1]))
+    return xrand, yrand
+
+
+class TH1(object):
+    "Minimal wrapper for ROOT TH1, for consistency and easy loading"
+
+    def __init__(self, *args):
+        self.th1 = load_hist(*args)
+
+    def __getattr__(self, attr):
+        "Forward all attributes to the contained TH1"
+        return getattr(self.th1, attr)
+
+
+class TH2(object):
+    "Minimal wrapper for ROOT TH2, for easy loading and to allow 2D sampling"
+
+    def __init__(self, *args):
+        self.th2 = load_hist(*args)
+        self.globalbins, self.globalbin_to_axisbin, self.cheights = get_sampling_vars(self.th2)
+
+    def GetRandom(self):
+        "A GetRandom that works for TH2s"
+        return get_random_xy(self.th2, self.globalbins, self.cheights, self.globalbin_to_axisbin)
+
+    def __getattr__(self, attr):
+        "Forward other attributes to the contained TH2"
+        return getattr(self.th2, attr)