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

Could you clarify what `%s_PX % particle` is supposed to mean? How would you use this function? — sidney, Jul 03 '16 at 18:36
Possible duplicate of [Concatenating variable names in Python](http://stackoverflow.com/questions/36898891/concatenating-variable-names-in-python) — jscs, Jul 03 '16 at 18:39
You are setting i as each item in MyTree, but your code doesn't do anything with i. Can you add more explanation or an example of what you want your output to be? — joel goldstick, Jul 03 '16 at 18:40
@sidney and @joel goldstick, SetXYZM() is a function in ROOT that sets the values of the four momenta of a given particle, so that for particle.SetXYZM(0, 0, 0, 0), the particle will have zero spatial momentum and (for natural units) zero mass (silly example). My difficulty lies in the fact that `Kplus_PX` and variables of the same form are only numbers (possibly long floats) stored in a TTree ( basically the record of a hadron collision). I need to input manually the momenta and mass, and I was hoping there was a way to make this easier to generalise. — Blaise Delaney, Jul 03 '16 at 18:47
I am looking to avoid having a function that has a massive `if` statement where I spell out explictly all the possible combinations of arguments of `SetXYZM()` for a particle of choice. — Blaise Delaney, Jul 03 '16 at 18:48
@BlaiseDelaney What do you mean by combinations of arguments of `SetXYZM()` for a particle of choice? I'm still unsure of what the function is doing. — sidney, Jul 03 '16 at 18:55
@sidney So each particle has associated a 4-momentum, an array of four elements. So, by pion.SetXYZM() I simply want to specify the 4 numbers is want to be in said array. Now, I want to be able to input as elements of this 4-element-array 4 variables associated with another candidate particle - in my example a kaon (`Kplus`). Is there a way to write a function that in practice appends for `particle` the ends `_PX` and alike so to have `SetXYZM(particle_PX etc.)`? My apologies for being unclear. — Blaise Delaney, Jul 03 '16 at 18:57

score 1 · Accepted Answer · answered Jul 12 '16 at 09:53

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.

Python function replacing part of variable

1 Answers1