Logo ROOT   6.10/09
Reference Guide
parse_CSV_file_with_TTree_ReadStream.py
Go to the documentation of this file.
1 ## \file
2 ## \ingroup tutorial_pyroot
3 ## This function provides an example of how one might
4 ## massage a csv data file to read into a ROOT TTree
5 ## via TTree::ReadStream. This could be useful if the
6 ## data read out from some DAQ program doesn't 'quite'
7 ## match the formatting expected by ROOT (e.g. comma-
8 ## separated, tab-separated with white-space strings,
9 ## headers not matching the expected format, etc.)
10 ##
11 ## This example is shipped with a data
12 ## file that looks like:
13 ##
14 ## ~~~{.cpp}
15 ## Date/Time Synchro Capacity Temp.Cold Head Temp. Electrode HV Supply Voltage Electrode 1 Electrode 2 Electrode 3 Electrode 4
16 ## # Example data to read out. Some data have oddities that might need to
17 ## # dealt with, including the 'NaN' in Electrode 4 and the empty string in Date/Time (last row)
18 ## 08112010.160622 7 5.719000E-10 8.790500 24.237700 -0.008332 0 0 0 0
19 ## 8112010.160626 7 5.710000E-10 8.828400 24.237500 -0.008818 0 0 0 0
20 ## 08112010.160626 7 5.719000E-10 8.828400 24.237500 -0.008818 0 0 0 0
21 ## 08112010.160627 7 5.719000E-10 9.014300 24.237400 -0.028564 0 0 0 NaN
22 ## 08112010.160627 7 5.711000E-10 8.786000 24.237400 -0.008818 0 0 0 0
23 ## 08112010.160628 7 5.702000E-10 8.786000 24.237400 -0.009141 0 0 0 0
24 ## 08112010.160633 7 5.710000E-10 9.016200 24.237200 -0.008818 0 0 0 0
25 ## 7 5.710000E-10 8.903400 24.237200 -0.008818 0 0 0 0
26 ## ~~~
27 ##
28 ## These data require some massaging, including:
29 ##
30 ## - Date/Time has a blank ('') entry that must be handled
31 ## - The headers are not in the correct format
32 ## - Tab-separated entries with additional white space
33 ## - NaN entries
34 ##
35 ## \macro_code
36 ##
37 ## \author Michael Marino
38 
39 import ROOT
40 import sys
41 import os
42 
43 def parse_CSV_file_with_TTree_ReadStream(tree_name, afile):
44 
45 
46  ROOT.gROOT.SetBatch()
47  # The mapping dictionary defines the proper branch names and types given a header name.
48  header_mapping_dictionary = {
49  'Date/Time' : ('Datetime' , str) ,
50  'Synchro' : ('Synchro' , int) ,
51  'Capacity' : ('Capacitance' , float) ,
52  'Temp.Cold Head' : ('TempColdHead' , float) ,
53  'Temp. Electrode' : ('TempElectrode' , float) ,
54  'HV Supply Voltage' : ('HVSupplyVoltage', float) ,
55  'Electrode 1' : ('Electrode1' , int) ,
56  'Electrode 2' : ('Electrode2' , int) ,
57  'Electrode 3' : ('Electrode3' , int) ,
58  'Electrode 4' : ('Electrode4' , int) ,
59  }
60 
61  type_mapping_dictionary = {
62  str : 'C',
63  int : 'I',
64  float : 'F'
65  }
66 
67 
68 
69  # Grab the header row of the file. In this particular example,
70  # the data are separated using tabs, but some of the header names
71  # include spaces and are not generally in the ROOT expected format, e.g.
72  #
73  # FloatData/F:StringData/C:IntData/I
74  #
75  # etc. Therefore, we grab the header_row of the file, and use
76  # a python dictionary to set up the appropriate branch descriptor
77  # line.
78 
79  # Open a file, grab the first line, strip the new lines
80  # and split it into a list along 'tab' boundaries
81  header_row = open(afile).readline().strip().split('\t')
82  # Create the branch descriptor
83  branch_descriptor = ':'.join([header_mapping_dictionary[row][0]+'/'+
84  type_mapping_dictionary[header_mapping_dictionary[row][1]]
85  for row in header_row])
86  #print branch_descriptor
87 
88  # Handling the input and output names. Using the same
89  # base name for the ROOT output file.
90  output_ROOT_file_name = os.path.splitext(afile)[0] + '.root'
91  output_file = ROOT.TFile(output_ROOT_file_name, 'recreate')
92  print "Outputting %s -> %s" % (afile, output_ROOT_file_name)
93 
94  output_tree = ROOT.TTree(tree_name, tree_name)
95  file_lines = open(afile).readlines()
96 
97  # Clean the data entries: remove the first (header) row.
98  # Ensure empty strings are tagged as such since
99  # ROOT doesn't differentiate between different types
100  # of white space. Therefore, we change all of these
101  # entries to 'empty'. Also, avoiding any lines that begin
102  # with '#'
103  file_lines = ['\t'.join([val if (val.find(' ') == -1 and val != '')
104  else 'empty' for val in line.split('\t')])
105  for line in file_lines[1:] if line[0] != '#' ]
106 
107  # Removing NaN, setting these entries to 0.0.
108  # Also joining the list of strings into one large string.
109  file_as_string = ('\n'.join(file_lines)).replace('NaN', str(0.0))
110  #print file_as_string
111 
112  # creating an istringstream to pass into ReadStream
113  istring = ROOT.istringstream(file_as_string)
114 
115  # Now read the stream
116  output_tree.ReadStream(istring, branch_descriptor)
117 
118  output_file.cd()
119  output_tree.Write()
120 
121 
122 if __name__ == '__main__':
123  if len(sys.argv) < 2:
124  print "Usage: %s file_to_parse.dat" % sys.argv[0]
125  sys.exit(1)
126  parse_CSV_file_with_TTree_ReadStream("example_tree", sys.argv[1])
127