Skip to content

Commit 50c223a

Browse files
authored
switch to original hepm3 conversion code (#8)
1 parent e34aa9f commit 50c223a

File tree

6 files changed

+131
-172
lines changed

6 files changed

+131
-172
lines changed

setup.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,8 @@ def get_description():
132132
'extern/HepMC3/include',
133133
'extern/pybind11/include',
134134
],
135-
language='c++')
135+
define_macros=[("HEPMC3_HEPEVT_NMXHEP", 10000)], # increase this if necessary
136+
language='c++',)
136137
],
137138
zip_safe=False,
138139
)

src/bindings.cpp

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -633,26 +633,24 @@ PYBIND11_MODULE(_bindings, m) {
633633
HEPEVT_Wrapper::set_hepevt_address((char*)&self);
634634
HEPEVT_Wrapper::zero_everything();
635635
})
636-
// GenEvent_to_HEPEVT is broken: sometimes produces wrong connections
637-
// .def("fill_from_genevent", [](HEPEVT& self, const GenEvent& evt) {
638-
// HepMC::HEPEVT_Wrapper::set_hepevt_address((char*)&self);
639-
// HepMC::HEPEVT_Wrapper::GenEvent_to_HEPEVT(&evt);
640-
// })
636+
.def("from_genevent", [](HEPEVT& self, const GenEvent& evt) {
637+
HEPEVT_Wrapper::set_hepevt_address((char*)&self);
638+
HEPEVT_Wrapper::GenEvent_to_HEPEVT(&evt);
639+
})
640+
.def("to_genevent", [](const HEPEVT& self, GenEvent& evt) {
641+
HEPEVT_Wrapper::set_hepevt_address((char*)&self);
642+
HEPEVT_Wrapper::HEPEVT_to_GenEvent(&evt);
643+
})
641644
.def("__str__", [](const HEPEVT& self) {
642645
HEPEVT_Wrapper::set_hepevt_address((char*)&self);
643646
std::ostringstream os;
644647
HEPEVT_Wrapper::print_hepevt(os);
645648
return os.str();
646649
})
647-
.def_property_readonly("ptr", [](const HEPEVT& self) { return (std::intptr_t)&self; })
650+
.def_property_readonly("ptr", [](const HEPEVT& self) { return (std::uintptr_t)&self; })
648651
.def_property_readonly_static("max_size", [](py::object){ return NMXHEP; } )
649652
;
650653

651-
m.def("fill_genevent_from_hepevt", fill_genevent_from_hepevt<double>,
652-
"evt"_a, "event_number"_a, "p"_a, "m"_a, "v"_a, "pid"_a, "parents"_a,
653-
"children"_a, "particle_status"_a, "vertex_status"_a,
654-
"momentum_scaling"_a = 1.0, "length_scaling"_a = 1.0);
655-
656654
m.def("content", [](const GenEvent& event) {
657655
std::ostringstream os;
658656
HepMC3::Print::content(os, event);

src/hepevt_wrapper.h

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,29 @@
11
#ifndef PYHEPMC_NG_HEPEVT_WRAPPER_H
22
#define PYHEPMC_NG_HEPEVT_WRAPPER_H
33

4-
#include <pybind11/pybind11.h>
4+
#include "pybind.h"
55
#include <pybind11/numpy.h>
66
#include "HepMC3/FourVector.h"
77
#include "HepMC3/GenEvent.h"
88
#include "HepMC3/GenParticle.h"
99
#include "HepMC3/GenVertex.h"
10+
#include "HepMC3/HEPEVT_Wrapper.h"
1011
#include <map>
1112
#include <utility>
1213

13-
template <typename RealType>
14+
// alternative conversion from HEPEVT to GenEvent, UNUSED
1415
bool fill_genevent_from_hepevt(HepMC3::GenEvent& evt,
1516
int event_number,
16-
pybind11::array_t<RealType> p_array, // N x 4
17-
pybind11::array_t<RealType> m_array, // N
18-
pybind11::array_t<RealType> v_array, // N x 4
17+
pybind11::array_t<HEPMC3_HEPEVT_PRECISION> p_array, // N x 4
18+
pybind11::array_t<HEPMC3_HEPEVT_PRECISION> m_array, // N
19+
pybind11::array_t<HEPMC3_HEPEVT_PRECISION> v_array, // N x 4
1920
pybind11::array_t<int> pid_array, // N
2021
pybind11::array_t<int> parents_array, // N x 2
2122
pybind11::array_t<int> /* children_array */, // N x 2, not used
2223
pybind11::array_t<int> particle_status_array, // N
2324
pybind11::array_t<int> vertex_status_array, // N
24-
RealType momentum_scaling,
25-
RealType length_scaling) {
25+
HEPMC3_HEPEVT_PRECISION momentum_scaling,
26+
HEPMC3_HEPEVT_PRECISION length_scaling) {
2627
using namespace HepMC3;
2728
namespace py = pybind11;
2829

@@ -98,7 +99,7 @@ bool fill_genevent_from_hepevt(HepMC3::GenEvent& evt,
9899
v->add_particle_in(evt.particles()[j]);
99100
v->set_status(vsta(i));
100101
evt.add_vertex(v);
101-
vi = vertex_map.insert(std::make_pair(parents, v)).first;
102+
vi = vertex_map.emplace(parents, v).first;
102103
}
103104
vi->second->add_particle_out(evt.particles()[i]);
104105
}

src/pybind.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ _T overload_cast(_T x) {
1919

2020
#define FUNC(name) m.def(#name, name)
2121
#define PROP_RO(name, cls) .def_property_readonly(#name, &cls::name)
22+
#define PROP_ROS(name, cls) .def_property_readonly_static(#name, &cls::name)
2223
#define PROP_RO_OL(name, cls, rval) .def_property_readonly(#name, overload_cast<rval, const cls>(&cls::name))
2324
#define PROP(name, cls) .def_property(#name, &cls::name, &cls::set_##name)
2425
#define PROP_OL(name, cls, rval) .def_property(#name, overload_cast<rval, const cls> (&cls::name), &cls::set_##name)

src/pyhepmc_ng/__init__.py

Lines changed: 69 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -88,72 +88,72 @@ def open(filename, mode="r", precision=None):
8888

8989
elif mode == "w":
9090
return WrappedAsciiWriter(filename, precision)
91-
92-
93-
def fill_genevent_from_hepevent_ptr(evt, ptr_as_int, max_size,
94-
int_type=ctypes.c_int,
95-
float_type=ctypes.c_double,
96-
hep_status_decoder=None,
97-
momentum_unit = 1,
98-
length_unit = 1):
99-
# struct HEPEVT // Fortran common block HEPEVT
100-
# {
101-
# int nevhep; // Event number
102-
# int nhep; // Number of entries in the event
103-
# int isthep[NMXHEP]; // Status code
104-
# int idhep [NMXHEP]; // PDG ID
105-
# int jmohep[NMXHEP][2]; // Idx of first and last mother
106-
# int jdahep[NMXHEP][2]; // Idx of first and last daughter
107-
# momentum_t phep [NMXHEP][5]; // Momentum: px, py, pz, e, m
108-
# momentum_t vhep [NMXHEP][4]; // Vertex: x, y, z, t
109-
# };
110-
111-
IntArray = int_type * max_size
112-
Int2 = int_type * 2
113-
Int2Array = Int2 * max_size
114-
Float4 = float_type * 4
115-
Float5 = float_type * 5
116-
Float4Array = Float4 * max_size
117-
Float5Array = Float5 * max_size
118-
119-
class HEPEVT(ctypes.Structure):
120-
_fields_ = (
121-
("event_number", int_type),
122-
("nentries", int_type),
123-
("status", IntArray),
124-
("pid", IntArray),
125-
("parents", Int2Array),
126-
("children", Int2Array),
127-
("pm", Float5Array),
128-
("v", Float4Array)
129-
)
130-
131-
h = ctypes.cast(ptr_as_int, ctypes.POINTER(HEPEVT))[0]
132-
133-
event_number = h.event_number
134-
n = h.nentries
135-
136-
import numpy as np
137-
138-
status = np.asarray(h.status)[:n]
139-
pm = np.asarray(h.pm)
140-
141-
if hep_status_decoder is None:
142-
particle_status = status
143-
vertex_status = np.zeros_like(status)
144-
else:
145-
particle_status, vertex_status = hep_status_decoder(status)
146-
147-
fill_genevent_from_hepevt(evt,
148-
event_number,
149-
pm[:n,:4],
150-
pm[:n,4],
151-
h.v[:n],
152-
h.pid[:n],
153-
h.parents[:n],
154-
h.children[:n],
155-
particle_status,
156-
vertex_status,
157-
momentum_unit,
158-
length_unit)
159-
return evt
91+
raise ValueError("mode must be r or w")
92+
93+
94+
_hepevt_buffer = HEPEVT()
95+
96+
def fill_genevent_from_hepevt(evt, **kwargs):
97+
"""
98+
Fills GenEvent from HEPEVT data.
99+
100+
Parameters
101+
----------
102+
event_number : int
103+
Current Event Number. Starts with 1.
104+
p : array(float) N x 4
105+
Array of 4-vectors (px, py, py, e).
106+
m : array(float) N
107+
Array of generated masses (virtual masses).
108+
v : array(float) N x 4
109+
Array of 4-vectors (x, y, z, t).
110+
pid : array(int) N
111+
Array of PDG IDs.
112+
parents : array(int) N x 2
113+
Array of (parent first index, parent last index). Indices start at 1.
114+
children : array(int) N x 2 (optional)
115+
Array of (child first index, child last index). Indices start at 1. May be none.
116+
status : array(int) N
117+
Array of particle status. See HEPMC3 docs for definition.
118+
momentum_scaling : float (optional)
119+
Momentum coordinates are stored in units of this number.
120+
vertex_scaling : float (optional)
121+
Vertex coordinates are stored in units of this number.
122+
123+
Notes
124+
-----
125+
The current implementation copies the input into an internal buffer. This is not as efficient as possible, but currently the only way to use the HepMC3 C++ code with input data which does not have exactly the memory layout of the HEPEVT struct defined in HepMC3/HEPEVT_Wrapper.h.
126+
127+
Calling this function is not thread-safe.
128+
"""
129+
event_number = kwargs["event_number"]
130+
p = kwargs["p"]
131+
m = kwargs["m"]
132+
v = kwargs["v"]
133+
pid = kwargs["pid"]
134+
parents = kwargs["parents"]
135+
children = kwargs.get("children", 0)
136+
status = kwargs["status"]
137+
momentum_scaling = kwargs.get("momentum_scaling", 1.0)
138+
vertex_scaling = kwargs.get("vertex_scaling", 1.0)
139+
140+
global _hepevt_buffer
141+
n = pid.shape[0]
142+
if n > _hepevt_buffer.max_size:
143+
raise ValueError(
144+
('Number of particles in event (%i) exceeds HepMC3 buffer size (%i).\n'
145+
'Change the line `define_macros={"HEPMC3_HEPEVT_NMXHEP": 50000}` in setup.py\n'
146+
'to a larger value and (re)compile pyhepmc_ng from scratch.') %
147+
(n, _hepevt_buffer.max_size))
148+
_hepevt_buffer.event_number = event_number
149+
_hepevt_buffer.nentries = n
150+
_hepevt_buffer.pm()[:n, :4] = p / momentum_scaling
151+
_hepevt_buffer.pm()[:n, 4] = m / momentum_scaling
152+
_hepevt_buffer.v()[:n] = v / vertex_scaling
153+
_hepevt_buffer.pid()[:n] = pid
154+
_hepevt_buffer.parents()[:n] = parents
155+
_hepevt_buffer.children()[:n] = children
156+
_hepevt_buffer.status()[:n] = status
157+
158+
evt.clear()
159+
_hepevt_buffer.to_genevent(evt)

tests/test_basic.py

Lines changed: 41 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -97,66 +97,52 @@ def evt():
9797
return evt
9898

9999

100-
def prepare_hepevt(evt):
100+
def test_hepevt(evt):
101101
particles = evt.particles
102102
h = hep.HEPEVT()
103-
h.event_number = evt.event_number
104-
h.nentries = len(evt.particles)
105-
pid = h.pid()
106-
sta = h.status()
107-
par = h.parents()
108-
chi = h.children()
109-
pm = h.pm()
110-
v = h.v()
111-
for i,p in enumerate(particles):
112-
sta[i] = p.status
113-
pid[i] = p.pid
114-
iparent = [q.id for q in p.parents]
115-
par[i] = (min(iparent), max(iparent)) if iparent else (0, 0)
116-
children = p.children
117-
ichildren = [q.id for q in children]
118-
chi[i] = (min(ichildren), max(ichildren)) if ichildren else (0, 0)
119-
pm[i,:4] = p.momentum
120-
pm[i,4] = p.generated_mass
121-
v[i] = p.production_vertex.position if p.production_vertex else (0, 0, 0, 0)
122-
return h
103+
h.from_genevent(evt)
104+
evt2 = hep.GenEvent()
105+
h.to_genevent(evt2)
106+
107+
# hepevt has no run_info, so we add it articially
108+
evt2.run_info = evt.run_info
109+
assert evt == evt2
123110

124111

125-
def test_hepevt(evt):
126-
particles = evt.particles
112+
def test_fill_genevent_from_hepevt(evt):
113+
h = hep.HEPEVT()
114+
h.from_genevent(evt)
115+
evt2 = hep.GenEvent()
116+
hep.fill_genevent_from_hepevt(evt2,
117+
event_number=h.event_number,
118+
status=h.status(),
119+
pid=h.pid(),
120+
p=h.pm()[:,:4] * 2,
121+
m=h.pm()[:,4] * 2,
122+
v=h.v() * 5,
123+
parents=h.parents(),
124+
momentum_scaling=2,
125+
vertex_scaling=5)
126+
# hepevt has no run_info, so we add it articially
127+
evt2.run_info = evt.run_info
128+
assert evt == evt2
127129

128-
h = prepare_hepevt(evt)
129-
assert h.event_number == evt.event_number
130-
assert h.nentries == len(evt.particles)
131-
pid = h.pid()
132-
sta = h.status()
133-
par = h.parents()
134-
chi = h.children()
135-
pm = h.pm()
136-
v = h.v()
137-
for i,p in enumerate(particles):
138-
assert sta[i] == p.status
139-
assert pid[i] == p.pid
140-
iparent = [q.id for q in p.parents]
141-
if iparent:
142-
assert tuple(par[i]) == (min(iparent), max(iparent))
143-
ichildren = [q.id for q in p.children]
144-
if ichildren:
145-
assert tuple(chi[i]) == (min(ichildren), max(ichildren))
146-
assert np.allclose(pm[i,:4], p.momentum)
147-
if p.production_vertex:
148-
assert np.allclose(v[i], p.production_vertex.position)
149-
assert str(h) == """ Event No.: 1
150-
Nr Type Parent(s) Daughter(s) Px Py Pz E Inv. M.
151-
1 2212 0 - 0 3 - 3 0.00 0.00 7000.00 7000.00 0.94
152-
2 2212 0 - 0 4 - 4 0.00 0.00 -7000.00 7000.00 0.94
153-
3 1 1 - 1 5 - 6 0.75 -1.57 32.19 32.24 0.00
154-
4 -2 2 - 2 5 - 6 -3.05 -19.00 -54.63 57.92 0.00
155-
5 -24 3 - 4 7 - 8 1.52 -20.68 -20.61 85.92 80.80
156-
6 22 3 - 4 0 - 0 -3.81 0.11 -1.83 4.23 0.00
157-
7 1 5 - 5 0 - 0 -2.44 28.82 6.08 29.55 0.01
158-
8 -2 5 - 5 0 - 0 3.96 -49.50 -26.69 56.37 0.01
159-
"""
130+
with pytest.raises(ValueError, match="exceeds HepMC3 buffer size"):
131+
n = h.max_size + 1
132+
x = np.zeros(n, dtype=int)
133+
p = np.zeros((n, 2))
134+
v = np.zeros((n, 4), dtype=float)
135+
hep.fill_genevent_from_hepevt(evt2,
136+
event_number=h.event_number,
137+
status=x,
138+
pid=x,
139+
p=v,
140+
m=v[:, 0],
141+
v=v,
142+
parents=p)
143+
144+
with pytest.raises(KeyError):
145+
hep.fill_genevent_from_hepevt(evt2)
160146

161147

162148
def test_sequence_access():
@@ -190,31 +176,3 @@ def test_weights():
190176
evt.weight(2)
191177
with pytest.raises(RuntimeError, match="no weight with given name"):
192178
evt.set_weight("c", 4)
193-
194-
195-
def test_fill_genevent_from_hepevt(evt):
196-
h = prepare_hepevt(evt)
197-
evt2 = hep.GenEvent()
198-
hep.fill_genevent_from_hepevt(evt2,
199-
h.event_number,
200-
h.pm()[:,:4] * 2,
201-
h.pm()[:,4] * 2,
202-
h.v() * 5,
203-
h.pid(),
204-
h.parents(),
205-
h.children(),
206-
h.status(), # particle status
207-
np.zeros_like(h.pid()), # vertex status
208-
2, 5)
209-
# hepevt has no run_info, so we add it articially
210-
evt2.run_info = evt.run_info
211-
assert evt == evt2
212-
213-
214-
def test_fill_genevent_from_hepevt_ptr(evt):
215-
h = prepare_hepevt(evt)
216-
evt2 = hep.GenEvent()
217-
hep.fill_genevent_from_hepevent_ptr(evt2, h.ptr, h.max_size)
218-
# hepevt has no run_info, so we add it articially
219-
evt2.run_info = evt.run_info
220-
assert evt == evt2

0 commit comments

Comments
 (0)