Hello,
I am trying to optimise the speed to a PyRoot script that combines the invariant masses of muon tracks in a TTree and then combines the resulting invariant mass with pairwise combinations of oppositely charged pion tracks. I have written a function that performs the task at hand, however, it is crucial for me to try to decrease the runtime.
Here is my current attempt:
from ROOT import *
from math import *
import heapq
import numpy as np
import itertools
from array import array
def MM(tree, hist)::
for entry in tree:
#access the momenta and other vars
buf_run = np.array(entry.run)
buf_evt = np.array(entry.evt)
buf_bunchID = np.array(entry.bunchid)
bufpx = np.array(entry.px)
bufpy = np.array(entry.py)
bufpz = np.array(entry.pz)
bufpt = np.array(entry.pt)
bufPIDp = np.array(entry.pidp)
bufPIDk = np.array(entry.pidk)
bufPIDpi = np.array(entry.pidpi)
isMuon = np.array(entry.muon)
longTrack = np.array(entry.nlong)
Charge = np.array(entry.charge)
if sum(isMuon)>=1 and sum(Charge[np.where(isMuon==1)])==0: #at least one muon track and sum of charges of muon tracks = 0
#create an empty 4-vector to track muon mass
nul_mu = TLorentzVector()
nul_mu.SetXYZM(0.,0.,0.,0.)
#for every track in the event
for i in range(len(bufpx)):
#add muon 4-vecs
if isMuon[i]==1:
mubuf_particle = TLorentzVector()
mubuf_particle.SetXYZM(bufpx[i], bufpy[i], bufpz[i], m_mu)
nul_mu+=mubuf_particle
#reconstruct mass around the J-Psi nominal mass value and at least 2 other non-muon (pion) tracks
if nul_mu.M() > (m_JPsi - 65.) and nul_mu.M() < (m_JPsi + 65.) and len(np.where(isMuon!=1)[0])>=2:
pi_idxs=np.where(isMuon!=1) #indices of pio tracks
temp=np.array(list(itertools.combinations(pi_idxs[0], 2))) #combinatorics of indices of hadron tracks, no repetitions, pairwise combinations
neutral_pairs_idxs = temp[np.where(np.sum(Charge[temp], axis=1)==0)]
for pair in neutral_pairs_idxs:
idx1 = pair[0]
idx2 = pair[1]
buf_piplus = TLorentzVector()
buf_piminus = TLorentzVector()
buf_piplus.SetXYZM( bufpx[idx1], bufpy[idx1], bufpz[idx1] , m_pi )
buf_piminus.SetXYZM( bufpx[idx2], bufpy[idx2], bufpz[idx2] , m_pi )
#Fill with JPsi and 2 hadron (opposite charge) tracks
#hist.Fill( (nul_mu+buf_piplus+buf_piminus).M())
#Fill mu+mu+pi+pi_MM - mu+mu_MM + nominal mass value of J/Psi (=3097), should improve S/B separation
#mumu_array[0]=(nul_mu+buf_piplus+buf_piminus).M() - nul_mu.M() + 3097
#tdata.Fill()
hist.Fill( (nul_mu+buf_piplus+buf_piminus).M() - nul_mu.M() + 3097)
In my attempt to speed to the code I use only the branches I am interested in
#worry only about branches of interest, in the hope to speed the code
c.SetBranchStatus("*", 0)
c.SetBranchStatus("px",1)
c.SetBranchStatus("py",1)
c.SetBranchStatus("pz",1)
c.SetBranchStatus("muon",1)
c.SetBranchStatus("charge",1)
Can anyone suggest a way to speed up the code either in reading the entries or perform the pairwise combinations?
I have been thinking of writing a Root script to perform this task, and I do not seem to manage to perform the pairwise combinations successfully. Can anyone suggest a way to do this?
Thanks,
bldelane
