This function provides an example of how one might massage a csv data file to read into a ROOT TTree via TTree::ReadStream.
This could be useful if the data read out from some DAQ program doesn't 'quite' match the formatting expected by ROOT (e.g. comma- separated, tab-separated with white-space strings, headers not matching the expected format, etc.)
import ROOT
import sys
import os
def parse_CSV_file_with_TTree_ReadStream(tree_name, afile):
ROOT.gROOT.SetBatch()
# The mapping dictionary defines the proper branch names and types given a header name.
header_mapping_dictionary = {
'Date/Time' : ('Datetime' , str) ,
'Synchro' : ('Synchro' , int) ,
'Capacity' : ('Capacitance' , float) ,
'Temp.Cold Head' : ('TempColdHead' , float) ,
'Temp. Electrode' : ('TempElectrode' , float) ,
'HV Supply Voltage' : ('HVSupplyVoltage', float) ,
'Electrode 1' : ('Electrode1' , int) ,
'Electrode 2' : ('Electrode2' , int) ,
'Electrode 3' : ('Electrode3' , int) ,
'Electrode 4' : ('Electrode4' , int) ,
}
type_mapping_dictionary = {
str : 'C',
int : 'I',
float : 'F'
}
# Grab the header row of the file. In this particular example,
# the data are separated using tabs, but some of the header names
# include spaces and are not generally in the ROOT expected format, e.g.
#
# FloatData/F:StringData/C:IntData/I
#
# etc. Therefore, we grab the header_row of the file, and use
# a python dictionary to set up the appropriate branch descriptor
# line.
# Open a file, grab the first line, strip the new lines
# and split it into a list along 'tab' boundaries
header_row = open(afile).readline().strip().split('\t')
# Create the branch descriptor
branch_descriptor = ':'.join([header_mapping_dictionary[row][0]+'/'+
type_mapping_dictionary[header_mapping_dictionary[row][1]]
for row in header_row])
#print branch_descriptor
# Handling the input and output names. Using the same
# base name for the ROOT output file.
output_ROOT_file_name = os.path.splitext(afile)[0] + '.root'
output_file = ROOT.TFile(output_ROOT_file_name, 'recreate')
print "Outputting %s -> %s" % (afile, output_ROOT_file_name)
output_tree = ROOT.TTree(tree_name, tree_name)
file_lines = open(afile).readlines()
# Clean the data entries: remove the first (header) row.
# Ensure empty strings are tagged as such since
# ROOT doesn't differentiate between different types
# of white space. Therefore, we change all of these
# entries to 'empty'. Also, avoiding any lines that begin
# with '#'
file_lines = ['\t'.join([val if (val.find(' ') == -1 and val != '')
else 'empty' for val in line.split('\t')])
for line in file_lines[1:] if line[0] != '#' ]
# Removing NaN, setting these entries to 0.0.
# Also joining the list of strings into one large string.
file_as_string = ('\n'.join(file_lines)).replace('NaN', str(0.0))
#print file_as_string
# creating an istringstream to pass into ReadStream
istring = ROOT.istringstream(file_as_string)
# Now read the stream
output_tree.ReadStream(istring, branch_descriptor)
output_file.cd()
output_tree.Write()
if __name__ == '__main__':
if len(sys.argv) < 2:
print "Usage: %s file_to_parse.dat" % sys.argv[0]
sys.exit(1)
parse_CSV_file_with_TTree_ReadStream("example_tree", sys.argv[1])