Python-функция, заменяющая часть переменной

Question

Я пишу код для проекта в области физики элементарных частиц (используя Pyroot).

В моем первом проекте я использую следующую строку

for i in MyTree:    

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

Это в основном присваивает пиону значения переменных в скобках, т.е. импульсы и инв. масса каона.

Помимо физики, я хотел бы написать функцию "вида":

def myfunc(particle):
    return %s_PX % particle

Я знаю, что это неправильно. Чего я хотел бы добиться, так это написать функцию, которая позволяет для данной частицы устанавливать в качестве аргументов SetXYZM частицы_PX, частицы_PY и т. Д.

Спасибо за помощь,

В

Answer 1

Для доступа к атрибутам класса из строковых переменных вы можете использовать Python 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)

По словам Джоша Касвелла, вы можете также получить доступ к именам переменных:

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

хотя это может быстро стать неприятным в зависимости от того, являются ли ваши переменные глобальными или локальными, а также локальными, в каком контексте.