ATLAS Offline Software
Loading...
Searching...
No Matches
CSV_InDetExporter.CSVDumper Class Reference
Collaboration diagram for CSV_InDetExporter.CSVDumper:

Public Member Functions

 __init__ (self, inputAOD, outputDir, dict_variables_types, treename="CollectionTree", nEvents=-1)
 WriteCSV (self, filename, dictionary)
 ArrayFloat3_to_CppArray (self, ArrayFloat3)
 ProcessEvent (self, evt)
 Run (self)

Public Attributes

 tree = xAOD.MakeTransientTree( TFile(inputAOD), treename )
# Getting number of events to run n_entries = self.tree.GetEntriesFast()
 outputDir = outputDir
# Getting number of events to run nEvents = self.n_entries
 dict_variables_types = dict_variables_types

Detailed Description

Definition at line 12 of file CSV_InDetExporter.py.

Constructor & Destructor Documentation

◆ __init__()

CSV_InDetExporter.CSVDumper.__init__ ( self,
inputAOD,
outputDir,
dict_variables_types,
treename = "CollectionTree",
nEvents = -1 )

Definition at line 14 of file CSV_InDetExporter.py.

14 def __init__(self, inputAOD, outputDir, dict_variables_types, treename="CollectionTree", nEvents=-1):
15
16 xAOD.Init() # Setting up ROOT tools
17 self.tree = xAOD.MakeTransientTree( TFile(inputAOD), treename )
18 self.n_entries = self.tree.GetEntriesFast()
19 self.outputDir = outputDir
20
21 if nEvents < 0 or nEvents > self.n_entries: # Getting number of events to run
22 self.nEvents = self.n_entries
23 else:
24 self.nEvents = nEvents
25 print(f"Running on {self.nEvents} events")
26
27 os.system("mkdir -p "+self.outputDir) # Creating directory to store files
28
29 # Container and Variables with their types that should be stored in csv format. Should be something like:
30 #{
31 # "ContainerName": { "var1": "int",
32 # "var2": "float",
33 # }
34 #}
35 self.dict_variables_types = dict_variables_types
36
void print(char *figname, TCanvas *c1)
StatusCode Init(const char *appname)
Function initialising ROOT/PyROOT for using the ATLAS EDM.
Definition Init.cxx:31

Member Function Documentation

◆ ArrayFloat3_to_CppArray()

CSV_InDetExporter.CSVDumper.ArrayFloat3_to_CppArray ( self,
ArrayFloat3 )

Definition at line 45 of file CSV_InDetExporter.py.

45 def ArrayFloat3_to_CppArray(self, ArrayFloat3): # Check ArrayFloat3 in https://gitlab.cern.ch/atlas/athena/-/blob/master/Event/xAOD/xAODInDetMeasurement/xAODInDetMeasurement/versions/SpacePoint_v1.h
46 try:
47 arr = ArrayFloat3.data()
48 except Exception:
49 arr = ArrayFloat3
50 arr.reshape((3,))
51 return list(arr)
52
53

◆ ProcessEvent()

CSV_InDetExporter.CSVDumper.ProcessEvent ( self,
evt )

Definition at line 54 of file CSV_InDetExporter.py.

54 def ProcessEvent(self, evt):
55
56 # Getting aux variables based on the implementation of https://gitlab.cern.ch/atlas/athena/-/commit/fc5baf9fd2bb28c56f115fc2107a3ff159f1945d
57 print("--> Event "+str(evt))
58 self.tree.GetEntry(evt)
59 EventNumber = self.tree.EventInfo.mcEventNumber()
60
61 # Check whether the container exists
62 for container in self.dict_variables_types.keys():
63 dict_lists = {}
64 dict_container = self.dict_variables_types[container]
65 try:
66 tp = getattr(self.tree, container)
67 except Exception:
68 print(".. Missing ", container)
69 continue
70
71 if tp.size() == 0:
72 print(".. Empty ", container)
73 continue
74
75 for var,fmt in dict_container.items():
76 try:
77 sp = tp.getConstDataSpan[ fmt ]( var)
78 except Exception:
79 # This is for the case arrays are either empty or have some trouble when accessed via getConstDataSpan. Used in excepction as makes the code slower
80 print("getConstDataSpan failed for variable ",var,fmt, container)
81 sp = [ getattr(element, var)() for element in tp ]
82
83 try:
84 # Convert list of std::vector<double> to the list of lists
85 list_of_lists = [list(std_vector) for std_vector in sp]
86 # Find the length of the longest vector
87 max_len = max(len(ll) for ll in list_of_lists)
88 # extend lists to the length of max_len
89 for ll in list_of_lists:
90 ll.extend([np.nan] * (max_len - len(ll)))
91 except Exception:
92 list_of_lists = sp # the sp is not iterable
93
94 # Convert the list of lists to a NumPy array
95 sp = np.array(list_of_lists)
96
97 if "ArrayFloat3" in fmt and len(sp) > 0: # Needs extra coding when dealing with xAOD::ArrayFloat3 instead of standard C++ array
98 sp = array( list( map(self.ArrayFloat3_to_CppArray, sp) ) )
99
100 if ("unsigned char" in fmt or "uint8" in fmt) and len(sp) > 0:
101 sp = sp.view(np.int32)
102
103 if sp.ndim == 1:
104 dict_lists[var] = sp
105 else:
106 # Then we have an array. Each element of the array will be a column in the csv file. [We want things flat]
107 for column in range(sp.shape[1]):
108 dict_lists[var+f"_at{column}"] = sp.T[column]
109
110 self.WriteCSV( filename=getCSVFilename(self.outputDir, container, EventNumber), dictionary=dict_lists )
111
112
#define max(a, b)
Definition cfImp.cxx:41
STL class.
STL class.

◆ Run()

CSV_InDetExporter.CSVDumper.Run ( self)

Definition at line 113 of file CSV_InDetExporter.py.

113 def Run(self):
114
115 for evt in range(self.nEvents):
116 self.ProcessEvent(evt)
117
118

◆ WriteCSV()

CSV_InDetExporter.CSVDumper.WriteCSV ( self,
filename,
dictionary )

Definition at line 37 of file CSV_InDetExporter.py.

37 def WriteCSV(self, filename, dictionary):
38
39 with open(filename, "w") as out:
40 writer = csv.writer(out)
41 writer.writerow(dictionary.keys())
42 writer.writerows( zip(*dictionary.values()) )
43 print(f"New file saved: {filename}")
44

Member Data Documentation

◆ dict_variables_types

CSV_InDetExporter.CSVDumper.dict_variables_types = dict_variables_types

Definition at line 35 of file CSV_InDetExporter.py.

◆ n_entries

# Getting number of events to run CSV_InDetExporter.CSVDumper.n_entries = self.tree.GetEntriesFast()

Definition at line 18 of file CSV_InDetExporter.py.

◆ nEvents

CSV_InDetExporter.CSVDumper.nEvents = self.n_entries

Definition at line 22 of file CSV_InDetExporter.py.

◆ outputDir

CSV_InDetExporter.CSVDumper.outputDir = outputDir

Definition at line 19 of file CSV_InDetExporter.py.

◆ tree

CSV_InDetExporter.CSVDumper.tree = xAOD.MakeTransientTree( TFile(inputAOD), treename )

Definition at line 17 of file CSV_InDetExporter.py.


The documentation for this class was generated from the following file: