Merge 3c358f7 into 6051dba

CMakeLists.txt

@@ -28,10 +28,7 @@ find_program(
         g++-11
         g++-10
         g++-9
-        g++-8
-        g++-7
-        g++-6
-        g++-5)
+        g++-8)
 if(mold_executable AND gxx_executable)
   message(STATUS "Using mold: ${mold_executable}")
   add_compile_options(-fuse-ld=mold)
@@ -40,7 +37,7 @@ endif()
 
 project(pyhepmc LANGUAGES CXX)
 
-if(CMAKE_CXX_COMPILER_ID STREQUAL GNU)
+if(GCC)
   # HepMC3 uses reinterpret_cast instead of const_cast
   add_compile_options(-Wno-strict-aliasing)
 endif()

docs/examples.rst

@@ -5,3 +5,4 @@ Examples
     :maxdepth: 1
 
     examples/basics
+    examples/processing

src/core.cpp

@@ -1,10 +1,12 @@
 #include "HepMC3/AssociatedParticle.h"
 #include "HepMC3/GenPdfInfo_fwd.h"
 #include "attributes_view.hpp"
+#include "geneventdata.hpp"
 #include "pointer.hpp"
 #include "pybind.hpp"
 #include "repr.hpp"
 #include <HepMC3/Attribute.h>
+#include <HepMC3/Data/GenEventData.h>
 #include <HepMC3/FourVector.h>
 #include <HepMC3/GenCrossSection.h>
 #include <HepMC3/GenEvent.h>
@@ -35,6 +37,7 @@ void register_io(py::module& m);
 
 namespace HepMC3 {
 
+bool operator==(const GenEventData& a, const GenEventData& b);
 bool operator==(const GenParticle& a, const GenParticle& b);
 bool operator==(const GenVertex& a, const GenVertex& b);
 bool operator==(const GenRunInfo& a, const GenRunInfo& b);
@@ -60,11 +63,51 @@ void from_hepevt(GenEvent& event, int event_number, py::array_t<double> px,
                  py::object parents, py::object children, py::object vx, py::object vy,
                  py::object vz, py::object vt);
 
+#define NP_ARRAY(kind, type, method)                              \
+  py::array_t<type> np_##kind##_##method(GenEvent& self) {        \
+    const auto& p = self.kind();                                  \
+    py::array_t<type> a(p.size());                                \
+    auto a2 = a.mutable_unchecked<1>();                           \
+    for (size_t i = 0; i < p.size(); ++i) a2[i] = p[i]->method(); \
+    return a;                                                     \
+  }
+
+#define NP_ARRAY2(kind, type, method1, method2)                              \
+  py::array_t<type> np_##kind##_##method1##_##method2(GenEvent& self) {      \
+    const auto& p = self.kind();                                             \
+    py::array_t<type> a(p.size());                                           \
+    auto a2 = a.mutable_unchecked<1>();                                      \
+    for (size_t i = 0; i < p.size(); ++i) a2[i] = p[i]->method1().method2(); \
+    return a;                                                                \
+  }
+
+NP_ARRAY(particles, int, id)
+NP_ARRAY(particles, int, pdg_id)
+NP_ARRAY(particles, int, status)
+NP_ARRAY(particles, bool, is_generated_mass_set)
+NP_ARRAY(particles, double, generated_mass)
+
+NP_ARRAY2(particles, double, momentum, px)
+NP_ARRAY2(particles, double, momentum, py)
+NP_ARRAY2(particles, double, momentum, pz)
+NP_ARRAY2(particles, double, momentum, e)
+
+NP_ARRAY(vertices, int, id)
+NP_ARRAY(vertices, int, status)
+NP_ARRAY(vertices, bool, has_set_position)
+
+NP_ARRAY2(vertices, double, position, x)
+NP_ARRAY2(vertices, double, position, y)
+NP_ARRAY2(vertices, double, position, z)
+NP_ARRAY2(vertices, double, position, t)
+
 } // namespace HepMC3
 
 PYBIND11_MODULE(_core, m) {
   using namespace HepMC3;
 
+  register_geneventdata_dtypes();
+
   py::module_ m_doc = py::module_::import("pyhepmc._doc");
   auto doc = py::cast<std::map<std::string, std::string>>(m_doc.attr("doc"));
 
@@ -206,6 +249,25 @@ PYBIND11_MODULE(_core, m) {
       // clang-format on
       ;
 
+  py::class_<GenEventData>(m, "GenEventData", DOC(GenEventData))
+      .def(py::init<>())
+      // clang-format off
+      PROP_RO2(weights, GenEventData)
+      PROP_RO2(vertices, GenEventData)
+      PROP_RO2(particles, GenEventData)
+      PROP_RO2(links1, GenEventData)
+      PROP_RO2(links2, GenEventData)
+      PROP_RO2(attribute_id, GenEventData)
+      PROP_RO2(attribute_name, GenEventData)
+      PROP_RO2(attribute_string, GenEventData)
+      ATTR(event_number, GenEventData)
+      ATTR(momentum_unit, GenEventData)
+      ATTR(length_unit, GenEventData)
+      ATTR(event_pos, GenEventData)
+      EQ(GenEventData)
+      // clang-format on
+      ;
+
   py::class_<GenHeavyIon, GenHeavyIonPtr, Attribute>(m, "GenHeavyIon", DOC(GenHeavyIon))
       .def(py::init<>())
       // clang-format off
@@ -497,7 +559,25 @@ PYBIND11_MODULE(_core, m) {
            "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(), DOC(GenEvent.from_hepevt))
+      .def("write_data", &GenEvent::write_data, "data"_a, DOC(GenEvent.write_data))
+      .def("read_data", &GenEvent::read_data, "data"_a, DOC(GenEvent.read_data))
       // clang-format off
+      DEF(np_particles_generated_mass)
+      DEF(np_particles_id)
+      DEF(np_particles_is_generated_mass_set)
+      DEF(np_particles_momentum_e)
+      DEF(np_particles_momentum_px)
+      DEF(np_particles_momentum_py)
+      DEF(np_particles_momentum_pz)
+      DEF(np_particles_pdg_id)
+      DEF(np_particles_status)
+      DEF(np_vertices_has_set_position)
+      DEF(np_vertices_id)
+      DEF(np_vertices_position_t)
+      DEF(np_vertices_position_x)
+      DEF(np_vertices_position_y)
+      DEF(np_vertices_position_z)
+      DEF(np_vertices_status)
       EQ(GenEvent)
       REPR(GenEvent)
       METH(clear, GenEvent)
@@ -592,9 +672,6 @@ PYBIND11_MODULE(_core, m) {
       // clang-format on
       ;
 
-  // py::class_<GenParticleData>(m, "GenParticleData");
-  // py::class_<GenVertexData>(m, "GenVertexData");
-
   m.def(
       "content",
       [](const GenEvent& event) {

src/equal.cpp

@@ -1,6 +1,9 @@
 #include "HepMC3/GenPdfInfo.h"
 #include "HepMC3/LHEFAttributes.h"
 #include "pointer.hpp"
+#include <HepMC3/Data/GenEventData.h>
+#include <HepMC3/Data/GenParticleData.h>
+#include <HepMC3/Data/GenVertexData.h>
 #include <HepMC3/GenEvent.h>
 #include <HepMC3/GenHeavyIon.h>
 #include <HepMC3/GenParticle.h>
@@ -14,6 +17,11 @@ bool operator!=(const T& a, const T& b) {
   return !operator==(a, b);
 }
 
+template <class T>
+bool operator==(const std::vector<T>& a, const std::vector<T>& b) {
+  return std::equal(a.begin(), a.end(), b.begin(), b.end());
+}
+
 // equality comparions used by unit tests
 bool is_close(const FourVector& a, const FourVector& b, double rel_eps = 1e-7) {
   auto is_close = [rel_eps](double a, double b) { return std::abs(a - b) < rel_eps; };
@@ -73,10 +81,8 @@ bool operator==(const GenRunInfo& a, const GenRunInfo& b) {
   const auto b_attr = b.attributes();
   return a.tools().size() == b.tools().size() &&
          a.weight_names().size() == b.weight_names().size() &&
-         a_attr.size() == b_attr.size() &&
-         std::equal(a.tools().begin(), a.tools().end(), b.tools().begin()) &&
-         std::equal(a.weight_names().begin(), a.weight_names().end(),
-                    b.weight_names().begin()) &&
+         a_attr.size() == b_attr.size() && a.tools() == b.tools() &&
+         a.weight_names() == b.weight_names() &&
          std::equal(a_attr.begin(), a_attr.end(), b_attr.begin(),
                     [](const std::pair<std::string, AttributePtr>& a,
                        const std::pair<std::string, AttributePtr>& b) {
@@ -146,4 +152,22 @@ bool operator==(const GenEvent& a, const GenEvent& b) {
   return equal_vertex_sets(a.vertices(), b.vertices());
 }
 
+bool operator==(const GenParticleData& a, const GenParticleData& b) {
+  return a.pid == b.pid && a.status == b.status && a.is_mass_set == b.is_mass_set &&
+         a.mass == b.mass && a.momentum == b.momentum;
+}
+
+bool operator==(const GenVertexData& a, const GenVertexData& b) {
+  return a.status == b.status && a.position == b.position;
+}
+
+bool operator==(const GenEventData& a, const GenEventData& b) {
+  return a.event_number == b.event_number && a.momentum_unit == b.momentum_unit &&
+         a.length_unit == b.length_unit && a.particles == b.particles &&
+         a.vertices == b.vertices && a.weights == b.weights &&
+         a.event_pos == b.event_pos && a.links1 == b.links1 && a.links2 == b.links2 &&
+         a.attribute_id == b.attribute_id && a.attribute_name == b.attribute_name &&
+         a.attribute_string == b.attribute_string;
+}
+
 } // namespace HepMC3

src/geneventdata.cpp

@@ -0,0 +1,50 @@
+#include "geneventdata.hpp"
+#include <HepMC3/Data/GenEventData.h>
+#include <HepMC3/Data/GenParticleData.h>
+#include <HepMC3/Data/GenVertexData.h>
+
+using namespace HepMC3;
+
+struct ParticleData {
+  int pid;
+  int status;
+  bool is_mass_set;
+  double mass;
+  double px, py, pz, e;
+};
+
+struct VertexData {
+  int status;
+  double x, y, z, t;
+};
+
+void register_geneventdata_dtypes() {
+  PYBIND11_NUMPY_DTYPE(ParticleData, pid, status, is_mass_set, mass, px, py, pz, e);
+  PYBIND11_NUMPY_DTYPE(VertexData, status, x, y, z, t);
+}
+
+#define MAKE(name, type)                                              \
+  py::object GenEventData_##name(py::object self) {                   \
+    auto& s = py::cast<GenEventData&>(self);                          \
+    py::ssize_t shape[1] = {static_cast<py::ssize_t>(s.name.size())}; \
+    auto dt = py::dtype::of<type>();                                  \
+    return py::array(dt, shape, s.name.data(), self);                 \
+  }
+
+MAKE(weights, double)
+MAKE(vertices, VertexData)
+MAKE(particles, ParticleData)
+MAKE(links1, int)
+MAKE(links2, int)
+
+#define MAKES(name)                                 \
+  py::object GenEventData_##name(py::object self) { \
+    auto& s = py::cast<GenEventData&>(self);        \
+    py::list l;                                     \
+    for (const auto& str : s.name) l.append(str);   \
+    return l;                                       \
+  }
+
+MAKES(attribute_id)
+MAKES(attribute_name)
+MAKES(attribute_string)

src/geneventdata.hpp

@@ -0,0 +1,17 @@
+#ifndef PYHEPMC_GENEVENTDATA_HPP
+#define PYHEPMC_GENEVENTDATA_HPP
+
+#include "pybind.hpp"
+
+py::object GenEventData_particles(py::object);
+py::object GenEventData_vertices(py::object);
+py::object GenEventData_weights(py::object);
+py::object GenEventData_links1(py::object);
+py::object GenEventData_links2(py::object);
+py::object GenEventData_attribute_id(py::object);
+py::object GenEventData_attribute_name(py::object);
+py::object GenEventData_attribute_string(py::object);
+
+void register_geneventdata_dtypes();
+
+#endif

src/pybind.hpp

@@ -5,6 +5,7 @@
 #include <pybind11/operators.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
 
 namespace py = pybind11;
 using namespace py::literals;
@@ -21,6 +22,7 @@ _T overload_cast(_T x) {
 
 #define FUNC(name) m.def(#name, name, DOC(name))
 #define PROP_RO(name, cls) .def_property_readonly(#name, &cls::name, DOC(cls.name))
+#define PROP_RO2(name, cls) .def_property_readonly(#name, cls##_##name, DOC(cls.name))
 #define PROP_ROS(name, cls) \
   .def_property_readonly_static(#name, &cls::name, DOC(cls.name))
 #define PROP_RO_OL(name, cls, rval)                                         \
@@ -40,5 +42,6 @@ _T overload_cast(_T x) {
 #define REPR(name) .def("__repr__", repr<name>)
 #define EQ(name) \
   .def("__eq__", py::overload_cast<const name&, const name&>(HepMC3::operator==))
+#define DEF(name) .def(#name, name)
 
 #endif

src/pyhepmc/__init__.py

@@ -30,13 +30,16 @@
 Missing functionality
 ---------------------
 
-- Not yet implemented: ``GenParticleData``, ``GenVertexData``, ``ReaderMT``.
-  These will be added in the future.
+- ``ReaderMT`` will not be implemented. If you want to use multi-threaded IO,
+  it is better to use the high-level threading API of Python to achieve this.
+  You can read the next GenEvent in a child thread from a file while you are
+  processing the current event in the main thread.
 
 """
 from pyhepmc._core import (
     Units,
     FourVector,
+    GenEventData,
     GenEvent,
     GenParticle,
     GenVertex,
@@ -69,6 +72,9 @@
 __all__ = (
     "Units",
     "FourVector",
+    "GenEventData",
+    "GenParticleData",
+    "GenVertexData",
     "GenEvent",
     "GenParticle",
     "GenVertex",