1 #include "ReaderuprootTree.h"
5 HEPMC3_DECLARE_READER_FILE(ReaderuprootTree)
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);
21 if (np_ret) len0 = PyArray_SHAPE(np_ret)[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]);
28 if (np_ret) Py_DECREF(np_ret);
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()));
45 PyObject *pReturn = PyObject_CallObject(pFunc, pArgs);
46 PyArrayObject *np_ret =
reinterpret_cast<PyArrayObject*
>(pReturn);
47 std::vector<std::string> out;
49 if (np_ret) len0 = PyArray_SHAPE(np_ret)[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));
55 for (
int i = 0; i < len; i++) {
56 std::wstring wa((c_out[i]));
57 std::string ret(wa.begin(), wa.end() );
62 if (np_ret) Py_DECREF(np_ret);
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;
81 const char *SomeModuleName =
"uproot4forhepmc3";
82 const char *SomeModuleCode = code.c_str();
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);
89 PyObject *pyValue = PyRun_String(SomeModuleCode, Py_file_input, localDict, localDict);
90 if (pyValue ==
nullptr) {
101 m_events_count(0),m_tree_name(treename.c_str()), m_branch_name(branchname.c_str()),m_tree(nullptr)
103 if (!
init(filename))
return;
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"
124 HEPMC3_ERROR(
"ReaderuprootTree: cannot initialize python modulr. Please check your uproot and/or numpy instalation.")
127 PyObject *pFuncInitFile =
nullptr;
128 PyObject *pFuncInitTree =
nullptr;
129 PyObject* pFuncEntries =
nullptr;
132 PyObject * pArgsFile =
nullptr;
133 PyObject * pArgsTree =
nullptr;
134 PyObject * pArgsEntries =
nullptr;
135 PyObject * pArgsGenRunInfo =
nullptr;
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);
151 pArgsTree = PyTuple_New(3);
152 PyTuple_SetItem(pArgsTree, 0,
m_file);
155 m_tree = PyObject_CallObject(pFuncInitTree, pArgsTree);
157 pArgsGenRunInfo = PyTuple_New(3);
158 PyTuple_SetItem(pArgsGenRunInfo, 0,
m_file);
160 PyTuple_SetItem(pArgsGenRunInfo, 2, Py_BuildValue(
"s#",
"GenRunInfo", 10));
161 m_genruninfo = PyObject_CallObject(pFuncInitTree, pArgsGenRunInfo);
166 pArgsEntries = PyTuple_New(1);
167 PyTuple_SetItem(pArgsEntries, 0,
m_tree);
168 PyObject * ret = PyObject_CallObject(pFuncEntries,pArgsEntries);
175 Py_DECREF(pFuncInitFile);
176 Py_DECREF(pFuncInitTree);
177 Py_DECREF(pArgsEntries);
179 Py_DECREF(pFuncEntries);
180 Py_DECREF(pArgsFile);
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");
210 auto event_pos_2_v = get_vector<double>(
m_tree,
"event_pos/event_pos.m_v2");
212 auto event_pos_3_v = get_vector<double>(
m_tree,
"event_pos/event_pos.m_v3");
214 auto event_pos_4_v = get_vector<double>(
m_tree,
"event_pos/event_pos.m_v4");
216 auto momentum_unit_v = get_vector<int>(
m_tree,
"momentum_unit");
218 auto length_unit_v = get_vector<int>(
m_tree,
"length_unit");
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);
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");
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");
256 for (
size_t k=0; k < particlesparticlespid.size(); k++)
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])
264 for (
size_t k=0; k < verticesverticesstatus.size(); k++)
297 Py_DECREF(m_genruninfo);
305 m_genruninfo =
nullptr;
308 Py_DECREF( PyImport_ImportModule(
"threading"));
318 ReaderuprootTree::~ReaderuprootTree()
327 long int m_tree_getEntries
number of processed events
std::vector< int > attribute_id
Attribute owner id.
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.
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.
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.
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.
FourVector event_pos
Event position.
std::vector< std::string > attribute_string
Attribute serialized as string.
std::string m_branch_name
Name of TBranch in TTree.
Stores serializable event information.
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.
bool read_event(GenEvent &evt) override
Read event from file.
Stores serializable vertex information.
Stores serializable run information.
std::shared_ptr< GenRunInfo > run_info() const
Get the global GenRunInfo object.
Stores serializable particle information.
bool skip(const int) override
skip events
std::vector< GenParticleData > particles
Particles.
PyObject * m_file
Python file handler.
void close() override
Close file.
PyObject * m_genruninfo
Python runInfo handler.
std::vector< double > weights
Weights.
std::vector< GenVertexData > vertices
Vertices.
void set_run_info(std::shared_ptr< GenRunInfo > run)
Set the global GenRunInfo object.
#define HEPMC3_ERROR(MESSAGE)
Macro for printing error messages.
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.
PyObject * m_tree
Python tree handler.