Python function replacing part of variable

Question

I am writing a code for a project in particle physics (using pyroot).

In my first draft, I use the following line

for i in MyTree:    

   pion.SetXYZM(K_plus_PX, K_plus_PY, K_plus_PZ,K_plus_MM)

This basically assigns to the pion the values of variables in the parenthesis, ie momenta and inv. mass of the kaon.

Physics aside, I would like to write a function "of the form":

def myfunc(particle):
    return %s_PX % particle

I know this is wrong. What I would like to achieve is to write a function that allows, for a given particle, to set particle_PX, particle_PY etc to be the arguments of SetXYZM.

Thank you for your help,

B

Answer 1

To access class attributes from string variables you can use python's getattr:

import ROOT
inputfile = ROOT.TFile.Open("somefile.root","read")
inputtree = inputfile.Get("NameOfTTree")
inputtree.Print()
# observe that there are branches
# K_plus_PX
# K_plus_PY
# K_plus_PZ
# K_plus_MM
# K_minus_PX
# K_minus_PY
# K_minus_PZ
# K_minus_MM
# pi_minus_PX
# pi_minus_PY
# pi_minus_PZ
# pi_minus_MM

def getx(ttree,particlename):
    return getattr(ttree,particlename+"_PX")
def gety(ttree,particlename):
    return getattr(ttree,particlename+"_PY")
def getz(ttree,particlename):
    return getattr(ttree,particlename+"_PZ")
def getm(ttree,particlename):
    return getattr(ttree,particlename+"_MM")
def getallfour(ttree,particlename):
    x = getattr(ttree,particlename+"_PX")
    y = getattr(ttree,particlename+"_PY")
    z = getattr(ttree,particlename+"_PZ")
    m = getattr(ttree,particlename+"_MM")
    return x,y,z,m


for entry in xrange(inputtree.GetEntries()):
    inputtree.GetEntry(entry)
    pion1 = ROOT.TLorentzVector()
    x = getx(inputtree,"K_plus")
    y = gety(inputtree,"K_plus")
    z = getz(inputtree,"K_plus")
    m = getm(inputtree,"K_plus")
    pion2.SetXYZM(x,y,z,m)
    x,y,z,m = getallfour(inputtree,"pi_minus")
    pion2 = ROOT.TLorentzVector()
    pion2.SetXYZM(x,y,z,m)

As linked by Josh Caswell, you can similarly access variable names:

def getx(particlename):
    x = globals()[partilcename+"_PX"]

though that might get nasty quickly as of whether your variables are global or local and for local, in which context.