94def modify_file(inputfile, muonfile) -> None:
95 """Add information from original muon to input simulation file."""
96 logging.warning(
97 f"vetoPoints & UpstreamTaggerPoints from the incoming muon (saved in {muonfile}) will be added to {inputfile}."
98 )
99
100 input_file = r.TFile.Open(inputfile, "read")
101 try:
102 input_tree = input_file.cbmsim
103 except Exception as e:
104 print(f"Error: {e}")
105 input_file.Close()
106 exit(1)
107
108
109 muon_file = r.TFile.Open(muonfile, "read")
110 try:
111 muon_tree = muon_file.DIS
112 except Exception as e:
113 print(f"Error: {e}")
114 muon_file.Close()
115 exit(1)
116
117 temp_filename = inputfile.replace(".root", ".tmp")
118 temp_file = r.TFile.Open(temp_filename, "recreate")
119 output_tree = input_tree.CloneTree(0)
120
121 combined_vetoPoint = r.TClonesArray("vetoPoint")
122 output_tree.SetBranchAddress("vetoPoint", combined_vetoPoint)
123
124 combined_UpstreamTaggerPoint = r.TClonesArray("UpstreamTaggerPoint")
125 output_tree.SetBranchAddress("UpstreamTaggerPoint", combined_UpstreamTaggerPoint)
126
127 table_data = []
128
129 for i, (input_event, muon_event) in enumerate(zip(input_tree, muon_tree)):
130 interaction_point = r.TVector3()
131 input_event.MCTrack[0].GetStartVertex(interaction_point)
132
133 combined_vetoPoint.Clear()
134
135 index = 0
136
137 for hit in input_event.vetoPoint:
138 if combined_vetoPoint.GetSize() == index:
139 combined_vetoPoint.Expand(index + 1)
140 combined_vetoPoint[index] = hit
141 index += 1
142
143 muoncount = 0
144 for hit in muon_event.muon_vetoPoints:
145 if hit.GetZ() < interaction_point.Z():
146 if combined_vetoPoint.GetSize() == index:
147 combined_vetoPoint.Expand(index + 1)
148 combined_vetoPoint[index] = hit
149 index += 1
150 muoncount += 1
151
152 combined_UpstreamTaggerPoint.Clear()
153
154 ubt_index = 0
155
156 for hit in input_event.UpstreamTaggerPoint:
157 if combined_UpstreamTaggerPoint.GetSize() == ubt_index:
158 combined_UpstreamTaggerPoint.Expand(ubt_index + 1)
159 combined_UpstreamTaggerPoint[ubt_index] = hit
160 ubt_index += 1
161
162 muon_ubtcount = 0
163 for hit in muon_event.muon_UpstreamTaggerPoints:
164 if hit.GetZ() < interaction_point.Z():
165 if combined_UpstreamTaggerPoint.GetSize() == ubt_index:
166 combined_UpstreamTaggerPoint.Expand(ubt_index + 1)
167 combined_UpstreamTaggerPoint[ubt_index] = hit
168 ubt_index += 1
169 muon_ubtcount += 1
170
171 table_data.append(
172 [
173 i,
174 len(muon_event.muon_vetoPoints),
175 len(input_event.vetoPoint),
176 muoncount,
177 len(combined_vetoPoint),
178 len(muon_event.muon_UpstreamTaggerPoints),
179 len(input_event.UpstreamTaggerPoint),
180 muon_ubtcount,
181 len(combined_UpstreamTaggerPoint),
182 ]
183 )
184 output_tree.Fill()
185
186 output_tree.Write("cbmsim", r.TObject.kOverwrite)
187 temp_file.Close()
188 input_file.Close()
189 muon_file.Close()
190
191 os.remove(inputfile)
192 os.rename(temp_filename, inputfile)
193 print(f"File updated with incoming muon info as {inputfile}")
194 print(tabulate(table_data, headers=headers, tablefmt="grid"))
195
196