HepMC3 event record library
ReaderuprootTree.cc
1// -*- C++ -*-
2//
3// This file is part of HepMC
4// Copyright (C) 2014-2023 The HepMC collaboration (see AUTHORS for details)
5//
6#include "ReaderuprootTree.h"
7namespace HepMC3
8{
9
10HEPMC3_DECLARE_READER_FILE(ReaderuprootTree)
11
12/// @brief obtain vector of objects using name and type
13template <class T> std::vector<T>
14ReaderuprootTree::get_vector(PyObject * file_name, const std::string& array_name, std::string desired_type) {
15 int i = m_events_count;
16 PyObject *pFunc = m_access_function;
17 PyObject * pArgs = PyTuple_New(4);
18 PyTuple_SetItem(pArgs, 0, file_name);
19 PyTuple_SetItem(pArgs, 1, Py_BuildValue("s#", array_name.c_str(), array_name.length()));
20 PyTuple_SetItem(pArgs, 2, Py_BuildValue("i", i));
21 PyTuple_SetItem(pArgs, 3, Py_BuildValue("s#", desired_type.c_str(), desired_type.length()));
22 PyObject *pReturn = PyObject_CallObject(pFunc, pArgs);
23 auto *np_ret = reinterpret_cast<PyArrayObject*>(pReturn);
24 std::vector<T> out;
25 int len0 = 0;
26 if (np_ret) len0 = PyArray_SHAPE(np_ret)[0];
27 if (len0 > 0) {
28 int len = PyArray_SHAPE(np_ret)[1];
29 T* c_out = reinterpret_cast<T*>(PyArray_DATA(np_ret));
30 for (int i = 0; i < len; i++) out.push_back(c_out[i]);
31 }
32 Py_DECREF(pArgs);
33 if (np_ret) Py_DECREF(np_ret);
34 return out;
35}
36
37/// @brief obtain vector of objects using name and type, specified for std::string
38template <>
39std::vector<std::string>
40ReaderuprootTree::get_vector<std::string>(PyObject * file_name, const std::string& array_name, std::string desired_type) {
41 if (desired_type.length() == 0) desired_type = "U500";
42 int i = m_events_count;
43 PyObject *pFunc = m_access_function;
44 PyObject * pArgs = PyTuple_New(4);
45 PyTuple_SetItem(pArgs, 0, file_name);
46 PyTuple_SetItem(pArgs, 1, Py_BuildValue("s#", array_name.c_str(), array_name.length()));
47 PyTuple_SetItem(pArgs, 2, Py_BuildValue("i", i));
48 PyTuple_SetItem(pArgs, 3, Py_BuildValue("s#", desired_type.c_str(), desired_type.length()));
49
50 PyObject *pReturn = PyObject_CallObject(pFunc, pArgs);
51 auto *np_ret = reinterpret_cast<PyArrayObject*>(pReturn);
52 std::vector<std::string> out;
53 int len0 = 0;
54 if (np_ret) len0 = PyArray_SHAPE(np_ret)[0];
55 if (len0>0) {
56 int len = PyArray_SHAPE(np_ret)[1];
57 using wc500 = wchar_t[500];
58 auto* c_out = reinterpret_cast<wc500*>(PyArray_DATA(np_ret));
59
60 for (int i = 0; i < len; i++) {
61 std::wstring wa((c_out[i]));
62 std::string ret(wa.begin(), wa.end() );
63 out.push_back(ret);
64 }
65 }
66 Py_DECREF(pArgs);
67 if (np_ret) Py_DECREF(np_ret);
68 return out;
69}
70
71PyObject* ReaderuprootTree::get_function(PyObject* m_python_module, const std::string& name)
72{
73 if (!m_python_module) return nullptr;
74 PyObject* pFuncInitFile = PyObject_GetAttrString(m_python_module, name.c_str());
75 if (!pFuncInitFile || !PyCallable_Check(pFuncInitFile)) {
76 Py_XDECREF(pFuncInitFile);
77 std::cout << name<<"is null or not callable" << std::endl;
78 return nullptr;
79 }
80 return pFuncInitFile;
81}
82
83
84PyObject* ReaderuprootTree::init_python_module(const std::string& code)
85{
86 const char *SomeModuleName = "uproot4forhepmc3";
87 const char *SomeModuleCode = code.c_str();
88 PyObject *m_python_module = PyModule_New(SomeModuleName);
89 PyModule_AddStringConstant(m_python_module, "__file__", "");
90 PyObject *localDict = PyModule_GetDict(m_python_module);
91 PyObject *builtins = PyEval_GetBuiltins();
92 PyDict_SetItemString(localDict, "__builtins__", builtins);
93
94 PyObject *pyValue = PyRun_String(SomeModuleCode, Py_file_input, localDict, localDict);
95 if (pyValue == nullptr) {
96 return nullptr;
97 }
98 Py_DECREF(pyValue);
99 return m_python_module;
100}
101
102ReaderuprootTree::ReaderuprootTree(const std::string &filename,const std::string &treename,const std::string &branchname):
103 m_events_count(0),m_tree_name(treename), m_branch_name(branchname),m_tree(nullptr)
104{
105 if (!init(filename)) return;
106}
107
108bool ReaderuprootTree::init(const std::string &filename)
109{
110
112
114
115 set_run_info(std::make_shared<GenRunInfo>());
116
117 Py_Initialize();
118 import_array();
119
121
122 R"EOT(
123import uproot
124import numpy
125def init_file(filename):
126 rootfile=uproot.open(str(filename))
127# print(rootfile.keys())
128 return rootfile
129
130def close_file(filename):
131 return filename.close()
132
133def init_tree(rootfile,treename,branchname):
134 tree=rootfile[str(treename)+"/"+str(branchname)]
135# print(tree.keys())
136 return tree
137
138def init_genruninfo(rootfile,treename,branchname):
139 tree=rootfile[str(treename)+"/"+str(branchname)]
140# print(tree.keys())
141 return tree
142
143def get_number_of_entries_in_tree(tree):
144 return tree.num_entries
145
146def get_array_from_tree(tree,branch,i,destype):
147 result=tree[str(branch)].array(library="np")[i]
148 if len(destype.strip()) == 0:
149 output=numpy.array([result])
150 else:
151 output=numpy.array([result], dtype=destype)
152# print("a.shape={}, a.dtype={}".format(output.shape, output.dtype))
153# print(branch,output)
154 return output
155)EOT"
156);
157 bool result =false;
158 if (!m_python_module) {
159 HEPMC3_ERROR( "ReaderuprootTree: cannot initialize python modulr. Please check your uproot and/or numpy instalation.")
160 return result;
161 }
162 PyObject *pFuncInitFile = nullptr;
163 PyObject *pFuncInitTree = nullptr;
164 PyObject* pFuncEntries = nullptr;
165
166
167 PyObject * pArgsFile = nullptr;
168 PyObject * pArgsTree = nullptr;
169 PyObject * pArgsEntries = nullptr;
170 PyObject * pArgsGenRunInfo = nullptr;
171
172 m_access_function = get_function(m_python_module, "get_array_from_tree");
173 pFuncInitFile = get_function(m_python_module, "init_file");
174 pFuncInitTree = get_function(m_python_module, "init_tree");
175 pFuncEntries = get_function(m_python_module, "get_number_of_entries_in_tree");
176
177
178
179if (m_access_function && pFuncInitFile && pFuncInitTree && pFuncEntries) {
180 pArgsFile = PyTuple_New(1);
181 PyTuple_SetItem(pArgsFile, 0, Py_BuildValue("s#", filename.c_str(), filename.length()));
182 m_file=PyObject_CallObject(pFuncInitFile,pArgsFile);
183
184
185 if (m_file){
186 pArgsTree = PyTuple_New(3);
187 PyTuple_SetItem(pArgsTree, 0, m_file);
188 PyTuple_SetItem(pArgsTree, 1, Py_BuildValue("s#", m_tree_name.c_str(), m_tree_name.length()));
189 PyTuple_SetItem(pArgsTree, 2, Py_BuildValue("s#", m_branch_name.c_str(), m_branch_name.length()));
190 m_tree = PyObject_CallObject(pFuncInitTree, pArgsTree);
191
192 pArgsGenRunInfo = PyTuple_New(3);
193 PyTuple_SetItem(pArgsGenRunInfo, 0, m_file);
194 PyTuple_SetItem(pArgsGenRunInfo, 1, Py_BuildValue("s#", m_tree_name.c_str(), m_tree_name.length()));
195 PyTuple_SetItem(pArgsGenRunInfo, 2, Py_BuildValue("s#", "GenRunInfo", 10));
196 m_genruninfo = PyObject_CallObject(pFuncInitTree, pArgsGenRunInfo);
197 }
198
199 if (m_tree){
201 pArgsEntries = PyTuple_New(1);
202 PyTuple_SetItem(pArgsEntries, 0, m_tree);
203 PyObject * ret = PyObject_CallObject(pFuncEntries,pArgsEntries);
204 m_tree_getEntries = PyLong_AsLong(ret);
205 Py_DECREF(ret);
206 }
207 if (m_tree && m_file && m_genruninfo && m_tree_getEntries) result=true;
208}
209
210 Py_DECREF(pFuncInitFile);
211 Py_DECREF(pFuncInitTree);
212 Py_DECREF(pArgsEntries);
213
214 Py_DECREF(pFuncEntries);
215 Py_DECREF(pArgsFile);
216 return result;
217}
218
219bool ReaderuprootTree::skip(const int n)
220{
223}
224
225
226
227bool ReaderuprootTree::read_event(GenEvent& evt)
228{
229 if (!m_python_module) return false;
230 if (m_events_count >= m_tree_getEntries) { m_events_count++; return false;}
231 m_event_data->particles.clear();
232 m_event_data->vertices.clear();
233 m_event_data->links1.clear();
234 m_event_data->links2.clear();
235 m_event_data->attribute_id.clear();
238
239 auto event_number_v = get_vector<int>(m_tree, "event_number");
240 if (event_number_v.empty()) { m_events_count++; return false;}
241 auto weights = get_vector<double>(m_tree, "weights");
242 auto event_pos_1_v = get_vector<double>(m_tree, "event_pos/event_pos.m_v1");
243 if (event_pos_1_v.empty()) { m_events_count++; return false;}
244 auto event_pos_2_v = get_vector<double>(m_tree,"event_pos/event_pos.m_v2");
245 if (event_pos_2_v.empty()) { m_events_count++; return false;}
246 auto event_pos_3_v = get_vector<double>(m_tree, "event_pos/event_pos.m_v3");
247 if (event_pos_3_v.empty()) { m_events_count++; return false;}
248 auto event_pos_4_v = get_vector<double>(m_tree, "event_pos/event_pos.m_v4");
249 if (event_pos_4_v.empty()) { m_events_count++; return false;}
250 auto momentum_unit_v = get_vector<int>(m_tree, "momentum_unit");
251 if (momentum_unit_v.empty()) { m_events_count++; return false;}
252 auto length_unit_v = get_vector<int>(m_tree, "length_unit");
253 if (length_unit_v.empty()) { m_events_count++; return false;}
254
255 auto event_number = event_number_v.at(0);
256 auto event_pos_1 = event_pos_1_v.at(0);
257 auto event_pos_2 = event_pos_2_v.at(0);
258 auto event_pos_3 = event_pos_3_v.at(0);
259 auto event_pos_4 = event_pos_4_v.at(0);
260 auto momentum_unit = momentum_unit_v.at(0);
261 auto length_unit = length_unit_v.at(0);
262
263 auto links1 = get_vector<int>(m_tree, "links1");
264 auto links2 = get_vector<int>(m_tree, "links2");
265 auto attribute_id = get_vector<int>(m_tree, "attribute_id");
266 auto attribute_name = get_vector<std::string>(m_tree, "attribute_name");
267 auto attribute_string = get_vector<std::string>(m_tree, "attribute_string");
268 auto particlesmomentumm_v1 = get_vector<double>(m_tree, "particles/particles.momentum.m_v1");
269 auto particlesmomentumm_v2 = get_vector<double>(m_tree, "particles/particles.momentum.m_v2");
270 auto particlesmomentumm_v3 = get_vector<double>(m_tree, "particles/particles.momentum.m_v3");
271 auto particlesmomentumm_v4 = get_vector<double>(m_tree, "particles/particles.momentum.m_v4");
272 auto particlesmass = get_vector<double>(m_tree, "particles/particles.mass");
273 auto particlesis_mass_set = get_vector<bool>(m_tree, "particles/particles.is_mass_set");
274 auto particlesparticlespid = get_vector<int>(m_tree, "particles/particles.pid");
275 auto particlesparticlesstatus = get_vector<int>(m_tree, "particles/particles.status");
276
277 auto verticespositionm_v1 = get_vector<double>(m_tree, "vertices/vertices.position.m_v1");
278 auto verticespositionm_v2 = get_vector<double>(m_tree, "vertices/vertices.position.m_v2");
279 auto verticespositionm_v3 = get_vector<double>(m_tree, "vertices/vertices.position.m_v3");
280 auto verticespositionm_v4 = get_vector<double>(m_tree, "vertices/vertices.position.m_v4");
281 auto verticesverticesstatus = get_vector<int>(m_tree, "vertices/vertices.status");
282
283 m_event_data->event_number = event_number;
284 m_event_data->momentum_unit = momentum_unit == 0?HepMC3::Units::MEV:HepMC3::Units::GEV;
285 m_event_data->length_unit = length_unit == 0?HepMC3::Units::MM:HepMC3::Units::CM;
286 m_event_data->event_pos = HepMC3::FourVector(event_pos_1, event_pos_2, event_pos_3, event_pos_4) ;
287 m_event_data->links1 = links1;
288 m_event_data->links2 = links2;
289
290 for (size_t k=0; k < particlesparticlespid.size(); k++)
291 {
292 HepMC3::GenParticleData p = { particlesparticlespid[k], particlesparticlesstatus[k], particlesis_mass_set[k], particlesmass[k],
293 HepMC3::FourVector(particlesmomentumm_v1[k], particlesmomentumm_v1[k], particlesmomentumm_v1[k], particlesmomentumm_v1[k])
294 };
295 m_event_data->particles.push_back(p);
297
298 for (size_t k=0; k < verticesverticesstatus.size(); k++)
299 {
300 HepMC3::GenVertexData v = { verticesverticesstatus[k], HepMC3::FourVector(verticespositionm_v1[k], verticespositionm_v2[k], verticespositionm_v3[k], verticespositionm_v4[k])};
301 m_event_data->vertices.push_back(v);
302 }
303 m_event_data->weights=weights;
304 m_event_data->attribute_id=attribute_id;
305 m_event_data->attribute_name=attribute_name;
306 m_event_data->attribute_string=attribute_string;
307 evt.read_data(*m_event_data);
308
310 m_run_info_data->tool_name.clear();
315
322
323 run_info()->read_data(*m_run_info_data);
324 evt.set_run_info(run_info());
326 return true;
327}
328
330{
331 Py_DECREF(m_genruninfo);
332 Py_DECREF(m_tree);
333 Py_DECREF(m_file); //This should close the file, right?
334 Py_DECREF(m_access_function);
335 Py_DECREF(m_python_module);
336
337 m_file = nullptr;
338 m_tree = nullptr;
339 m_genruninfo = nullptr;
340 m_access_function = nullptr;
341 m_python_module = nullptr;
342 Py_DECREF( PyImport_ImportModule("threading")); //If someone, at some point would document it in CPython...
343 Py_Finalize();
344}
345
347{
348 if (!m_python_module) return true;
349 if (m_events_count >= m_tree_getEntries) return true;
350 return false;
351}
352ReaderuprootTree::~ReaderuprootTree()
353{
354 if (m_event_data) {delete m_event_data; m_event_data=nullptr;}
355 if (m_run_info_data) {delete m_run_info_data; m_run_info_data=nullptr;}
356}
357
358} // namespace HepMC3
359
360
#define HEPMC3_ERROR(MESSAGE)
Macro for printing error messages.
Definition: Errors.h:24
Generic 4-vector.
Definition: FourVector.h:36
virtual void set_run_info(std::shared_ptr< GenRunInfo > run)
Set the global GenRunInfo object.
Definition: Reader.h:56
virtual std::shared_ptr< GenRunInfo > run_info() const
Get the global GenRunInfo object.
Definition: Reader.h:44
ReaderuprootTree(const std::string &filename, const std::string &treename="hepmc3_tree", const std::string &branchname="hepmc3_event")
Constructor with tree and branch names.
PyObject * m_genruninfo
Python runInfo handler.
bool read_event(GenEvent &evt) override
Read event from file.
long int m_tree_getEntries
number of processed events
std::vector< T > get_vector(PyObject *file_name, const std::string &array_name, std::string desired_type="")
Get arrays.
PyObject * m_access_function
Python access function for arrays.
bool failed() override
Get file error state.
bool skip(const int) override
skip events
PyObject * m_tree
Python tree handler.
bool init(const std::string &filename)
init routine
GenEventData * m_event_data
Pointer to structure that holds event data.
int m_events_count
Events count. Needed to read the tree.
void close() override
Close file.
std::string m_tree_name
Name of TTree.
PyObject * get_function(PyObject *, const std::string &)
Get python functions.
std::string m_branch_name
Name of TBranch in TTree.
PyObject * init_python_module(const std::string &)
Init python module.
GenRunInfoData * m_run_info_data
Pointer to structure that holds run info data.
PyObject * m_python_module
Python module.
PyObject * m_file
Python file handler.
HepMC3 main namespace.
std::vector< std::string > ReaderuprootTree::get_vector< std::string >(PyObject *file_name, const std::string &array_name, std::string desired_type)
obtain vector of objects using name and type, specified for std::string
Stores serializable event information.
Definition: GenEventData.h:26
std::vector< GenVertexData > vertices
Vertices.
Definition: GenEventData.h:32
std::vector< int > links2
Second id of the vertex links.
Definition: GenEventData.h:52
int event_number
Event number.
Definition: GenEventData.h:27
std::vector< std::string > attribute_string
Attribute serialized as string.
Definition: GenEventData.h:56
std::vector< GenParticleData > particles
Particles.
Definition: GenEventData.h:31
std::vector< int > links1
First id of the vertex links.
Definition: GenEventData.h:51
std::vector< std::string > attribute_name
Attribute name.
Definition: GenEventData.h:55
Units::LengthUnit length_unit
Length unit.
Definition: GenEventData.h:29
std::vector< int > attribute_id
Attribute owner id.
Definition: GenEventData.h:54
FourVector event_pos
Event position.
Definition: GenEventData.h:35
std::vector< double > weights
Weights.
Definition: GenEventData.h:33
Units::MomentumUnit momentum_unit
Momentum unit.
Definition: GenEventData.h:28
Stores serializable particle information.
Stores serializable run information.
std::vector< std::string > tool_name
Tool names.
std::vector< std::string > tool_version
Tool versions.
std::vector< std::string > attribute_string
Attribute serialized as string.
std::vector< std::string > attribute_name
Attribute name.
std::vector< std::string > tool_description
Tool descriptions.
std::vector< std::string > weight_names
Weight names.
Stores serializable vertex information.
Definition: GenVertexData.h:22