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