From 3fb41da85705ebd264052da424a2fbe3b33e8721 Mon Sep 17 00:00:00 2001 From: Hans Dembinski Date: Thu, 1 Sep 2022 13:34:18 +0200 Subject: [PATCH] Attribute support (#34) The missing attribute support is added for GenRunInfo, GenEvent, GenParticle, GenVertex. The following derived classes are added or completed: - GenPdfInfo, GenHeavyIon, GenCrossSection, HEPRUPAttribute, HEPEUPAttribute C++ Attributes are converted to native Python types were possible: - all C++ primitives (numbers, bool) and std::string are mapped to corresponding python types - std::vector of C++ primitives and string are mapped to list of corresponding python types --- .github/workflows/docs.yml | 47 +++++ .github/workflows/sphinx.yml | 31 --- .gitmodules | 6 + CMakeLists.txt | 4 +- docs/citation.rst | 8 + docs/conf.py | 10 +- docs/index.rst | 1 + extern/HepMC3 | 2 +- extern/cpp-member-accessor | 1 + extern/mp11 | 1 + setup.cfg | 1 + src/attribute_conversion.cpp | 117 +++++++++++ src/attributes_view.cpp | 91 +++++++++ src/attributes_view.hpp | 68 +++++++ src/core.cpp | 344 ++++++++++++++++++++++---------- src/from_hepevt.cpp | 9 +- src/io.cpp | 21 +- src/pointer.hpp | 22 ++ src/{pybind.h => pybind.hpp} | 4 +- src/pyhepmc/__init__.py | 11 +- src/pyhepmc/_attributes.py | 48 +++++ src/pyhepmc/_doc.py | 4 + src/pyhepmc/io.py | 3 + src/runinfo_attributes_view.cpp | 65 ++++++ tests/test_basic.py | 154 ++++++++++++++ tests/test_init.py | 2 + 26 files changed, 910 insertions(+), 165 deletions(-) create mode 100644 .github/workflows/docs.yml delete mode 100644 .github/workflows/sphinx.yml create mode 100644 docs/citation.rst create mode 160000 extern/cpp-member-accessor create mode 160000 extern/mp11 create mode 100644 src/attribute_conversion.cpp create mode 100644 src/attributes_view.cpp create mode 100644 src/attributes_view.hpp create mode 100644 src/pointer.hpp rename src/{pybind.h => pybind.hpp} (96%) create mode 100644 src/pyhepmc/_attributes.py create mode 100644 src/runinfo_attributes_view.cpp diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..1ff70a1 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,47 @@ +name: docs + +on: + pull_request: + release: + types: [published] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref }} + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + with: + submodules: true + fetch-depth: 3 + - uses: actions/setup-python@v2 + with: + python-version: "3.9" + - run: python -m pip install -v -e .[doc] + - run: python docs/build.py + - uses: actions/upload-pages-artifact@v1 + with: + path: 'docs/_build/html' + + deploy: + if: github.ref_type == 'tag' || github.ref_name == 'main' + needs: build + # Set permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages + permissions: + contents: read + pages: write + id-token: write + + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + + runs-on: ubuntu-latest + + steps: + - uses: actions/configure-pages@v2 + - uses: actions/deploy-pages@v1 diff --git a/.github/workflows/sphinx.yml b/.github/workflows/sphinx.yml deleted file mode 100644 index 2328ce7..0000000 --- a/.github/workflows/sphinx.yml +++ /dev/null @@ -1,31 +0,0 @@ -name: doc - -on: - pull_request: - push: - branches: [main] - release: - types: [published] - -concurrency: - group: ${{ github.workflow }}-${{ github.head_ref }} - cancel-in-progress: true - -jobs: - doc: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - with: - submodules: true - fetch-depth: 3 - - uses: actions/setup-python@v2 - with: - python-version: "3.9" - - run: python -m pip install -v -e .[doc] - - run: python docs/build.py - - uses: peaceiris/actions-gh-pages@v3 - if: github.ref_type == 'tag' || github.ref_name == 'main' - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: docs/_build/html diff --git a/.gitmodules b/.gitmodules index ed6dfc3..368b228 100644 --- a/.gitmodules +++ b/.gitmodules @@ -4,3 +4,9 @@ [submodule "extern/HepMC3"] path = extern/HepMC3 url = https://gitlab.cern.ch/hepmc/HepMC3.git +[submodule "extern/mp11"] + path = extern/mp11 + url = https://github.com/boostorg/mp11.git +[submodule "extern/cpp-member-accessor"] + path = extern/cpp-member-accessor + url = https://github.com/hliberacki/cpp-member-accessor.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 6ac8454..e46ea6e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,7 +10,7 @@ if(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY) endif() set(CMAKE_CXX_STANDARD - 11 + 14 CACHE STRING "C++ version selection") set(CMAKE_CXX_STANDARD_REQUIRED ON) # optional, ensure standard is supported set(CMAKE_CXX_EXTENSIONS OFF) # optional, keep compiler extensions off @@ -27,6 +27,8 @@ file(GLOB SOURCES "src/*.cpp") file(GLOB SOURCES2 "extern/HepMC3/src/*.cc") pybind11_add_module(_core MODULE ${SOURCES} ${SOURCES2}) +target_include_directories(_core PRIVATE extern/cpp-member-accessor/include) +target_include_directories(_core PRIVATE extern/mp11/include) target_include_directories(_core PRIVATE extern/HepMC3/include) target_compile_definitions(_core PRIVATE HepMC3_EXPORTS=1) diff --git a/docs/citation.rst b/docs/citation.rst new file mode 100644 index 0000000..312a9ba --- /dev/null +++ b/docs/citation.rst @@ -0,0 +1,8 @@ +Citation +======== + +If you use pyhepmc, please cite the HepMC3 library on which pyhepmc is build, and pyhepmc itself on Zenodo: + +* `A. Buckley and others, Comput.Phys.Commun. 260 (2021) 107310 `_ + +* `pyhepmc on Zenodo `_ diff --git a/docs/conf.py b/docs/conf.py index 2526466..619d522 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -6,6 +6,7 @@ # -- Project information ----------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information import pyhepmc +import os project = "pyhepmc" copyright = "2022, Hans Dembinski" @@ -28,9 +29,16 @@ # -- Options for HTML output ------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output -html_theme = "classic" html_static_path = ["_static"] +on_rtd = os.environ.get("READTHEDOCS", None) == "True" +if not on_rtd: + # Import and set the theme if we're building docs locally + import sphinx_rtd_theme + + html_theme = "sphinx_rtd_theme" + html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] + # Autodoc options autodoc_member_order = "groupwise" autodoc_mock_imports = ["numpy", "particle", "graphviz"] diff --git a/docs/index.rst b/docs/index.rst index 42c0683..87ce5f8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -4,3 +4,4 @@ :hidden: reference + citation diff --git a/extern/HepMC3 b/extern/HepMC3 index c524bfd..591bccc 160000 --- a/extern/HepMC3 +++ b/extern/HepMC3 @@ -1 +1 @@ -Subproject commit c524bfde132d1778b16fa2967388ebfa0e1b83bc +Subproject commit 591bccc2fdb65bfed9ab97a85eda57c0c1073355 diff --git a/extern/cpp-member-accessor b/extern/cpp-member-accessor new file mode 160000 index 0000000..e72109f --- /dev/null +++ b/extern/cpp-member-accessor @@ -0,0 +1 @@ +Subproject commit e72109f1400a9b70cdaaae1bf0e9192900dcc332 diff --git a/extern/mp11 b/extern/mp11 new file mode 160000 index 0000000..f6133a9 --- /dev/null +++ b/extern/mp11 @@ -0,0 +1 @@ +Subproject commit f6133a9f1f965d89676a33c4a39b3df09373b929 diff --git a/setup.cfg b/setup.cfg index f31afef..9bd07da 100644 --- a/setup.cfg +++ b/setup.cfg @@ -34,6 +34,7 @@ test = particle doc = sphinx + sphinx-rtd-theme [flake8] max-line-length = 90 diff --git a/src/attribute_conversion.cpp b/src/attribute_conversion.cpp new file mode 100644 index 0000000..d7f3f33 --- /dev/null +++ b/src/attribute_conversion.cpp @@ -0,0 +1,117 @@ +#include "pointer.hpp" +#include "pybind.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace HepMC3 { + +py::object attribute_to_python(AttributePtr a) { + using namespace boost::mp11; + + // this implementation must cover all C++ attribute types derived from Attribute + + py::object result; + + using RawTypes = mp_list; + // use mp_identity to make sure that default ctor is a noop + using Types = mp_transform; + // this loop has exactly one match if any + mp_for_each([&](auto t) { + using AttributeType = typename decltype(t)::type; + if (auto x = std::dynamic_pointer_cast(a)) result = py::cast(x); + }); + + if (!result) { + + using RawTypes = + mp_list; + + using Types = mp_transform; + // this loop has exactly one match if any + mp_for_each([&](auto t) { + using AttributeType = typename decltype(t)::type; + if (auto x = std::dynamic_pointer_cast(a)) + result = py::cast(x->value()); + }); + } + + if (!result) throw std::runtime_error("Attribute not convertible to Python type"); + return result; +} + +AttributePtr attribute_from_python(py::object obj) { + using namespace boost::mp11; + + AttributePtr result; + + if (py::isinstance(obj)) { + result = py::cast(obj); + } else if (py::isinstance(obj)) { + result = py::cast(obj); + } else if (py::isinstance(obj)) { + result = py::cast(obj); + } else if (py::isinstance(obj)) { + result = py::cast(obj); + } else if (py::isinstance(obj)) { + result = py::cast(obj); + } + + if (!result) { + using Types = mp_list, + mp_list, + mp_list, + mp_list>; + + mp_for_each([&](auto t) { + using T = decltype(t); + using AT = mp_at_c; + using PT = mp_at_c; + using CT = mp_at_c; + if (!result && py::isinstance(obj)) { + auto a = std::make_shared(); + a->set_value(py::cast(obj)); + result = a; + } + }); + } + + if (!result && py::isinstance(obj) && py::len(obj) > 0) { + using Types = mp_list, + mp_list, + mp_list>; + + py::int_ zero(0); + mp_for_each([&](auto t) { + using T = decltype(t); + using AT = mp_at_c; + using PT = mp_at_c; + using CT = std::vector>; + auto item = py::reinterpret_borrow(obj[zero]); + if (!result && py::isinstance(item)) { + auto a = std::make_shared(); + a->set_value(py::cast(obj)); + result = a; + } + }); + } + + if (!result) throw std::runtime_error("Python type not convertible to Attribute"); + return result; +} + +} // namespace HepMC3 diff --git a/src/attributes_view.cpp b/src/attributes_view.cpp new file mode 100644 index 0000000..86d9779 --- /dev/null +++ b/src/attributes_view.cpp @@ -0,0 +1,91 @@ +#include "attributes_view.hpp" +#include "pointer.hpp" +#include "pybind.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// To avoid the superfluous copy, we use the legal crowbar +// to access the private attribute map of GenEvent +MEMBER_ACCESSOR(MA1, HepMC3::GenEvent, m_attributes, + HepMC3::AttributesView::AttributeMap) + +namespace HepMC3 { +py::object AttributesView::Iter::next() { + while (it_ != end_) { + auto& amap2 = it_->second; + if (amap2.find(id_) != amap2.end()) return py::cast(it_++->first); + ++it_; + } + throw py::stop_iteration(); +} + +AttributesView::Iter AttributesView::iter() { + auto& amap = attributes(); + return {amap.begin(), amap.end(), id_}; +} + +py::object AttributesView::getitem(py::str name) { + auto& amap = attributes(); + auto it = amap.find(py::cast(name)); + if (it != amap.end()) { + auto& amap2 = it->second; + auto jt = amap2.find(id_); + if (jt != amap2.end()) return attribute_to_python(jt->second); + } + throw py::key_error(name); + return {}; +} + +void AttributesView::setitem(py::str name, py::object value) { + auto& amap = attributes(); + auto& amap2 = amap[py::cast(name)]; + amap2.emplace(id_, attribute_from_python(value)); +} + +void AttributesView::delitem(py::str name) { + auto& amap = attributes(); + auto it = amap.find(py::cast(name)); + if (it != amap.end()) { + auto& amap2 = it->second; + auto jt = amap2.find(id_); + if (jt != amap2.end()) { + amap2.erase(jt); + return; + } + } + throw py::key_error(name); +} + +bool AttributesView::contains(py::str name) { + auto& amap = attributes(); + auto it = amap.find(py::cast(name)); + if (it == amap.end()) return false; + auto& amap2 = it->second; + return amap2.find(id_) != amap2.end(); +} + +AttributesView::AttributeMap& AttributesView::attributes() { + auto ref = accessor::accessMember(*event_); + return ref.get(); +} + +py::ssize_t AttributesView::len() { + py::size_t n = 0; + auto& amap = attributes(); + for (auto& kv : amap) { + auto& amap2 = kv.second; + auto it = amap2.find(id_); + if (it != amap2.end()) ++n; + } + return n; +} + +} // namespace HepMC3 diff --git a/src/attributes_view.hpp b/src/attributes_view.hpp new file mode 100644 index 0000000..5c36671 --- /dev/null +++ b/src/attributes_view.hpp @@ -0,0 +1,68 @@ +#ifndef PYHEPMC_ATTRIBUTEMAPVIEW_HPP +#define PYHEPMC_ATTRIBUTEMAPVIEW_HPP + +#include "HepMC3/Attribute.h" +#include "pointer.hpp" +#include "pybind.hpp" +#include +#include +#include + +namespace HepMC3 { + +class GenRunInfo; +class GenEvent; +class Attribute; + +struct AttributesView { + using AttributeIdMap = std::map; + using AttributeMap = std::map; + + GenEvent* event_; + int id_; + + struct Iter { + using iterator = AttributeMap::iterator; + iterator it_, end_; + int id_; + py::object next(); + }; + + Iter iter(); + py::object getitem(py::str name); + void setitem(py::str name, py::object value); + void delitem(py::str name); + bool contains(py::str name); + py::ssize_t len(); + AttributeMap& attributes(); +}; + +struct RunInfoAttributesView { + using AttributeMap = std::map; + using iterator = typename AttributeMap::iterator; + + GenRunInfoPtr run_info_; + + struct Iter { + using iterator = AttributeMap::iterator; + iterator it_, end_; + py::object next(); + }; + + Iter iter(); + py::object getitem(py::str name); + void setitem(py::str name, py::object value); + void delitem(py::str name); + bool contains(py::str name); + py::ssize_t len(); + AttributeMap& attributes(); + iterator find(py::str name); + iterator end(); +}; + +py::object attribute_to_python(AttributePtr a); +AttributePtr attribute_from_python(py::object obj); + +} // namespace HepMC3 + +#endif diff --git a/src/core.cpp b/src/core.cpp index 4c9291f..57b1a7f 100644 --- a/src/core.cpp +++ b/src/core.cpp @@ -1,25 +1,29 @@ -#include "pybind.h" - -#include "HepMC3/Attribute.h" -#include "HepMC3/FourVector.h" -#include "HepMC3/GenCrossSection.h" -#include "HepMC3/GenEvent.h" -#include "HepMC3/GenHeavyIon.h" -#include "HepMC3/GenParticle.h" -#include "HepMC3/GenPdfInfo.h" -#include "HepMC3/GenRunInfo.h" -#include "HepMC3/GenVertex.h" -#include "HepMC3/Print.h" -#include "HepMC3/Units.h" - +#include "attributes_view.hpp" +#include "pointer.hpp" +#include "pybind.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include #include #include #include #include +#include +#include #include #include #include +#include #include // #include "GzReaderAscii.h" @@ -46,7 +50,10 @@ std::ostream& ostream_range(std::ostream& os, Iterator begin, Iterator end, namespace HepMC3 { -using GenRunInfoPtr = std::shared_ptr; +template +bool operator!=(const T& a, const T& b) { + return !operator==(a, b); +} // equality comparions used by unit tests bool is_close(const FourVector& a, const FourVector& b, double rel_eps = 1e-7) { @@ -61,10 +68,6 @@ bool operator==(const GenParticle& a, const GenParticle& b) { is_close(a.momentum(), b.momentum()); } -bool operator!=(const GenParticle& a, const GenParticle& b) { - return !operator==(a, b); -} - // compares all real qualities of both particle sets, // but ignores the .id() fields and the particle order bool equal_particle_sets(const std::vector& a, @@ -87,8 +90,6 @@ bool operator==(const GenVertex& a, const GenVertex& b) { 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_sets(const std::vector& a, @@ -108,10 +109,6 @@ bool operator==(const GenRunInfo::ToolInfo& a, const GenRunInfo::ToolInfo& b) { return a.name == b.name && a.version == b.version && a.description == b.description; } -bool operator!=(const GenRunInfo::ToolInfo& a, const GenRunInfo::ToolInfo& b) { - return !operator==(a, b); -} - bool operator==(const GenRunInfo& a, const GenRunInfo& b) { const auto a_attr = a.attributes(); const auto b_attr = b.attributes(); @@ -122,8 +119,8 @@ bool operator==(const GenRunInfo& a, const GenRunInfo& b) { std::equal(a.weight_names().begin(), a.weight_names().end(), b.weight_names().begin()) && std::equal(a_attr.begin(), a_attr.end(), b_attr.begin(), - [](const std::pair>& a, - const std::pair>& b) { + [](const std::pair& a, + const std::pair& b) { if (a.first != b.first) return false; if (bool(a.second) != bool(b.second)) return false; if (!a.second) return true; @@ -134,8 +131,6 @@ bool operator==(const GenRunInfo& a, const GenRunInfo& b) { }); } -bool operator!=(const GenRunInfo& a, const GenRunInfo& b) { return !operator==(a, b); } - bool operator==(const GenEvent& a, const GenEvent& b) { // incomplete: // missing comparison of GenHeavyIon, GenPdfInfo, GenCrossSection @@ -157,7 +152,17 @@ bool operator==(const GenEvent& a, const GenEvent& b) { return equal_vertex_sets(a.vertices(), b.vertices()); } -bool operator!=(const GenEvent& a, const GenEvent& b) { return !operator==(a, b); } +bool operator==(const GenHeavyIon& a, const GenHeavyIon& b) { + return a.Ncoll_hard == b.Ncoll_hard && a.Npart_proj == b.Npart_proj && + a.Npart_targ == b.Npart_targ && a.Ncoll == b.Ncoll && + a.N_Nwounded_collisions == b.N_Nwounded_collisions && + a.Nwounded_N_collisions == b.Nwounded_N_collisions && + a.Nwounded_Nwounded_collisions == b.Nwounded_Nwounded_collisions && + a.impact_parameter == b.impact_parameter && + a.event_plane_angle == b.event_plane_angle && + a.sigma_inel_NN == b.sigma_inel_NN && a.centrality == b.centrality && + a.user_cent_estimate == b.user_cent_estimate; +} template std::ostream& repr(std::ostream& os, const std::shared_ptr& x) { @@ -255,6 +260,24 @@ inline std::ostream& repr(std::ostream& os, const HepMC3::GenEvent& x) { return os; } +inline int gencrosssection_validate_index(GenCrossSection& cs, py::object obj) { + auto idx = py::cast(obj); + const auto size = cs.event() ? std::max(cs.event()->weights().size(), 1ul) : 1ul; + if (idx < 0) idx += size; + if (idx < 0 || idx >= size) throw py::index_error(); + return idx; +} + +inline std::string gencrosssection_validate_name(GenCrossSection& cs, py::object obj) { + auto name = py::cast(obj); + if (cs.event() && cs.event()->run_info()) { + const auto& wnames = cs.event()->run_info()->weight_names(); + if (std::find(wnames.begin(), wnames.end(), name) != wnames.end()) return name; + } + throw py::key_error(name); + return {}; +} + 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, @@ -374,54 +397,6 @@ PYBIND11_MODULE(_core, m) { FUNC(delta_r2_rap); FUNC(delta_r_rap); - py::class_ clsGenRunInfo(m, "GenRunInfo", DOC(GenRunInfo)); - clsGenRunInfo.def(py::init<>()) - .def_property("tools", - overload_cast&, GenRunInfo>( - &GenRunInfo::tools), - [](GenRunInfo& self, py::sequence seq) { - self.tools() = py::cast>(seq); - }) - .def("__repr__", - [](const GenRunInfo& self) { - std::ostringstream os; - repr(os, self); - return os.str(); - }) - .def(py::self == py::self) - // clang-format off - PROP(weight_names, GenRunInfo) - PROP_RO(attributes, GenRunInfo) - // clang-format on - ; - - py::class_(clsGenRunInfo, "ToolInfo", DOC(GenRunInfo.ToolInfo)) - .def(py::init(), "name"_a, "version"_a, - "description"_a) - .def(py::init([](py::sequence seq) { - if (py::len(seq) != 3) throw py::value_error("length != 3"); - return new GenRunInfo::ToolInfo({py::cast(seq[0]), - py::cast(seq[1]), - py::cast(seq[2])}); - })) - .def("__repr__", - [](const GenRunInfo::ToolInfo& self) { - std::ostringstream os; - repr(os, self); - return os.str(); - }) - .def(py::self == py::self) - // clang-format off - ATTR(name, GenRunInfo::ToolInfo) - ATTR(version, GenRunInfo::ToolInfo) - ATTR(description, GenRunInfo::ToolInfo) - // clang-format on - ; - - py::implicitly_convertible(); - - using AttributePtr = std::shared_ptr; - py::class_(m, "Attribute", DOC(Attribute)) .def_property_readonly( "particle", overload_cast(&Attribute::particle), @@ -443,24 +418,28 @@ PYBIND11_MODULE(_core, m) { // clang-format on ; - // py::class_(m, "IntAttribute", DOC(IntAttribute)) - // .def(py::init<>()) - // .def(py::init(), "val"_a) - // .def("__str__", - // [](const IntAttribute& self) { - // std::string s; - // self.to_string(s); - // return s; - // }) - // // clang-format off - // METH(from_string, IntAttribute) - // PROP(value, IntAttribute) - // // clang-format on - // ; - - py::class_(m, "GenHeavyIon", DOC(GenHeavyIon)) + py::class_(m, "HEPRUPAttribute", + DOC(HEPRUPAttribute)) + .def(py::init<>()) + // clang-format off + ATTR(heprup, HEPRUPAttribute) + // clang-format on + ; + + py::class_(m, "HEPEUPAttribute", + DOC(HEPEUPAttribute)) .def(py::init<>()) // clang-format off + METH(momentum, HEPEUPAttribute) + ATTR(hepeup, HEPEUPAttribute) + // clang-format on + ; + + py::class_(m, "GenHeavyIon", DOC(GenHeavyIon)) + .def(py::init<>()) + .def("__eq__", py::overload_cast( + HepMC3::operator==)) + // clang-format off ATTR(Ncoll_hard, GenHeavyIon) ATTR(Npart_proj, GenHeavyIon) ATTR(Npart_targ, GenHeavyIon) @@ -482,7 +461,7 @@ PYBIND11_MODULE(_core, m) { // clang-format on ; - py::class_(m, "GenPdfInfo", DOC(GenPdfInfo)) + py::class_(m, "GenPdfInfo", DOC(GenPdfInfo)) .def(py::init<>()) .def_property( "parton_id1", [](const GenPdfInfo& self) { return self.parton_id[0]; }, @@ -519,16 +498,25 @@ PYBIND11_MODULE(_core, m) { // clang-format on ; - py::class_(m, "GenCrossSection", + py::class_(m, "GenCrossSection", DOC(GenCrossSection)) - .def(py::init<>()) + .def(py::init([](double cs, double cse, long acc, long att) { + auto p = std::make_shared(); + p->set_cross_section(cs, cse, acc, att); + return p; + }), + "cross_section"_a, "cross_section_error"_a, "accepted_events"_a, + "attempted_events"_a) .def( "xsec", [](GenCrossSection& self, py::object obj) { + // need to do checks for invalid indices because they are missing in C++ if (py::isinstance(obj)) { - return self.xsec(py::cast(obj)); + auto idx = gencrosssection_validate_index(self, obj); + return self.xsec(idx); } else if (py::isinstance(obj)) { - return self.xsec(py::cast(obj)); + auto name = gencrosssection_validate_name(self, obj); + return self.xsec(name); } else throw py::type_error("int or str required"); return 0.0; @@ -538,9 +526,11 @@ PYBIND11_MODULE(_core, m) { "xsec_err", [](GenCrossSection& self, py::object obj) { if (py::isinstance(obj)) { - return self.xsec_err(py::cast(obj)); + auto idx = gencrosssection_validate_index(self, obj); + return self.xsec_err(idx); } else if (py::isinstance(obj)) { - return self.xsec_err(py::cast(obj)); + auto name = gencrosssection_validate_name(self, obj); + return self.xsec_err(name); } else throw py::type_error("int or str required"); return 0.0; @@ -550,9 +540,11 @@ PYBIND11_MODULE(_core, m) { "set_xsec", [](GenCrossSection& self, py::object obj, double value) { if (py::isinstance(obj)) { - self.set_xsec(py::cast(obj), value); + auto idx = gencrosssection_validate_index(self, obj); + self.set_xsec(idx, value); } else if (py::isinstance(obj)) { - self.set_xsec(py::cast(obj), value); + auto name = gencrosssection_validate_name(self, obj); + self.set_xsec(name, value); } else throw py::type_error("int or str required"); }, @@ -561,9 +553,11 @@ PYBIND11_MODULE(_core, m) { "set_xsec_err", [](GenCrossSection& self, py::object obj, double value) { if (py::isinstance(obj)) { - self.set_xsec_err(py::cast(obj), value); + auto idx = gencrosssection_validate_index(self, obj); + self.set_xsec_err(idx, value); } else if (py::isinstance(obj)) { - self.set_xsec_err(py::cast(obj), value); + auto name = gencrosssection_validate_name(self, obj); + self.set_xsec_err(name, value); } else throw py::type_error("int or str required"); }, @@ -575,10 +569,98 @@ PYBIND11_MODULE(_core, m) { // clang-format on ; - py::class_(m, "GenEvent", DOC(GenEvent)) - .def(py::init, Units::MomentumUnit, - Units::LengthUnit>(), - "run"_a, "momentum_unit"_a = Units::GEV, "length_unit"_a = Units::MM) + py::class_ clsAttributesView(m, "AttributesView", + DOC(AttributesView)); + clsAttributesView + + .def("__getitem__", &AttributesView::getitem) + .def("__setitem__", &AttributesView::setitem) + .def("__delitem__", &AttributesView::delitem) + .def("__contains__", &AttributesView::contains) + .def("__iter__", &AttributesView::iter) + .def("__len__", &AttributesView::len); + + py::class_(clsAttributesView, "Iter") + .def("next", &AttributesView::Iter::next) + .def("__next__", &AttributesView::Iter::next) + .def("__iter__", [](AttributesView::Iter& self) { return self; }); + + py::class_ clsRunInfoAttributesView( + m, "RunInfoAttributesView", DOC(RunInfoAttributesView)); + clsRunInfoAttributesView + + .def("__getitem__", &RunInfoAttributesView::getitem) + .def("__setitem__", &RunInfoAttributesView::setitem) + .def("__delitem__", &RunInfoAttributesView::delitem) + .def("__contains__", &RunInfoAttributesView::contains) + .def("__iter__", &RunInfoAttributesView::iter) + .def("__len__", &RunInfoAttributesView::len); + + py::class_(clsRunInfoAttributesView, "Iter") + .def("next", &RunInfoAttributesView::Iter::next) + .def("__next__", &RunInfoAttributesView::Iter::next) + .def("__iter__", [](RunInfoAttributesView::Iter& self) { return self; }); + + py::class_ clsGenRunInfo(m, "GenRunInfo", DOC(GenRunInfo)); + + clsGenRunInfo.def(py::init<>()) + .def_property("tools", + overload_cast&, GenRunInfo>( + &GenRunInfo::tools), + [](GenRunInfo& self, py::sequence seq) { + self.tools() = py::cast>(seq); + }) + .def_property( + "attributes", [](GenRunInfoPtr self) { return RunInfoAttributesView{self}; }, + [](GenRunInfoPtr self, py::dict obj) { + auto amv = RunInfoAttributesView{self}; + py::cast(amv).attr("clear")(); + for (const auto& kv : obj) { + amv.setitem(py::cast(kv.first), + py::reinterpret_borrow(kv.second)); + } + }, + DOC(attributes)) + .def("__repr__", + [](const GenRunInfo& self) { + std::ostringstream os; + repr(os, self); + return os.str(); + }) + .def(py::self == py::self) + // clang-format off + PROP(weight_names, GenRunInfo) + // clang-format on + ; + + py::class_(clsGenRunInfo, "ToolInfo", DOC(GenRunInfo.ToolInfo)) + .def(py::init(), "name"_a, "version"_a, + "description"_a) + .def(py::init([](py::sequence seq) { + if (py::len(seq) != 3) throw py::value_error("length != 3"); + return new GenRunInfo::ToolInfo({py::cast(seq[0]), + py::cast(seq[1]), + py::cast(seq[2])}); + })) + .def("__repr__", + [](const GenRunInfo::ToolInfo& self) { + std::ostringstream os; + repr(os, self); + return os.str(); + }) + .def(py::self == py::self) + // clang-format off + ATTR(name, GenRunInfo::ToolInfo) + ATTR(version, GenRunInfo::ToolInfo) + ATTR(description, GenRunInfo::ToolInfo) + // clang-format on + ; + + py::implicitly_convertible(); + + py::class_(m, "GenEvent", DOC(GenEvent)) + .def(py::init(), "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) .def_property( @@ -608,7 +690,7 @@ PYBIND11_MODULE(_core, m) { [](GenEvent& self, const std::string& name, double v) { self.weight(name) = v; }, - "name"_a, "value"_a, DOC(GenEvent.weight)) + "name"_a, "value"_a, DOC(GenEvent.set_weight)) .def_property("heavy_ion", overload_cast(&GenEvent::heavy_ion), &GenEvent::set_heavy_ion, DOC(GenEvent.heavy_ion)) @@ -619,6 +701,20 @@ PYBIND11_MODULE(_core, m) { "cross_section", overload_cast(&GenEvent::cross_section), &GenEvent::set_cross_section, DOC(GenEvent.cross_section)) + .def_property( + "attributes", + [](GenEvent& self) { + return AttributesView{&self, 0}; + }, + [](GenEvent& self, py::dict obj) { + auto amv = AttributesView{&self, 0}; + py::cast(amv).attr("clear")(); + for (const auto& kv : obj) { + amv.setitem(py::cast(kv.first), + py::reinterpret_borrow(kv.second)); + } + }, + DOC(attributes)) .def("reserve", &GenEvent::reserve, "particles"_a, "vertices"_a = 0, DOC(GenEvent.reserve)) .def(py::self == py::self) @@ -668,6 +764,20 @@ PYBIND11_MODULE(_core, m) { repr(os, self); return os.str(); }) + .def_property( + "attributes", + [](GenParticle& self) { + return AttributesView{self.parent_event(), self.id()}; + }, + [](GenParticle& self, py::dict obj) { + auto amv = AttributesView{self.parent_event(), self.id()}; + py::cast(amv).attr("clear")(); + for (const auto& kv : obj) { + amv.setitem(py::cast(kv.first), + py::reinterpret_borrow(kv.second)); + } + }, + DOC(attributes)) // clang-format off PROP_RO_OL(parent_event, GenParticle, const GenEvent*) PROP_RO(in_event, GenParticle) @@ -696,6 +806,20 @@ PYBIND11_MODULE(_core, m) { repr(os, self); return os.str(); }) + .def_property( + "attributes", + [](GenVertex& self) { + return AttributesView{self.parent_event(), self.id()}; + }, + [](GenVertex& self, py::dict obj) { + auto amv = AttributesView{self.parent_event(), self.id()}; + py::cast(amv).attr("clear")(); + for (const auto& kv : obj) { + amv.setitem(py::cast(kv.first), + py::reinterpret_borrow(kv.second)); + } + }, + DOC(GenVertex.attributes)) // clang-format off PROP_RO_OL(parent_event, GenVertex, const GenEvent*) PROP_RO(in_event, GenVertex) diff --git a/src/from_hepevt.cpp b/src/from_hepevt.cpp index 9f04183..08b8884 100644 --- a/src/from_hepevt.cpp +++ b/src/from_hepevt.cpp @@ -1,8 +1,7 @@ -#include "HepMC3/GenEvent.h" -#include "HepMC3/GenParticle.h" -#include "HepMC3/GenVertex.h" -#include "pybind.h" - +#include "pybind.hpp" +#include +#include +#include #include #include #include diff --git a/src/io.cpp b/src/io.cpp index c05044a..0e8ec67 100644 --- a/src/io.cpp +++ b/src/io.cpp @@ -1,19 +1,16 @@ -#include "pybind.h" - -#include "HepMC3/GenRunInfo.h" -#include "HepMC3/ReaderAscii.h" -#include "HepMC3/ReaderAsciiHepMC2.h" -#include "HepMC3/ReaderHEPEVT.h" -#include "HepMC3/ReaderLHEF.h" -#include "HepMC3/WriterAscii.h" -#include "HepMC3/WriterAsciiHepMC2.h" -#include "HepMC3/WriterHEPEVT.h" - +#include "pybind.hpp" +#include +#include +#include +#include +#include +#include +#include +#include #ifdef HEPMC3_ROOTIO #include "HepMC3/ReaderRoot.h" #include "HepMC3/ReaderRootTree.h" #endif - #include #include #include diff --git a/src/pointer.hpp b/src/pointer.hpp new file mode 100644 index 0000000..439b641 --- /dev/null +++ b/src/pointer.hpp @@ -0,0 +1,22 @@ +#ifndef PYHEPMC_POINTER_HPP +#define PYHEPMC_POINTER_HPP + +#include "HepMC3/GenRunInfo.h" +#include + +namespace HepMC3 { +class Attribute; +class HEPEUPAttribute; +class HEPRUPAttribute; +class GenRunInfo; +class GenEvent; + +using GenEventPtr = std::shared_ptr; +using AttributePtr = std::shared_ptr; +using HEPEUPAttributePtr = std::shared_ptr; +using HEPRUPAttributePtr = std::shared_ptr; +using GenRunInfoPtr = std::shared_ptr; +using AttributePtr = std::shared_ptr; +} // namespace HepMC3 + +#endif diff --git a/src/pybind.h b/src/pybind.hpp similarity index 96% rename from src/pybind.h rename to src/pybind.hpp index 0a69c49..1aef48b 100644 --- a/src/pybind.h +++ b/src/pybind.hpp @@ -1,5 +1,5 @@ -#ifndef PYHEPMC_PYBIND_HEADER -#define PYHEPMC_PYBIND_HEADER +#ifndef PYHEPMC_PYBIND_HPP +#define PYHEPMC_PYBIND_HPP #include #include diff --git a/src/pyhepmc/__init__.py b/src/pyhepmc/__init__.py index 939185a..a833d0f 100644 --- a/src/pyhepmc/__init__.py +++ b/src/pyhepmc/__init__.py @@ -10,6 +10,10 @@ convertible to :class:`FourVector` und :class:`ToolInfo`. - In addition to the C++ Reader/Writer classes, we offer an easy to use :func:`open`. It can read any supported format and writes in HepMC3 format. +- Attributes for :class:`GenRunInfo`, :class:`GenEvent`, :class:`GenParticle`, + :class:`GenVertex` can be accessed via a dict-like view returned by the + ``attributes`` property. The view automatically converts between native C++ + types to native Python types. - The ``Print`` class is missing, but :func:`listing` and :func:`content` are present as free functions. - The member functions ``delta_X`` of :class:`FourVector` are free functions @@ -25,8 +29,6 @@ - Not yet implemented: ``GenParticleData``, ``GenVertexData``, ``ReaderMT``, ``ReaderGZ``, ``Setup``, ``WriterGZ``. These will be added in the future. -- Generic ``Attribute`` s for :class:`GenEvent`, :class:`GenParticle`, - :class:`GenVertex`, :class:`GenRunInfo` are not yet implemented. """ # flake8: F401 from ._core import ( # noqa: F401 @@ -39,6 +41,8 @@ GenRunInfo, GenPdfInfo, GenCrossSection, + HEPRUPAttribute, + HEPEUPAttribute, equal_vertex_sets, equal_particle_sets, content, @@ -54,6 +58,7 @@ from .io import open as open # noqa: F401 from ._version import __version__ as __version__ # noqa: F401 import typing as _tp +from . import _attributes # noqa __all__ = ( "Units", @@ -65,6 +70,8 @@ "GenRunInfo", "GenPdfInfo", "GenCrossSection", + "HEPRUPAttribute", + "HEPEUPAttribute", "equal_vertex_sets", "equal_particle_sets", "content", diff --git a/src/pyhepmc/_attributes.py b/src/pyhepmc/_attributes.py new file mode 100644 index 0000000..86f3eab --- /dev/null +++ b/src/pyhepmc/_attributes.py @@ -0,0 +1,48 @@ +from ._core import AttributesView, RunInfoAttributesView +import typing as _tp + + +def _clear(self: _tp.Any) -> None: + for key in list(self): + del self[key] + + +def _items( + self: _tp.Any, +) -> _tp.Generator[_tp.Tuple[str, _tp.Any], None, None]: + for name in self: + yield name, self[name] + + +def _eq(self: _tp.Any, other: _tp.Dict[str, _tp.Any]) -> bool: + if len(self) != len(other): + return False + for k, v in self.items(): + if k not in other: + return False + if v != other[k]: + return False + return True + + +def _repr(self: _tp.Any) -> str: + s = f"<{self.__class__.__name__}>{{" + first = True + for k, v in self.items(): + if not first: + s += ", " + first = False + s += f"{k!r}: {v!r}" + s += "}" + return s + + +AttributesView.clear = _clear +AttributesView.items = _items +AttributesView.__eq__ = _eq +AttributesView.__repr__ = _repr + +RunInfoAttributesView.clear = _clear +RunInfoAttributesView.items = _items +RunInfoAttributesView.__eq__ = _eq +RunInfoAttributesView.__repr__ = _repr diff --git a/src/pyhepmc/_doc.py b/src/pyhepmc/_doc.py index 87b6563..84bbb73 100644 --- a/src/pyhepmc/_doc.py +++ b/src/pyhepmc/_doc.py @@ -114,6 +114,10 @@ If this attribute is not set, returns None. """, "GenEvent.vertices": "Access list of vertices.", + "attributes": """Access attributes with a dict-like view. + + It is possible to read and write attributes. Primitive C++ types (and vectors therefore) are converted from/to native Python types. + """, } doc.update(override) diff --git a/src/pyhepmc/io.py b/src/pyhepmc/io.py index 9bf8748..f1c876b 100644 --- a/src/pyhepmc/io.py +++ b/src/pyhepmc/io.py @@ -46,6 +46,9 @@ def __next__(self) -> GenEvent: raise StopIteration return evt + def __iter__(self) -> "_Iter": + return self + next = __next__ diff --git a/src/runinfo_attributes_view.cpp b/src/runinfo_attributes_view.cpp new file mode 100644 index 0000000..7ae3db6 --- /dev/null +++ b/src/runinfo_attributes_view.cpp @@ -0,0 +1,65 @@ +#include "attributes_view.hpp" +#include "pointer.hpp" +#include "pybind.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// To avoid the superfluous copy, we use the legal crowbar +// to access the private attribute map of GenRunInfo +MEMBER_ACCESSOR(MA2, HepMC3::GenRunInfo, m_attributes, + HepMC3::RunInfoAttributesView::AttributeMap) + +namespace HepMC3 { +py::object RunInfoAttributesView::Iter::next() { + if (it_ != end_) return py::cast(it_++->first); + throw py::stop_iteration(); +} + +RunInfoAttributesView::Iter RunInfoAttributesView::iter() { + auto& m = attributes(); + return {m.begin(), m.end()}; +} + +py::object RunInfoAttributesView::getitem(py::str name) { + auto it = find(name); + if (it == attributes().end()) throw py::key_error(name); + return attribute_to_python(it->second); +} + +void RunInfoAttributesView::setitem(py::str name, py::object value) { + auto& amap = attributes(); + amap[py::cast(name)] = attribute_from_python(value); +} + +void RunInfoAttributesView::delitem(py::str name) { + auto it = find(name); + if (it == attributes().end()) throw py::key_error(name); + attributes().erase(it); +} + +bool RunInfoAttributesView::contains(py::str name) { + return find(name) != attributes().end(); +} + +RunInfoAttributesView::AttributeMap& RunInfoAttributesView::attributes() { + auto ref = accessor::accessMember(*run_info_); + return ref.get(); +} + +py::ssize_t RunInfoAttributesView::len() { + return static_cast(attributes().size()); +} + +RunInfoAttributesView::iterator RunInfoAttributesView::find(py::str name) { + return attributes().find(py::cast(name)); +} + +} // namespace HepMC3 diff --git a/tests/test_basic.py b/tests/test_basic.py index 45002a5..edb4c6c 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -98,6 +98,160 @@ def evt(): return evt +def test_GenHeavyIon(): + hi = hep.GenHeavyIon() + assert hi == hep.GenHeavyIon() + hi.Ncoll_hard = 3 + assert hi != hep.GenHeavyIon() + + +def test_GenCrossSection(): + cs = hep.GenCrossSection(1.2, 0.2, 3, 10) + assert cs.event is None + cs.set_xsec(0, 1.2) + with pytest.raises(KeyError): + cs.set_xsec_err("foo", 0.2) + with pytest.raises(KeyError): + cs.xsec("foo") + assert cs.xsec(0) == 1.2 + assert cs.xsec_err(0) == 0.2 + assert cs.xsec() == 1.2 + with pytest.raises(IndexError): + assert cs.xsec(1) + + evt = hep.GenEvent() + evt.run_info = hep.GenRunInfo() + evt.run_info.weight_names = ("foo", "bar") # optional + evt.weights = [1.0, 2.0] + evt.cross_section = cs + assert evt.cross_section.event is evt + cs = evt.cross_section + cs.set_xsec("foo", 1.3) + cs.set_xsec("bar", 2.3) + assert cs.xsec("foo") == 1.3 + assert cs.xsec("bar") == 2.3 + assert cs.xsec(0) == 1.3 + assert cs.xsec(1) == 2.3 + with pytest.raises(IndexError): + cs.xsec(2) + with pytest.raises(KeyError): + cs.xsec("baz") + with pytest.raises(KeyError): + cs.xsec_err("baz") + + +def test_attributes_0(): + ri = hep.GenRunInfo() + att = ri.attributes + assert att == {} + assert len(att) == 0 + assert repr(att) == r"{}" + att["foo"] = 1 + att["bar"] = "xy" + att["baz"] = True + assert att["foo"] == 1 + assert att["bar"] == "xy" + assert att["baz"] is True + with pytest.raises(KeyError): + att["xyz"] + assert len(att) == 3 + assert att == {"baz": True, "foo": 1, "bar": "xy"} + # AttributeMapView has sorted keys + assert repr(att) == r"{'bar': 'xy', 'baz': True, 'foo': 1}" + + del att["bar"] + assert len(att) == 2 + assert att == {"baz": True, "foo": 1} + + keys = [k for k in att] + assert keys == ["baz", "foo"] + + assert len(ri.attributes) == 2 + assert ri.attributes == att + att.clear() + assert len(att) == 0 + assert att == {} + + del ri + # att should keep GenRunInfo alive through shared_ptr + assert len(att) == 0 + + +def test_attributes_1(): + pass + + +def test_attributes_2(evt): + att = evt.attributes + assert att == {} + assert len(att) == 0 + assert repr(att) == r"{}" + att["foo"] = 1 + att["bar"] = "xy" + att["baz"] = True + assert att["foo"] == 1 + assert att["bar"] == "xy" + assert att["baz"] is True + with pytest.raises(KeyError): + att["xyz"] + assert len(att) == 3 + assert att == {"baz": True, "foo": 1, "bar": "xy"} + # AttributeMapView has sorted keys + assert repr(att) == r"{'bar': 'xy', 'baz': True, 'foo': 1}" + + del att["bar"] + assert len(att) == 2 + assert att == {"baz": True, "foo": 1} + + keys = [k for k in att] + assert keys == ["baz", "foo"] + + assert len(evt.attributes) == 2 + assert evt.attributes == att + att.clear() + assert len(att) == 0 + assert att == {} + + +@pytest.mark.parametrize( + "value", + [ + True, + 2, + 1.5, + "baz", + [1, 2], + ["foo", "bar"], + [True, False], + hep.GenCrossSection(1.2, 0.2, 3, 10), + hep.GenHeavyIon(), + hep.GenPdfInfo(), + hep.HEPRUPAttribute(), + hep.HEPEUPAttribute(), + ], +) +def test_attributes_3(evt, value): + p1 = evt.particles[0] + assert p1.id == 1 + assert p1.attributes == {} + p1.attributes = {"foo": value} + assert p1.attributes == {"foo": value} + p1.attributes = {"bar": value} + assert p1.attributes == {"bar": value} + p1.attributes = {} + assert p1.attributes == {} + + v1 = evt.vertices[0] + assert v1.id != p1.id + assert v1.attributes == {} + v1.attributes = {"foo": value} + assert v1.attributes == {"foo": value} + v1.attributes = {"bar": value} + assert v1.attributes == {"bar": value} + v1.attributes = {} + assert v1.attributes == {} + + def test_FourVector(): a = hep.FourVector(1, 2, 3, 4) b = hep.FourVector([1, 2, 3, 4]) diff --git a/tests/test_init.py b/tests/test_init.py index 08114e1..cdf0a19 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -13,6 +13,8 @@ def test_init(): core_names.remove("stringstream") core_names.remove("Attribute") + core_names.remove("AttributesView") + core_names.remove("RunInfoAttributesView") names = set() for name in dir(pyhepmc):