def test_rootcint():
generate('map<int,vector<float> >', '<vector>;<map>')
generate('map<int,vector<int> >', '<vector>;<map>')
generate('vector<TLorentzVector>', '<vector>;TLorentzVector.h')
ROOT.std.map('int,vector<float>')
ROOT.std.map('int,vector<int>')
ROOT.std.vector('TLorentzVector')
temp = CPPType.from_string('vector<vector<vector<int> > >')
temp.ensure_built()
stl.vector('vector<map<int, string> >')
stl.map("string", "string")
stl.map(stl.string, stl.string)
stl.map(int, stl.string)
stl.map(stl.string, int)
stl.map("string", ROOT.TLorentzVector)
histmap = stl.map("string", ROOT.TH1D)()
a = ROOT.TH1D("a", "a", 10, -1, 1)
histmap["a"] = a
StrHist = stl.pair(stl.string, "TH1*")
generate('pair<map<string,TH1*>::iterator,bool>', '<map>;<TH1.h>')
histptrmap = stl.map(stl.string, "TH1*")()
histptrmap.insert(StrHist("test", a))
assert histptrmap["test"] is a
class rEvent(TreeModel):
id = IntCol()
xhits = stl.vector("float")
yhits = stl.vector("float")
nHits = IntCol()
nTracks = IntCol()
nGoodTracks = IntCol()
def test_rootcint():
generate('map<int,vector<float> >', '<vector>;<map>')
generate('map<int,vector<int> >', '<vector>;<map>')
generate('vector<TLorentzVector>', '<vector>;TLorentzVector.h')
ROOT.std.map('int,vector<float>')
ROOT.std.map('int,vector<int>')
ROOT.std.vector('TLorentzVector')
temp = CPPType.from_string('vector<vector<vector<int> > >')
temp.ensure_built()
stl.vector('vector<map<int, string> >')
stl.map("string", "string")
stl.map(stl.string, stl.string)
stl.map(int, stl.string)
stl.map(stl.string, int)
stl.map("string", ROOT.TLorentzVector)
histmap = stl.map("string", ROOT.TH1D)()
a = ROOT.TH1D("a", "a", 10, -1, 1)
histmap["a"] = a
StrHist = stl.pair(stl.string, "TH1*")
generate('pair<map<string,TH1*>::iterator,bool>', '<map>;<TH1.h>')
histptrmap = stl.map(stl.string, "TH1*")()
histptrmap.insert(StrHist("test", a))
assert histptrmap["test"] is a
class ObjectB(TreeModel):
# A tree object collection
x = stl.vector('int')
y = stl.vector('float')
vect = stl.vector('TLorentzVector')
# collection size
n = IntCol()
def setUp(self):
self.ntuple_types = ['int', 'bool', 'uint', 'float']
self.ntuple_vtypes = ['vec:{}'.format(t) for t in self.ntuple_types]
# from
# https://github.com/rootpy/rootpy/blob/master/rootpy/tree/treetypes.py
self.root_types = [IntCol, BoolCol, UIntCol, FloatCol]
self.root_vtypes = [
vector('int'), vector('bool'), vector('uint'), vector('float')]
def setUp(self):
self.ntuple_types = ['int', 'bool', 'uint', 'float']
self.ntuple_vtypes = ['vec:{}'.format(t) for t in self.ntuple_types]
# from
# https://github.com/rootpy/rootpy/blob/master/rootpy/tree/treetypes.py
self.root_types = [IntCol, BoolCol, UIntCol, FloatCol]
self.root_vtypes = [
vector('int'),
vector('bool'),
vector('uint'),
vector('float')
]
def __init__(self, tree, **kwargs):
super(PileupReweight_xAOD, self).__init__(**kwargs)
self.tree = tree
self.pileup_tool = ROOT.CP.PileupReweightingTool('pileup_tool')
mc_vec = stl.vector('string')()
data_vec = stl.vector('string')()
mc_vec.push_back('PileupReweighting/mc14v1_defaults.prw.root')
data_vec.push_back('lumi/2012/ilumicalc_histograms_None_200842-215643.root')
self.pileup_tool.setProperty('ConfigFiles', mc_vec)
self.pileup_tool.setProperty('LumiCalcFiles', data_vec)
self.pileup_tool.setProperty('int')('DefaultChannel', 161656)
self.pileup_tool.initialize()
class Event(TreeModel):
# properties of particle "a"
a_x = FloatCol()
a_y = FloatCol()
a_z = FloatCol()
# properties of particle "b"
b_x = FloatCol()
b_y = FloatCol()
b_z = FloatCol()
# a collection of particles
col_x = stl.vector("float")
col_y = stl.vector("float")
col_z = stl.vector("float")
col_n = IntCol()
# a TLorentzVector
p = LorentzVector
i = IntCol()
class rTrack(TreeModel):
id = IntCol()
event_id = IntCol()
hit_indices = stl.vector("int")
residuals = stl.vector("float")
lincoords = stl.vector("float")
color = IntCol()
length = FloatCol()
rho = FloatCol()
theta = FloatCol()
x0 = FloatCol()
y0 = FloatCol()
x1 = FloatCol()
y1 = FloatCol()
nHits = IntCol()
R2 = FloatCol()
is_good = IntCol()
def write_to_root(tracks, laser):
""" Writes tracks and laser data to a root file which is readable by the reconstruction algorithm """
from rootpy.vector import Vector3
import rootpy.stl as stl
from rootpy.tree import Tree
from rootpy.io import root_open
laser_entry, laser_exit = laser
Vec = stl.vector(Vector3)
track = Vec()
laserentry = Vector3(laser_entry[0], laser_entry[1], laser_entry[2])
laserexit = Vector3(laser_exit[0], laser_exit[1], laser_exit[2])
f = root_open("test.root", "recreate")
track_tree = Tree('tracks')
laser_tree = Tree('lasers')
track_tree.create_branches({'track': stl.vector(Vector3)})
laser_tree.create_branches({'entry': Vector3,
'exit': Vector3})
for k in range(10):
print(k)
for i in range(1000):
track.push_back(Vector3(i, k, k * i))
track_tree.track = track
track.clear()
laser_tree.entry = Vector3(0, 0, 0)
laser_tree.exit = Vector3(k, k, k)
track_tree.fill()
laser_tree.fill()
track_tree.write()
laser_tree.write()
f.close()
def getClusters(self, eff=False):
"""Returns the clusters (all or effective)"""
clusters = stl.vector( 'TLorentzVector' )()
n = self._tree.tau_cluster_n[self._recoIndex]
if eff: n = int(math.ceil(self._tree.tau_numEffTopoClusters[self._recoIndex]))
for i in range(0, n):
E = self._tree.tau_cluster_E[self._recoIndex][i]
eta = self._tree.tau_cluster_eta_atTJVA[self._recoIndex][i]
phi = self._tree.tau_cluster_phi_atTJVA[self._recoIndex][i]
pt = E/math.cosh(eta)
vect = TLorentzVector()
vect.SetPtEtaPhiM(pt, eta, phi, 0)
clusters.push_back(vect)
return clusters
def getClusters(self, eff=False):
"""Returns the clusters (all or effective)"""
clusters = stl.vector('TLorentzVector')()
n = self._tree.tau_cluster_n[self._recoIndex]
if eff:
n = int(
math.ceil(self._tree.tau_numEffTopoClusters[self._recoIndex]))
for i in range(0, n):
E = self._tree.tau_cluster_E[self._recoIndex][i]
eta = self._tree.tau_cluster_eta_atTJVA[self._recoIndex][i]
phi = self._tree.tau_cluster_phi_atTJVA[self._recoIndex][i]
pt = E / math.cosh(eta)
vect = TLorentzVector()
vect.SetPtEtaPhiM(pt, eta, phi, 0)
clusters.push_back(vect)
return clusters
def getTracks(self):
"""Returns the seedCalo tracks"""
tracks = stl.vector( 'TLorentzVector' )()
d0 = []
z0 = []
n = self._tree.tau_track_n[self._recoIndex]
for i in range(0, n):
# pt = self._tree.tau_track_atTJVA_pt[self._recoIndex][i]
# eta = self._tree.tau_track_atTJVA_eta[self._recoIndex][i]
# phi = self._tree.tau_track_atTJVA_phi[self._recoIndex][i]
pt = self._tree.tau_track_atTJVA_pt[self._recoIndex][i]
eta = self._tree.tau_track_atTJVA_eta[self._recoIndex][i]
phi = self._tree.tau_track_atTJVA_phi[self._recoIndex][i]
m = 140
d0.append(self._tree.tau_track_atTJVA_d0[self._recoIndex][i])
z0.append(self._tree.tau_track_atTJVA_z0[self._recoIndex][i])
if pt > 0.0 and eta < 5.0:
vec = TLorentzVector()
vec.SetPtEtaPhiM(pt, eta, phi, 0.)
tracks.push_back(vec)
return tracks,d0,z0
def getTracks(self):
"""Returns the seedCalo tracks"""
tracks = stl.vector('TLorentzVector')()
d0 = []
z0 = []
n = self._tree.tau_track_n[self._recoIndex]
for i in range(0, n):
# pt = self._tree.tau_track_atTJVA_pt[self._recoIndex][i]
# eta = self._tree.tau_track_atTJVA_eta[self._recoIndex][i]
# phi = self._tree.tau_track_atTJVA_phi[self._recoIndex][i]
pt = self._tree.tau_track_atTJVA_pt[self._recoIndex][i]
eta = self._tree.tau_track_atTJVA_eta[self._recoIndex][i]
phi = self._tree.tau_track_atTJVA_phi[self._recoIndex][i]
m = 140
d0.append(self._tree.tau_track_atTJVA_d0[self._recoIndex][i])
z0.append(self._tree.tau_track_atTJVA_z0[self._recoIndex][i])
if pt > 0.0 and eta < 5.0:
vec = TLorentzVector()
vec.SetPtEtaPhiM(pt, eta, phi, 0.)
tracks.push_back(vec)
return tracks, d0, z0
def add_branch(arr, filename, tree = 'bTag_AntiKt2PV0TrackJets', branchname = 'jet_mv2c20_new'):
'''
writes the newly evaluated mv2 scores into a branch in a Friend Tree
# --------------------------------------------------------------------------------
# -- *WARNING*: Make sure the file you are trying to modify is *NOT* already open!
# Otherwise, instead of adding a branch to that file, you will
# corrupt the file!
# --------------------------------------------------------------------------------
Args:
-----
arr: array containg newly evaluated mv2 scores
filename: .root file where new branch will be added
tree: (optional) name of TTree that will get a new Friend
branchname: (optional) name of the new branch
'''
# -- Check if file already exists:
if not os.path.exists(filename):
print '[WARNING] file not found, creating new file'
# -- Open file:
f = root_open(filename, "update")
# -- Extract TTree
T = f[tree]
# -- Need to figure out dtype in order to save the branch correctly
# -- If dtype is wrong, ROOT returns garbage, so this below is important!
# -- Case of branch being event level values:
if 'float' in str(type(arr[0])):
if '64' in str(type(arr[0])):
dtype = 'double'
else:
dtype = 'float'
elif 'int' in str(type(arr[0])):
dtype = 'int'
# -- Case of branch being jet level list:
elif hasattr(arr[0], '__iter__'):
if 'float' in str(type(arr[0][0])):
if '64' in str(type(arr[0][0])):
dtype = 'double'
else:
dtype = 'float'
elif 'int' in str(type(arr[0][0])):
dtype = 'int'
else:
raise TypeError('Nested type `{}` not supported'.format(str(type(arr[0][0]))))
dtype = 'vector<{}>'.format(dtype)
else:
raise TypeError('Type `{}` not supported'.format(str(type(arr[0]))))
sys.stdout.write('Detected dtype: {}\n'.format(dtype))
sys.stdout.flush()
# -- Create friend:
T_friend = Tree(tree+'_Friend')
# -- Add new branch to friend tree:
T_friend.create_branches({branchname : dtype})
# -- Fill the branch:
sys.stdout.write('Filling branch "{}" ... \n'.format(branchname))
sys.stdout.flush()
for i, branch4event in enumerate(arr):
if 'vector' in dtype:
buf = stl.vector(dtype.replace('vector<', '').replace('>', ''))()
_ = [buf.push_back(e) for e in branch4event]
exec('T_friend.{} = buf'.format(branchname))
else:
exec('T_friend.{} = branch4event'.format(branchname))
T_friend.Fill()
# -- Write out the tree and close the file:
sys.stdout.write('Finalizing and closing file "{}" \n'.format(filename))
sys.stdout.flush()
T.AddFriend(T_friend, tree+'_Friend')
T_friend.Write()
f.Write()
f.Close()
def test_multiple(self):
variables = []
for _ in range(5):
variables.append(convert_type('vec:float'))
for v in variables:
self.assertEqual(v, vector('float'))
import rootpy.stl as stl
import ROOT
# Create a vector type
StrVector = stl.vector(stl.string)
# Instantiate
strvector = StrVector()
strvector.push_back("Hello")
# etc.
MapStrRoot = stl.map(stl.string, ROOT.TH1D)
MapStrRootPtr = stl.map(stl.string, "TH1D*")
StrHist = stl.pair(stl.string, "TH1*")
m = MapStrRootPtr()
a = ROOT.TH1D('t1', 't1', 10, 0, 1)
m.insert(StrHist("test", a))
print m
import ROOT as root
from ROOT import gROOT
import sys, os.path, os
from rootpy import stl
from rootpy.io import File
from rootpy.tree import Tree, TreeModel, TreeChain
from rootpy.plotting import Hist, Canvas
import neriX_config, neriX_led_datasets, neriX_pmt_gain_corrections
import threading, pickle, time, array
# MUST INCLUDE TYPES TO AVOID SEG FAULT
stl.vector(stl.vector('float'))
stl.vector(stl.vector('int'))
def convert_name_to_unix_time(fileName):
if fileName[-5:] == '.root':
fileName = fileName[:-5]
# check that we know position of dates in name
if fileName[0:5] == 'nerix':
sTimeTaken = '20' + fileName[6:17] + '00'
return int(time.mktime(time.strptime(sTimeTaken, '%Y%m%d_%H%M%S')))
elif fileName[0:8] == 'ct_nerix':
sTimeTaken = '20' + fileName[9:20] + '00'
return int(time.mktime(time.strptime(sTimeTaken, '%Y%m%d_%H%M%S')))
else:
print 'Only know how to handle nerix and ct_nerix files currently'
print 'Please add to convert_name_to_unix_time function in order to handle the new case'
sys.exit()
class GenTree(TreeModel):
p_NuGen = FloatCol()
pt_NuGen = FloatCol()
phi_NuGen = FloatCol()
eta_NuGen = FloatCol()
aperture_Gen = FloatCol()
p_Gen = stl.vector(float)
pt_Gen = stl.vector(float)
px_Gen = stl.vector(float)
py_Gen = stl.vector(float)
pz_Gen = stl.vector(float)
eta_Gen = stl.vector(float)
phi_Gen = stl.vector(float)
theta_Gen = stl.vector(float)
mass_Gen = stl.vector(float)
e_Gen = stl.vector(float)
p_Reco = stl.vector(float)
pt_Reco = stl.vector(float)
px_Reco = stl.vector(float)
py_Reco = stl.vector(float)
pz_Reco = stl.vector(float)
vx_Reco = stl.vector(float)
vy_Reco = stl.vector(float)
vz_Reco = stl.vector(float)
eta_Reco = stl.vector(float)
etaErr_Reco = stl.vector(float)
phi_Reco = stl.vector(float)
phiErr_Reco = stl.vector(float)
theta_Reco = stl.vector(float)
thetaErr_Reco = stl.vector(float)
idx = stl.vector(int)
pdgId = stl.vector(int)
dr = stl.vector(float)
dx_Gen = stl.vector(float)
dy_Gen = stl.vector(float)
dz_Gen = stl.vector(float)
algo = stl.vector(int)
# 'dxy', 'dxyError', 'dz', 'dzError'
dxy_RecoPV = stl.vector(float)
dxyErr_RecoPV = stl.vector(float)
dz_RecoPV = stl.vector(float)
dzErr_RecoPV = stl.vector(float)
kaonCharged = IntCol()
pionCharged = IntCol()
kaonNeutrals = IntCol()
pionNeutrals = IntCol()
muons = IntCol()
electrons = IntCol()
neutrals = IntCol()
charged = IntCol()
dMeson = IntCol()
pt_Jet = FloatCol()
eta_Jet = FloatCol()
phi_Jet = FloatCol()
mass_Jet = FloatCol()
et_Jet = FloatCol()
deposits_Jet = stl.vector(float)
import ROOT
from rootpy.tree.filtering import *
from rootpy.extern.hep import pdg
from rootpy import stl
VectorTLorentzVector = stl.vector("TLorentzVector")
from itertools import ifilter
from math import *
from . import datasets
from .corrections import reweight_ggf
from .units import GeV
from .tautools import TauDecay
from . import jetcleaning
from . import utils
from . import log; log = log[__name__]
from goodruns import GRL
BCH_TOOLS = []
class GRLFilter(EventFilter):
def __init__(self, grl, **kwargs):
super(GRLFilter, self).__init__(**kwargs)
if isinstance(grl, GRL):
self.grl = grl
else:
import rootpy.stl as stl
import ROOT
# Create a vector type
StrVector = stl.vector(stl.string)
# Instantiate
strvector = StrVector()
strvector.push_back("Hello")
# etc.
MapStrRoot = stl.map(stl.string, ROOT.TH1D)
MapStrRootPtr = stl.map(stl.string, "TH1D*")
StrHist = stl.pair(stl.string, "TH1*")
m = MapStrRootPtr()
a = ROOT.TH1D('t1', 't1', 10, 0, 1)
m.insert(StrHist("test", a))
print m
def test_multiple(self):
variables = []
for _ in range(5):
variables.append(convert_type('vec:float'))
for v in variables:
self.assertEqual(v, vector('float'))
class Event(TreeModel):
x = stl.vector('TLorentzVector')
i = IntCol()
base_dir = '/home/data/uboone/laser/processed/'
side = "upstream"
downsample = 1
track_filenames = ["/sim/laser-tracks-24.npy"]
laser_filenames = ["/sim/laser-data-24.npy"]
track_paths = [base_dir + filename for filename in track_filenames]
laser_paths = [base_dir + filename for filename in laser_filenames]
sides = {"upstream": 1, "downstream": 2}
output_file = base_dir + "beams/" + "laserbeams-24.root"
Vec = stl.vector(Vector3)
track = Vec()
with root_open(output_file, "recreate"):
track_tree = Tree('tracks')
laser_tree = Tree('lasers')
track_tree.create_branches({'track': stl.vector(Vector3), 'event': 'I'})
laser_tree.create_branches({
'entry': Vector3,
'exit': Vector3,
'side': 'I'
}) # 1: upstream, 2: down
for track_file, laser_file in zip(track_paths, laser_paths):
track_data = np.load(track_file)
laser_data = np.load(laser_file)
def add_branch(arr,
filename,
tree='bTag_AntiKt2PV0TrackJets',
branchname='jet_mv2c20_new'):
'''
writes the newly evaluated mv2 scores into a branch in a Friend Tree
# --------------------------------------------------------------------------------
# -- *WARNING*: Make sure the file you are trying to modify is *NOT* already open!
# Otherwise, instead of adding a branch to that file, you will
# corrupt the file!
# --------------------------------------------------------------------------------
Args:
-----
arr: array containg newly evaluated mv2 scores
filename: .root file where new branch will be added
tree: (optional) name of TTree that will get a new Friend
branchname: (optional) name of the new branch
'''
# -- Check if file already exists:
if not os.path.exists(filename):
print '[WARNING] file not found, creating new file'
# -- Open file:
f = root_open(filename, "update")
# -- Extract TTree
T = f[tree]
# -- Need to figure out dtype in order to save the branch correctly
# -- If dtype is wrong, ROOT returns garbage, so this below is important!
# -- Case of branch being event level values:
if 'float' in str(type(arr[0])):
if '64' in str(type(arr[0])):
dtype = 'double'
else:
dtype = 'float'
elif 'int' in str(type(arr[0])):
dtype = 'int'
# -- Case of branch being jet level list:
elif hasattr(arr[0], '__iter__'):
if 'float' in str(type(arr[0][0])):
if '64' in str(type(arr[0][0])):
dtype = 'double'
else:
dtype = 'float'
elif 'int' in str(type(arr[0][0])):
dtype = 'int'
else:
raise TypeError('Nested type `{}` not supported'.format(
str(type(arr[0][0]))))
dtype = 'vector<{}>'.format(dtype)
else:
raise TypeError('Type `{}` not supported'.format(str(type(arr[0]))))
sys.stdout.write('Detected dtype: {}\n'.format(dtype))
sys.stdout.flush()
# -- Create friend:
T_friend = Tree(tree + '_Friend')
# -- Add new branch to friend tree:
T_friend.create_branches({branchname: dtype})
# -- Fill the branch:
sys.stdout.write('Filling branch "{}" ... \n'.format(branchname))
sys.stdout.flush()
for i, branch4event in enumerate(arr):
if 'vector' in dtype:
buf = stl.vector(dtype.replace('vector<', '').replace('>', ''))()
_ = [buf.push_back(e) for e in branch4event]
exec('T_friend.{} = buf'.format(branchname))
else:
exec('T_friend.{} = branch4event'.format(branchname))
T_friend.Fill()
# -- Write out the tree and close the file:
sys.stdout.write('Finalizing and closing file "{}" \n'.format(filename))
sys.stdout.flush()
T.AddFriend(T_friend, tree + '_Friend')
T_friend.Write()
f.Write()
f.Close()
import ROOT
from rootpy.tree.filtering import *
from rootpy.extern.hep import pdg
from rootpy import stl
VectorTLorentzVector = stl.vector("TLorentzVector")
Vector = stl.vector('float')
from itertools import ifilter
from math import *
from array import array as carray
from xaod import TOOLS
from . import datasets
# from .corrections import reweight_ggf
from .units import GeV
from .tautools import TauDecay
from . import utils
from . import store_helper
from . import log; log = log[__name__]
from goodruns import GRL
BCH_TOOLS = []
class GRLFilter(EventFilter):
def __init__(self, grl, **kwargs):
super(GRLFilter, self).__init__(**kwargs)
if isinstance(grl, GRL):
#gc.disable()
"""
import warnings
#warnings.filterwarnings('error')
# this turns runtime warnings into errors to
# catch unexpected behavior
warnings.simplefilter('error', RuntimeWarning)
# this turns off a deprecation warning in general but aimed
# at corners plot creation (should try to only put there)
# and move back to all warnings as errors)
warnings.simplefilter('ignore', DeprecationWarning)
warnings.simplefilter('ignore', FutureWarning)
"""
# MUST INCLUDE TYPES TO AVOID SEG FAULT
stl.vector(stl.vector('float'))
stl.vector(stl.vector('double'))
stl.vector(stl.vector('int'))
# create dot product in c++ for speed
# and thread safety
C.register_file('../../python_modules/mc_code/c_safe_dot.C', ['safe_dot'])
def reduce_method(m):
return (getattr, (m.__self__, m.__func__.__name__))
class fit_nr(object):
def __init__(self,
class Decorations(TreeModel):
jet_ispileup = stl.vector('bool')
# a TLorentzVector
p = LorentzVector
i = IntCol()
# Make two files, each with a Tree called "test"
print "Creating test tree in chaintest1.root"
tmp_dir = tempfile.mkdtemp()
f = root_open(tmp_dir + "/chaintest1.root", "recreate")
branches = {
'x': FloatCol(),
'y': FloatCol(),
'z': FloatCol(),
'i': FloatCol(),
'vi': stl.vector('float'),
'vx': FloatArrayCol(4),
'vy': stl.vector('float'), }
# print branches
MyTreeModel = TreeModelMeta('MyTreeModel', (TreeModel,), branches)
tree = Tree("test", model=MyTreeModel)
# tree.create_branches(branches)
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
for vi in range(4):
tree.vi.push_back(vi**2)
tree.vy.push_back(vi**3)
