Python/src/HepMC3.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/Io/HepMC3/HepMC3InputConverter.hpp"
0010 #include "ActsExamples/Io/HepMC3/HepMC3OutputConverter.hpp"
0011 #include "ActsExamples/Io/HepMC3/HepMC3Reader.hpp"
0012 #include "ActsExamples/Io/HepMC3/HepMC3Util.hpp"
0013 #include "ActsExamples/Io/HepMC3/HepMC3Writer.hpp"
0014 #include "ActsExamples/Utilities/MultiplicityGenerators.hpp"
0015 #include "ActsPython/Utilities/Helpers.hpp"
0016 #include "ActsPython/Utilities/Macros.hpp"
0017
0018 #include <pybind11/pybind11.h>
0019 #include <pybind11/stl.h>
0020 #include <pybind11/stl/filesystem.h>
0021
0022 namespace py = pybind11;
0023 using namespace pybind11::literals;
0024
0025 using namespace Acts;
0026 using namespace ActsExamples;
0027
0028 namespace ActsPython {
0029 void addHepMC3(Context& ctx) {
0030   auto [m, mex] = ctx.get("main", "examples");
0031
0032   auto hepmc3 = mex.def_submodule("_hepmc3");
0033
0034   ACTS_PYTHON_DECLARE_WRITER(HepMC3Writer, hepmc3, "HepMC3Writer", outputPath,
0035                              inputEvent, compression, maxEventsPending,
0036                              writeEventsInOrder);
0037
0038   // Declare the HepMC3Reader class first
0039   auto reader =
0040       py::class_<HepMC3Reader, IReader, std::shared_ptr<HepMC3Reader>>(
0041           hepmc3, "HepMC3Reader")
0042           .def(py::init<const HepMC3Reader::Config&, Acts::Logging::Level>(),
0043                py::arg("config"), py::arg("level"))
0044           .def_property_readonly("config", &HepMC3Reader::config);
0045
0046   // Expose Input struct as a nested class of HepMC3Reader
0047   py::class_<HepMC3Reader::Input>(reader, "Input")
0048       .def(py::init<>())
0049       .def(py::init([](const std::filesystem::path& path,
0050                        std::shared_ptr<const MultiplicityGenerator>
0051                            multiplicityGenerator) {
0052              HepMC3Reader::Input inp;
0053              inp.path = path;
0054              inp.multiplicityGenerator = multiplicityGenerator;
0055              return inp;
0056            }),
0057            py::arg("path"), py::arg("multiplicityGenerator"))
0058       .def_readwrite("path", &HepMC3Reader::Input::path)
0059       .def_readwrite("multiplicityGenerator",
0060                      &HepMC3Reader::Input::multiplicityGenerator)
0061       // Factory methods for convenience
0062       .def_static(
0063           "Fixed",
0064           [](const std::filesystem::path& path, std::size_t n) {
0065             HepMC3Reader::Input inp;
0066             inp.path = path;
0067             inp.multiplicityGenerator =
0068                 std::make_shared<FixedMultiplicityGenerator>(n);
0069             return inp;
0070           },
0071           py::arg("path"), py::arg("n") = 1,
0072           "Create Input with FixedMultiplicityGenerator")
0073       .def_static(
0074           "Poisson",
0075           [](const std::filesystem::path& path, double mean) {
0076             HepMC3Reader::Input inp;
0077             inp.path = path;
0078             inp.multiplicityGenerator =
0079                 std::make_shared<PoissonMultiplicityGenerator>(mean);
0080             return inp;
0081           },
0082           py::arg("path"), py::arg("mean"),
0083           "Create Input with PoissonMultiplicityGenerator");
0084
0085   auto config = py::class_<HepMC3Reader::Config>(reader, "Config")
0086                     .def(py::init<>(), "Default constructor");
0087   // Now configure the HepMC3Reader itself
0088   ACTS_PYTHON_STRUCT(config, inputs, inputPath, outputEvent, printListing,
0089                      numEvents, checkEventNumber, maxEventBufferSize,
0090                      vertexGenerator, randomNumbers);
0091
0092   ACTS_PYTHON_DECLARE_ALGORITHM(HepMC3OutputConverter, hepmc3,
0093                                 "HepMC3OutputConverter", inputParticles,
0094                                 inputVertices, outputEvent);
0095
0096   ACTS_PYTHON_DECLARE_ALGORITHM(
0097       HepMC3InputConverter, hepmc3, "HepMC3InputConverter", inputEvent,
0098       outputParticles, outputVertices, printListing, checkConsistency,
0099       mergePrimaries, primaryVertexSpatialThreshold, vertexSpatialThreshold,
0100       mergeSecondaries);
0101
0102   {
0103     using enum HepMC3Util::Compression;
0104     py::enum_<HepMC3Util::Compression>(hepmc3, "Compression")
0105         .value("none", none)
0106         .value("zlib", zlib)
0107         .value("lzma", lzma)
0108         .value("bzip2", bzip2)
0109         .value("zstd", zstd);
0110   }
0111
0112   {
0113     using enum HepMC3Util::Format;
0114     py::enum_<HepMC3Util::Format>(hepmc3, "Format")
0115         .value("ascii", ascii)
0116         .value("root", root);
0117   }
0118
0119   hepmc3.def("availableCompressionModes", []() {
0120     auto modes = HepMC3Util::availableCompressionModes();
0121     return std::vector(modes.begin(), modes.end());
0122   });
0123
0124   hepmc3.def("availableFormats", []() {
0125     auto formats = HepMC3Util::availableFormats();
0126     return std::vector(formats.begin(), formats.end());
0127   });
0128
0129   hepmc3.def("compressionExtension", &HepMC3Util::compressionExtension);
0130   hepmc3.def("compressionFromFilename", &HepMC3Util::compressionFromFilename,
0131              py::arg("filename"));
0132   hepmc3.def("formatFromFilename", &HepMC3Util::formatFromFilename,
0133              py::arg("filename"));
0134
0135   // HepMC3 normalize function and result
0136   {
0137     auto result =
0138         py::class_<HepMC3Util::NormalizeResult>(hepmc3, "NormalizeResult")
0139             .def(py::init<>())
0140             .def_readonly("numEvents", &HepMC3Util::NormalizeResult::numEvents)
0141             .def_readonly("outputFiles",
0142                           &HepMC3Util::NormalizeResult::outputFiles)
0143             .def_readonly("totalInputSize",
0144                           &HepMC3Util::NormalizeResult::totalInputSize)
0145             .def_readonly("totalOutputSize",
0146                           &HepMC3Util::NormalizeResult::totalOutputSize)
0147             .def_readonly("totalReadTime",
0148                           &HepMC3Util::NormalizeResult::totalReadTime)
0149             .def_readonly("totalWriteTime",
0150                           &HepMC3Util::NormalizeResult::totalWriteTime);
0151   }
0152
0153   hepmc3.def("normalizeFiles", &HepMC3Util::normalizeFiles,
0154              py::arg("inputFiles"), py::arg("singleOutputPath") = std::nullopt,
0155              py::arg("outputDir") = ".", py::arg("outputPrefix") = "events",
0156              py::arg("eventsPerFile") = 10000,
0157              py::arg("maxEvents") = std::nullopt,
0158              py::arg("format") = HepMC3Util::Format::ascii,
0159              py::arg("compression") = HepMC3Util::Compression::none,
0160              py::arg("compressionLevel") = 6, py::arg("verbose") = false);
0161 }
0162 }  // namespace ActsPython