127 lines
4.7 KiB
Python
127 lines
4.7 KiB
Python
import root_pandas as rp
|
|
import numpy as np
|
|
from uproot3_methods import TLorentzVectorArray as lorentzarr
|
|
from uproot3_methods import TLorentzVector as lorentz
|
|
from tqdm import tqdm
|
|
import sys
|
|
from shutil import copyfile
|
|
import os
|
|
from ROOT import TFile, TTree, TList
|
|
|
|
#fname = sys.argv[1]
|
|
#tname = sys.argv[2]
|
|
|
|
def makeNewKstar(fname,tname,MC):
|
|
d = rp.read_root(fname+'.root', key=tname)
|
|
|
|
# load DTF momenta in arrays of 4-vectors
|
|
fourmom = {}
|
|
for pi in ['mu_plus', 'mu_minus', 'pi_zero_resolved', 'K_plus']:
|
|
fourmom[pi] = lorentzarr.from_cartesian(d[pi+'_PX_DTF'], d[pi+'_PY_DTF'], d[pi+'_PZ_DTF'], d[pi+'_PE_DTF'])
|
|
|
|
# calculate the rescaling (k)
|
|
Bplus_M_PDG = 5279.34
|
|
kvalues = []
|
|
for i in tqdm(range(len(d)), desc='calculating pi0 rescaling'):
|
|
ks = np.linspace(0.2, 2, 1000)
|
|
ppi0 = fourmom['pi_zero_resolved'][i]
|
|
pi0_rescaled = lorentzarr.from_xyzm(ppi0.x*ks, ppi0.y*ks, ppi0.z*ks, fourmom['pi_zero_resolved'][i].mass)
|
|
Bmass_rescaled = (pi0_rescaled + fourmom['K_plus'][i] + fourmom['mu_plus'][i] + fourmom['mu_minus'][i]).mass
|
|
bestk = ks[np.argmin(abs(Bmass_rescaled-Bplus_M_PDG))]
|
|
kvalues.append(bestk)
|
|
|
|
# now apply the rescaling
|
|
pi0_p3 = fourmom['pi_zero_resolved'].p3 * kvalues
|
|
fourmom['pi0_rescaled'] = lorentzarr.from_xyzm(pi0_p3.x, pi0_p3.y, pi0_p3.z, fourmom['pi_zero_resolved'].mass)
|
|
|
|
d['pi_zero_resolved_PX_scaled'] = fourmom['pi0_rescaled'].x
|
|
d['pi_zero_resolved_PY_scaled'] = fourmom['pi0_rescaled'].y
|
|
d['pi_zero_resolved_PZ_scaled'] = fourmom['pi0_rescaled'].z
|
|
d['pi_zero_resolved_PT_scaled'] = fourmom['pi0_rescaled'].pt
|
|
d['pi_zero_resolved_PE_scaled'] = fourmom['pi0_rescaled'].E
|
|
d['K_star_plus_M_scaled'] = (fourmom['pi0_rescaled'] + fourmom['K_plus']).mass
|
|
|
|
if(MC):
|
|
truefourmom = {}
|
|
for pi in ['K_star_plus']:
|
|
truefourmom[pi] = lorentzarr.from_cartesian(d[pi+'_TRUEP_X'], d[pi+'_TRUEP_Y'], d[pi+'_TRUEP_Z'], d[pi+'_TRUEP_E'])
|
|
|
|
d['K_star_plus_M_TRUE'] = truefourmom['K_star_plus'].mass
|
|
|
|
d.to_root(fname+'_pi0rescaled.root', key=tname, store_index=False)
|
|
return
|
|
|
|
|
|
|
|
def cutFile(fname, tname):
|
|
d = rp.read_root(fname+'_pi0rescaled.root', key=tname,where='K_star_plus_M_scaled>791.66 & K_star_plus_M_scaled<991.66') #It takes & not &&!!!!!!!!
|
|
d.to_root(fname+'_pi0rescaledCut.root', key=tname, store_index=False)
|
|
copyfile(fname+'_pi0rescaledCut.root',fname+'.root')
|
|
os.remove(fname+'_pi0rescaledCut.root')
|
|
|
|
def mergeFiles(fname1, fname2, outname, tname):
|
|
tfile1 = TFile(fname1, 'read')
|
|
tfile2 = TFile(fname2, 'read')
|
|
ttree1 = tfile1.Get(tname)
|
|
ttree2 = tfile2.Get(tname)
|
|
treeList = TList()
|
|
treeList.append(ttree1)
|
|
treeList.append(ttree2)
|
|
outFile = TFile(outname,'recreate')
|
|
outFile.cd()
|
|
outTree = TTree.MergeTrees(treeList)
|
|
outFile.Write()
|
|
outFile.Close()
|
|
return
|
|
|
|
|
|
years = [2011,2012,2015,2016,2017,2018]
|
|
treeName = "DecayTree"
|
|
mainPath = "/home/lhcb/kopecna/B2KstarMuMu_clean/Data/Tuples"
|
|
#MC
|
|
for pol in ["down","up"]:
|
|
for yr in years:
|
|
if (yr==2015): continue
|
|
fileName = mainPath+"/MC/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
|
|
makeNewKstar(fileName,treeName,True)
|
|
cutFile(fileName,treeName)
|
|
|
|
#Ref
|
|
for pol in ["down","up"]:
|
|
for yr in years:
|
|
if (yr>2016): continue
|
|
fileName = mainPath+"/RefMC/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
|
|
makeNewKstar(fileName,treeName,True)
|
|
cutFile(fileName,treeName)
|
|
|
|
#PHSP
|
|
for pol in ["down","up"]:
|
|
for yr in years:
|
|
fileName = mainPath+"/PHSP/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
|
|
makeNewKstar(fileName,treeName,True)
|
|
cutFile(fileName,treeName)
|
|
|
|
#Data
|
|
for pol in ["down","up"]:
|
|
for yr in years:
|
|
fileName = mainPath+"/Data/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
|
|
makeNewKstar(fileName,treeName,False)
|
|
cutFile(fileName,treeName)
|
|
|
|
|
|
'''
|
|
#Inclusive
|
|
for pol in ["down","up"]:
|
|
for yr in years:
|
|
if (yr==2015): continue
|
|
if (yr>2016): continue
|
|
fileName = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoXJpsi/"+str(yr)+pol+"/"+str(yr)+pol+"_pi0Resolved"
|
|
makeNewKstar(fileName,treeName,True)
|
|
cutFile(fileName,treeName)
|
|
|
|
for yr in [2011,2012,2016]:
|
|
fileName1 = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoXJpsi/"+str(yr)+"down/"+str(yr)+"down_pi0Resolved.root"
|
|
fileName1 = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoXJpsi/"+str(yr)+"up/"+str(yr)+"up_pi0Resolved.root"
|
|
outname = "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/"+str(yr)+"_Inclusive_BDTinput.root"
|
|
mergeFiles(fileName1, fileName1, outname, "DecayTreeTruthMatched")
|
|
''' |