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()
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:
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)