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

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

parent a964e288
ParticleGun documentation
-------------------------
See https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/ParticleGunForAtlas
for some coherent documentation that should be kept up to date.
ParticleGun TODO list
---------------------
NOTHING!
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"
#! -*- 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")
#! -*- 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")
#! -*- 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")
#! -*- 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")
#! -*- 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")
#! -*- 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")
# 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
# 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)
This diff is collapsed.
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