Merge be24eda into 6051dba

CMakeLists.txt

@@ -1,12 +1,5 @@
 cmake_minimum_required(VERSION 3.13...3.17)
 
-if(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY)
-  message(
-    STATUS
-      "You should call this through setup.py so that all variables are set; python setup.py build_ext --inplace"
-  )
-endif()
-
 find_program(ccache_executable ccache)
 if(ccache_executable)
   message(STATUS "Using ccache: ${ccache_executable}")
@@ -28,10 +21,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 +30,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/bench.cpp

@@ -0,0 +1,18 @@
+#include <HepMC3/GenEvent.h>
+#include <HepMC3/GenParticle.h>
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+using namespace HepMC3;
+
+void register_bench(py::module& m) {
+  m.def("_sum_energy_of_protons", [](GenEvent& event) {
+    double esum = 0;
+    for (auto&& p : event.particles()) {
+      if (std::abs(p->pdg_id()) != 2212) continue;
+      if (p->status() != 1) continue;
+      esum += p->momentum().e();
+    }
+    return esum;
+  });
+}

src/core.cpp

@@ -1,10 +1,13 @@
 #include "HepMC3/AssociatedParticle.h"
 #include "HepMC3/GenPdfInfo_fwd.h"
 #include "attributes_view.hpp"
+#include "geneventdata.hpp"
+#include "numpy_api.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>
@@ -32,9 +35,11 @@
 // #include "GzReaderAscii.h"
 
 void register_io(py::module& m);
+void register_bench(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);
@@ -65,6 +70,8 @@ void from_hepevt(GenEvent& event, int event_number, py::array_t<double> px,
 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 +213,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,6 +523,9 @@ 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))
+      .def_property_readonly("numpy", [](py::object self) { return NumpyAPI(self); })
       // clang-format off
       EQ(GenEvent)
       REPR(GenEvent)
@@ -592,9 +621,6 @@ PYBIND11_MODULE(_core, m) {
       // clang-format on
       ;
 
-  // py::class_<GenParticleData>(m, "GenParticleData");
-  // py::class_<GenVertexData>(m, "GenVertexData");
-
   m.def(
       "content",
       [](const GenEvent& event) {
@@ -631,4 +657,6 @@ PYBIND11_MODULE(_core, m) {
   FUNC(equal_vertex_sets);
 
   register_io(m);
+  register_bench(m);
+  register_numpy_api(m);
 }

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/numpy_api.cpp

@@ -0,0 +1,59 @@
+#include "numpy_api.hpp"
+
+#define NP_ARRAY(cls, kind, type, method)                           \
+  .def_property_readonly(#method, [](cls& self) {                   \
+    using namespace HepMC3;                                         \
+    const GenEvent& event = py::cast<const GenEvent&>(self.event_); \
+    const auto& x = event.kind();                                   \
+    py::array_t<type> a(x.size());                                  \
+    auto a2 = a.mutable_unchecked<1>();                             \
+    for (size_t i = 0; i < x.size(); ++i) a2[i] = x[i]->method();   \
+    return a;                                                       \
+  })
+
+#define NP_ARRAY2(cls, kind, type, method1, method2)                         \
+  .def_property_readonly(#method2, [](cls& self) {                           \
+    using namespace HepMC3;                                                  \
+    const GenEvent& event = py::cast<const GenEvent&>(self.event_);          \
+    const auto& x = event.kind();                                            \
+    py::array_t<type> a(x.size());                                           \
+    auto a2 = a.mutable_unchecked<1>();                                      \
+    for (size_t i = 0; i < x.size(); ++i) a2[i] = x[i]->method1().method2(); \
+    return a;                                                                \
+  })
+
+void register_numpy_api(py::module& m) {
+  py::class_<ParticlesAPI>(m, "ParticlesAPI")
+      .def(py::init<py::object>())
+      // clang-format off
+        NP_ARRAY(ParticlesAPI, particles, int, id)
+        NP_ARRAY(ParticlesAPI, particles, int, pid)
+        NP_ARRAY(ParticlesAPI, particles, int, status)
+        NP_ARRAY(ParticlesAPI, particles, bool, is_generated_mass_set)
+        NP_ARRAY(ParticlesAPI, particles, double, generated_mass)
+        NP_ARRAY2(ParticlesAPI, particles, double, momentum, px)
+        NP_ARRAY2(ParticlesAPI, particles, double, momentum, py)
+        NP_ARRAY2(ParticlesAPI, particles, double, momentum, pz)
+        NP_ARRAY2(ParticlesAPI, particles, double, momentum, e)
+      // clang-format on
+      ;
+
+  py::class_<VerticesAPI>(m, "VerticesAPI")
+      .def(py::init<py::object>())
+      // clang-format off
+        NP_ARRAY(VerticesAPI, vertices, int, id)
+        NP_ARRAY(VerticesAPI, vertices, int, status)
+        NP_ARRAY(VerticesAPI, vertices, bool, has_set_position)
+        NP_ARRAY2(VerticesAPI, vertices, double, position, x)
+        NP_ARRAY2(VerticesAPI, vertices, double, position, y)
+        NP_ARRAY2(VerticesAPI, vertices, double, position, z)
+        NP_ARRAY2(VerticesAPI, vertices, double, position, t)
+      // clang-format on
+      ;
+
+  py::class_<NumpyAPI>(m, "NumpyAPI")
+      .def_property_readonly("particles",
+                             [](NumpyAPI& self) { return ParticlesAPI(self.event_); })
+      .def_property_readonly("vertices",
+                             [](NumpyAPI& self) { return VerticesAPI(self.event_); });
+}

src/numpy_api.hpp

@@ -0,0 +1,26 @@
+#ifndef PYHEPMC_NUMPY_API_HPP
+#define PYHEPMC_NUMPY_API_HPP
+
+#include <HepMC3/GenEvent.h>
+#include <HepMC3/GenParticle.h>
+#include <HepMC3/GenVertex.h>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+struct NumpyAPI {
+  py::object event_;
+  NumpyAPI(py::object event) : event_{event} {}
+};
+
+struct ParticlesAPI : NumpyAPI {
+  using NumpyAPI::NumpyAPI;
+};
+struct VerticesAPI : NumpyAPI {
+  using NumpyAPI::NumpyAPI;
+};
+
+void register_numpy_api(py::module& m);
+
+#endif

src/pointer.hpp

@@ -12,13 +12,17 @@ class HEPRUPAttribute;
 class GenRunInfo;
 class GenEvent;
 class AssociatedParticle;
+class GenParticle;
+class GenVertex;
 
 using GenEventPtr = std::shared_ptr<GenEvent>;
 using AttributePtr = std::shared_ptr<Attribute>;
 using HEPEUPAttributePtr = std::shared_ptr<HEPEUPAttribute>;
 using HEPRUPAttributePtr = std::shared_ptr<HEPRUPAttribute>;
 using GenRunInfoPtr = std::shared_ptr<GenRunInfo>;
 using AssociatedParticlePtr = std::shared_ptr<AssociatedParticle>;
+using ConstGenParticlePtr = std::shared_ptr<const GenParticle>;
+using ConstGenVertexPtr = std::shared_ptr<const GenVertex>;
 } // namespace HepMC3
 
 #endif