Cicada  v1.4.1
Project 8 Processed-Data File Format Library
ReadKTOutputFile.py
Go to the documentation of this file.
1 #!/usr/bin/env python
2 
3 '''
4 Name: ReadKTOutputFile
5 Author: MG
6 Date: 2018/05/10
7 Purpose:
8  Example of how to open and read a Katydid output file.
9  Can read
10 '''
11 
12 import CicadaPy
13 CicadaPy.loadLibraries(True)
14 
15 from ROOT import TFile, TTreeReader, TTreeReaderValue
16 import sys
17 
18 def ReadKTOutputFile(filename,var,katydid=False,objectType="TMultiTrackEventData",name="multiTrackEvents:Event"):
19  '''
20  Read a ROOT file and returns specified content.
21  filename (str): name of the ROOT file
22  var (str/list): Variable to extract from the object
23  katydid (bool): Katydid (true) or Cicada (false) namespace
24  objectType (str): class name of the object to read
25  name (str:str): names of the tree and of the object to read, separated by a ":"
26  '''
27  # Open the ROOT file
28  file = TFile.Open(filename)
29  if not file:
30  raise FileNotFoundError("File {} does not exist".format(filename))
31  # Extract tree and object names
32  # names[0]: name of the tree
33  # names[1]: name of the object in the tree
34  names = name.split(":")
35  # Check if tree is in file
36  if not file.GetListOfKeys().Contains(str(names[0])):
37  print("Error: no tree {} in file".format(str(names[0])))
38  return {}
39  # Extract tree from file
40  tree = file.Get(str(names[0]))
41  # Create TTreeReader
42  treeReader = TTreeReader(tree)
43  # Create object TMultiTrackEventData to "point" to the object "Event" in the tree
44  if katydid:
45  import ROOT.Katydid as KT
46  else:
47  import ROOT.Cicada as KT
48  if hasattr(KT,objectType):
49  multiTrackEvents = TTreeReaderValue(getattr(KT,objectType))(treeReader, names[1])
50  else:
51  print("Couldn't find type <{}> in module; exiting!".format(objectType))
52  sys.exit(1)
53 
54  result = {}
55  if isinstance(var,list):
56  for key in var:
57  result.update({key:[]})
58  elif isinstance(var,str):
59  result.update({var: []})
60  else:
61  print("{} not a list or a string; exiting!".format(var))
62 
63  for var,_ in result.items():
64  # Analyze var to see if we are looking for a subTree
65  subArg = var.split(".")
66 
67  n=0
68  # Go through the events
69  if len(subArg)==1:
70  while treeReader.Next():
71  function = getattr(multiTrackEvents,"Get{}".format(subArg[0]))
72  result[var].append(function())
73  n+=1
74  else:
75  while treeReader.Next():
76  function = getattr(multiTrackEvents,"Get{}".format(subArg[0]))
77  obj = function()
78  subList = []
79  for i in range(multiTrackEvents.GetTracks().GetEntries()):
80  function = getattr(getattr(obj,"At")(i),"Get{}".format(subArg[1]))
81  subList.append(function())
82  result[var].append(subList)
83  n+=1
84  treeReader.Restart()
85  if n==0:
86  print("Error: no data found; wrong branch? or wrong namespace (Cicada/Katydid) -- maybe, try '-k'/add True as last argument?")
87  if len(result.keys())==1:
88  return result[list(result.keys())[0]]
89  return result
90 
91 if __name__ =="__main__":
92  print('\n P8 Katydid output file reader example\n')
93  import argparse
94  p = argparse.ArgumentParser(description='''
95  Katydid output file reader script
96  ''')
97  p.add_argument('-i','--input',
98  help='ROOT input file',
99  type=str)
100  p.add_argument('-b','--branch',
101  help='Branch to read in the TMultiTrackEventData object',
102  nargs='+',
103  default="StartTimeInAcq")
104  p.add_argument('-k','--katydid',
105  help='Flag stating Katydid namespace should be used instead of Cicada',
106  action='store_true',
107  required=False,
108  default=False)
109  p.add_argument('-t','--type',
110  help='Object Type to be read',
111  type=str,
112  required=False,
113  default="TMultiTrackEventData")
114  p.add_argument('-n','--name',
115  help='Object and Tree Name to be read, separated by a ":" (ex: multiTrackEvents:Event)',
116  type=str,
117  required=False,
118  default="multiTrackEvents:Event")
119  args = p.parse_args()
120 
121  result = ReadKTOutputFile(args.input, args.branch, katydid=args.katydid, objectType=args.type, name=args.name)
122  if isinstance(result,dict):
123  nElements=len(result[list(result.keys())[0]])
124  else:
125  nElements = len(result)
126  print("Found {} elements for {}".format(nElements,args.branch))
def ReadKTOutputFile(filename, var, katydid=False, objectType="TMultiTrackEventData", name="multiTrackEvents:Event")