diff --git a/setup.cfg b/setup.cfg index de57026..adbc550 100644 --- a/setup.cfg +++ b/setup.cfg @@ -28,6 +28,7 @@ python_requires = >=3.6 [options.extras_require] test = + numpy pytest pytest-cov graphviz diff --git a/src/core.cpp b/src/core.cpp index 0749104..94c6693 100644 --- a/src/core.cpp +++ b/src/core.cpp @@ -66,8 +66,8 @@ bool operator!=(const GenParticle& a, const GenParticle& b) { // compares all real qualities of both particle sets, // but ignores the .id() fields and the particle order -bool equal_particle_set(const std::vector& a, - const std::vector& b) { +bool equal_particle_sets(const std::vector& a, + const std::vector& b) { if (a.size() != b.size()) return false; auto unmatched = b; for (auto&& ai : a) { @@ -82,16 +82,16 @@ bool equal_particle_set(const std::vector& a, // compares all real qualities of two vertices, but ignores the .id() field bool operator==(const GenVertex& a, const GenVertex& b) { return a.status() == b.status() && is_close(a.position(), b.position()) && - equal_particle_set(a.particles_in(), b.particles_in()) && - equal_particle_set(a.particles_out(), b.particles_out()); + equal_particle_sets(a.particles_in(), b.particles_in()) && + equal_particle_sets(a.particles_out(), b.particles_out()); } bool operator!=(const GenVertex& a, const GenVertex& b) { return !operator==(a, b); } // compares all real qualities of both vertex sets, // but ignores the .id() fields and the vertex order -bool equal_vertex_set(const std::vector& a, - const std::vector& b) { +bool equal_vertex_sets(const std::vector& a, + const std::vector& b) { if (a.size() != b.size()) return false; auto unmatched = b; for (auto&& ai : a) { @@ -153,7 +153,7 @@ bool operator==(const GenEvent& a, const GenEvent& b) { } // if all vertices compare equal, then also all particles are equal - return equal_vertex_set(a.vertices(), b.vertices()); + return equal_vertex_sets(a.vertices(), b.vertices()); } bool operator!=(const GenEvent& a, const GenEvent& b) { return !operator==(a, b); } @@ -263,6 +263,13 @@ inline std::ostream& repr(std::ostream& os, const HepMC3::GenEvent& x) { repr(os, x.run_info()) << ")"; return os; } + +void from_hepevt(GenEvent& event, int event_number, py::array_t px, + py::array_t py, py::array_t pz, py::array_t en, + py::array_t m, py::array_t pid, py::array_t status, + py::object parents, py::object children, py::object vx, py::object vy, + py::object vz, py::object vt); + } // namespace HepMC3 PYBIND11_MODULE(_core, m) { @@ -279,7 +286,6 @@ PYBIND11_MODULE(_core, m) { // GenEvent // GenParticle // GenVertex - // fill_genevent_from_hepevt // print_hepevt // print_content // print_listing @@ -433,8 +439,6 @@ PYBIND11_MODULE(_core, m) { "run"_a, "momentum_unit"_a = Units::GEV, "length_unit"_a = Units::MM) .def(py::init(), "momentum_unit"_a = Units::GEV, "length_unit"_a = Units::MM) - PROP_RO_OL(particles, GenEvent, const std::vector&) - PROP_RO_OL(vertices, GenEvent, const std::vector&) .def_property("weights", overload_cast&, GenEvent>(&GenEvent::weights), [](GenEvent& self, py::sequence seq) { @@ -462,9 +466,6 @@ PYBIND11_MODULE(_core, m) { "name"_a, "value"_a) .def_property_readonly("weight_names", [](const GenEvent& self) { return self.weight_names(); }) - PROP(run_info, GenEvent) PROP(event_number, GenEvent) - PROP_RO(momentum_unit, GenEvent) PROP_RO(length_unit, GenEvent) - METH(set_units, GenEvent) .def_property("heavy_ion", overload_cast(&GenEvent::heavy_ion), &GenEvent::set_heavy_ion) @@ -474,15 +475,8 @@ PYBIND11_MODULE(_core, m) { .def_property( "cross_section", overload_cast(&GenEvent::cross_section), - &GenEvent::set_cross_section) METH(event_pos, GenEvent) - PROP_RO_OL(beams, GenEvent, std::vector) - METH_OL(add_vertex, GenEvent, void, GenVertexPtr) - METH_OL(add_particle, GenEvent, void, GenParticlePtr) - METH(set_beam_particles, GenEvent) - METH_OL(remove_vertex, GenEvent, void, GenVertexPtr) - METH_OL(remove_particle, GenEvent, void, GenParticlePtr) + &GenEvent::set_cross_section) .def("reserve", &GenEvent::reserve, "particles"_a, "vertices"_a = 0) - METH(clear, GenEvent) .def(py::self == py::self) .def("__repr__", [](GenEvent& self) { @@ -490,11 +484,35 @@ PYBIND11_MODULE(_core, m) { repr(os, self); return os.str(); }) - .def("__str__", [](GenEvent& self) { - std::ostringstream os; - HepMC3::Print::listing(os, self, 2); - return os.str(); - }); + .def("__str__", + [](GenEvent& self) { + std::ostringstream os; + HepMC3::Print::listing(os, self, 2); + return os.str(); + }) + + .def("from_hepevt", from_hepevt, "event_number"_a, "px"_a, "py"_a, "pz"_a, "en"_a, + "m"_a, "pid"_a, "status"_a, "parents"_a = py::none(), + "children"_a = py::none(), "vx"_a = py::none(), "vy"_a = py::none(), + "vz"_a = py::none(), "vt"_a = py::none()) + // clang-format off + METH(clear, GenEvent) + PROP(run_info, GenEvent) + PROP(event_number, GenEvent) + PROP_RO(momentum_unit, GenEvent) + PROP_RO(length_unit, GenEvent) + METH(set_units, GenEvent) + METH(event_pos, GenEvent) + PROP_RO_OL(beams, GenEvent, std::vector) + METH_OL(add_vertex, GenEvent, void, GenVertexPtr) + METH_OL(add_particle, GenEvent, void, GenParticlePtr) + METH(set_beam_particles, GenEvent) + METH_OL(remove_vertex, GenEvent, void, GenVertexPtr) + METH_OL(remove_particle, GenEvent, void, GenParticlePtr) + PROP_RO_OL(particles, GenEvent, const std::vector&) + PROP_RO_OL(vertices, GenEvent, const std::vector&) + // clang-format on + ; py::class_(m, "GenParticle") .def(py::init(), @@ -554,5 +572,8 @@ PYBIND11_MODULE(_core, m) { return os.str(); }); + FUNC(equal_particle_sets); + FUNC(equal_vertex_sets); + register_io(m); } diff --git a/src/from_hepevt.cpp b/src/from_hepevt.cpp new file mode 100644 index 0000000..9f04183 --- /dev/null +++ b/src/from_hepevt.cpp @@ -0,0 +1,182 @@ +#include "HepMC3/GenEvent.h" +#include "HepMC3/GenParticle.h" +#include "HepMC3/GenVertex.h" +#include "pybind.h" + +#include +#include +#include +#include +#include + +// template <> +// struct std::hash> { +// std::size_t operator()(const std::pair& p) const noexcept { +// auto h1 = std::hash{}(p.first); +// auto h2 = std::hash{}(p.second); +// return h1 ^ (h2 << 1); // or use boost::hash_combine +// } +// }; + +template <> +struct std::less> { + bool operator()(const std::pair& a, + const std::pair& b) const noexcept { + if (a.first == b.first) return a.second < b.second; + return a.first < b.first; + } +}; + +void normalize(int& m1, int& m2) { + // normalize mother range, see + // https://pythia.org/latest-manual/ParticleProperties.html + // m1 == m2 == 0 : no mothers + // m1 == m2 > 0 : elastic scattering + // m1 > 0, m2 == 0 : one mother + // m1 < m2, both > 0: interaction + // m2 < m1, both > 0: same, needs swapping + + if (m1 > 0 && m2 > 0 && m2 < m1) std::swap(m1, m2); + + if (m1 > 0 && m2 == 0) m2 = m1; + + --m1; // fortran to c index +} + +namespace HepMC3 { + +void connect_parents_and_children(GenEvent& event, bool parents, + py::array_t parents_or_children, py::object vx, + py::object vy, py::object vz, py::object vt) { + + if (parents_or_children.request().ndim != 2) + throw std::runtime_error("parents or children must be 2D"); + + auto rco = parents_or_children.unchecked<2>(); + + const std::vector& particles = event.particles(); + const int n = particles.size(); + if (rco.shape(0) != n || rco.shape(1) != 2) + throw std::runtime_error("parents or children must have shape (N, 2)"); + + // find unique vertices: + // particles with same parents or children share one vertex + // if parents: map from parents to children + // if children: map from children to parents + std::map, std::vector> vmap; + for (int i = 0; i < n; ++i) { + if (rco(i, 0) == 0 && rco(i, 1) == 0) continue; + vmap[std::make_pair(rco(i, 0), rco(i, 1))].push_back(i); + } + + const int has_vertex = !vx.is_none() + !vy.is_none() + !vz.is_none() + !vt.is_none(); + + if (has_vertex && has_vertex != 4) + throw std::runtime_error("if one of vx, vy, vz, vt is set, all must be set"); + + py::array_t avx; + py::array_t avy; + py::array_t avz; + py::array_t avt; + + if (has_vertex) { + avx = py::cast>(vx); + avy = py::cast>(vy); + avz = py::cast>(vz); + avt = py::cast>(vt); + + if (avx.request().ndim != 1 || avy.request().ndim != 1 || avz.request().ndim != 1 || + avt.request().ndim != 1) + throw std::runtime_error("vx, vy, vz, vt must be 1D"); + } + + auto x = avx.unchecked<1>(); + auto y = avy.unchecked<1>(); + auto z = avz.unchecked<1>(); + auto t = avt.unchecked<1>(); + + if (has_vertex && + (x.shape(0) != n || y.shape(0) != n || z.shape(0) != n || t.shape(0) != n)) + throw std::runtime_error("vx, vy, vz, vt must have same length"); + + for (const auto& vi : vmap) { + + int m1 = vi.first.first; + int m2 = vi.first.second; + // there must be at least one parent or child when we arrive here... + normalize(m1, m2); + assert(m1 < m2); + + // ...with at least one child or parent + const auto& co = vi.second; + assert(!co.empty()); + + FourVector pos; + if (has_vertex) { + // we assume this is a production vertex + // if parent, co.front() is location of production vertex of first child + // if child, we use location of first child m1 + const int i = parents ? co.front() : m1; + pos.set(x(i), y(i), z(i), t(i)); + } + + GenVertexPtr v{new GenVertex(pos)}; + + if (parents) { + for (int k = m1; k < m2; ++k) v->add_particle_in(particles.at(k)); + for (const auto k : co) v->add_particle_out(particles.at(k)); + } else { + for (int k = m1; k < m2; ++k) v->add_particle_out(particles.at(k)); + for (const auto k : co) v->add_particle_in(particles.at(k)); + } + + event.add_vertex(v); + } +} + +void from_hepevt(GenEvent& event, int event_number, py::array_t px, + py::array_t py, py::array_t pz, py::array_t en, + py::array_t m, py::array_t pid, py::array_t status, + py::object parents, py::object children, py::object vx, py::object vy, + py::object vz, py::object vt) { + if (px.request().ndim != 1) throw std::runtime_error("px must be 1D"); + if (py.request().ndim != 1) throw std::runtime_error("py must be 1D"); + if (pz.request().ndim != 1) throw std::runtime_error("pz must be 1D"); + if (en.request().ndim != 1) throw std::runtime_error("en must be 1D"); + if (m.request().ndim != 1) throw std::runtime_error("m must be 1D"); + if (pid.request().ndim != 1) throw std::runtime_error("pid must be 1D"); + if (status.request().ndim != 1) throw std::runtime_error("status must be 1D"); + + auto rpx = px.unchecked<1>(); + auto rpy = py.unchecked<1>(); + auto rpz = pz.unchecked<1>(); + auto ren = en.unchecked<1>(); + auto rm = m.unchecked<1>(); + auto rpid = pid.unchecked<1>(); + auto rsta = status.unchecked<1>(); + + const int n = pid.shape(0); + if (rpx.shape(0) != n || rpy.shape(0) != n || rpz.shape(0) != n || + ren.shape(0) != n || rm.shape(0) != n) + throw std::runtime_error("px, py, pz, en, m, pid, status must have same length"); + + event.clear(); + event.reserve(n); + event.set_event_number(event_number); + + for (int i = 0; i < n; ++i) { + GenParticlePtr p{ + new GenParticle(FourVector(rpx(i), rpy(i), rpz(i), ren(i)), rpid(i), rsta(i))}; + p->set_generated_mass(rm(i)); + event.add_particle(p); + } + + const bool have_parents = !parents.is_none(); + const bool have_children = !children.is_none(); + if (have_parents || have_children) + connect_parents_and_children( + event, have_parents, + py::cast>(have_parents ? parents : children), vx, vy, vz, vt); +} + +} // namespace HepMC3 diff --git a/src/pybind.h b/src/pybind.h index 275a3a6..bf0a21b 100644 --- a/src/pybind.h +++ b/src/pybind.h @@ -1,7 +1,7 @@ #ifndef PYHEPMC_PYBIND_HEADER #define PYHEPMC_PYBIND_HEADER -// #include +#include #include #include #include diff --git a/tests/test_basic.py b/tests/test_basic.py index 79177d9..befd7e6 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -4,38 +4,37 @@ @pytest.fixture() def evt(): - # - # In this example we will place the following event into HepMC "by hand" - # - # name status pdg_id parent Px Py Pz Energy Mass - # 1 !p+! 1 2212 0,0 0.000 0.000 7000.000 7000.000 0.938 - # 2 !p+! 2 2212 0,0 0.000 0.000-7000.000 7000.000 0.938 - # ========================================================================= - # 3 !d! 3 1 1,1 0.750 -1.569 32.191 32.238 0.000 - # 4 !u~! 4 -2 2,2 -3.047 -19.000 -54.629 57.920 0.000 - # 5 !W-! 5 -24 3,4 1.517 -20.68 -20.605 85.925 80.799 - # 6 !gamma! 6 22 3,4 -3.813 0.113 -1.833 4.233 0.000 - # 7 !d! 7 1 5,5 -2.445 28.816 6.082 29.552 0.010 - # 8 !u~! 8 -2 5,5 3.962 -49.498 -26.687 56.373 0.006 - - # now we build the graph, which will looks like - # p7 # - # p1 / # - # \v1__p3 p5---v4 # - # \_v3_/ \ # - # / \ p8 # - # v2__p4 \ # - # / p6 # - # p2 # - # # - - # - # Finally, the event gets a weight. - # + r""" + In this example we will generate the following event by hand + + name status pdg_id parent Px Py Pz Energy Mass + 1 !p+! 1 2212 0,0 0.000 0.000 7000.000 7000.000 0.938 + 2 !p+! 2 2212 0,0 0.000 0.000 -7000.000 7000.000 0.938 + ========================================================================= + 3 !d! 3 1 1,1 0.750 -1.569 32.191 32.238 0.000 + 4 !u~! 4 -2 2,2 -3.047 -19.000 -54.629 57.920 0.000 + 5 !W-! 5 -24 3,4 1.517 -20.68 -20.605 85.925 80.799 + 6 !gamma! 6 22 3,4 -3.813 0.113 -1.833 4.233 0.000 + 7 !d! 7 1 5,5 -2.445 28.816 6.082 29.552 0.010 + 8 !u~! 8 -2 5,5 3.962 -49.498 -26.687 56.373 0.006 + + The corresponding graph looks like this + + p7 + p1 / + \v1__p3 p5---v4 + \_v3_/ \ + / \ p8 + v2__p4 \ + / p6 + p2 + + Finally, the event gets a weight. + """ evt = hep.GenEvent(hep.Units.GEV, hep.Units.MM) evt.event_number = 1 - # px py pz e pdgid status + # px py pz e pdgid status p1 = hep.GenParticle((0.0, 0.0, 7000.0, 7000.0), 2212, 1) p1.generated_mass = 0.938 p2 = hep.GenParticle((0.0, 0.0, -7000.0, 7000.0), 2212, 2) @@ -110,6 +109,79 @@ def test_GenEvent(evt): assert p4.momentum == (-3.047, -19.0, -54.629, 57.920) +@pytest.mark.parametrize("use_parent", (True, False)) +def test_GenEvent_from_hepevt(use_parent, evt): + status = [p.status for p in evt.particles] + pid = [p.pid for p in evt.particles] + px = [p.momentum[0] for p in evt.particles] + py = [p.momentum[1] for p in evt.particles] + pz = [p.momentum[2] for p in evt.particles] + en = [p.momentum[3] for p in evt.particles] + m = [p.generated_mass for p in evt.particles] + vx = vy = vz = vt = [0, 0, 1, 2, 3, 3, 4, 4] + + assert len(m) == 8 + + # p7 + # p1 / + # \v1__p3 p5---v4 + # \_v3_/ \ + # / \ p8 + # v2__p4 \ + # / p6 + # p2 + + parents = [(0, 0), (0, 0), (1, 1), (2, 2), (3, 4), (3, 4), (5, 5), (5, 5)] + children = [(3, 0), (4, 0), (5, 6), (5, 6), (7, 8), (0, 0), (0, 0), (0, 0)] + + ev = hep.GenEvent() + if use_parent: + ev.from_hepevt( + evt.event_number, + px, + py, + pz, + en, + m, + pid, + status, + parents, + None, + vx, + vy, + vz, + vt, + ) + else: + ev.from_hepevt( + evt.event_number, + px, + py, + pz, + en, + m, + pid, + status, + None, + children, + vx, + vy, + vz, + vt, + ) + + # cannot be taken from HepEvt record, but is set for evt + ev.weights = [1.0] + ev.run_info = hep.GenRunInfo() + ev.run_info.weight_names = ["0"] + + assert len(ev.particles) == len(evt.particles) + assert ev.particles == evt.particles + assert len(ev.vertices) == len(evt.vertices) + assert hep.equal_vertex_sets(ev.vertices, evt.vertices) + assert ev == evt + + def test_sequence_access(): evt = hep.GenEvent() evt.add_particle(hep.GenParticle())