src/plugins/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 using namespace ActsPython;
0028
0029 PYBIND11_MODULE(ActsExamplesPythonBindingsHepMC3, hepmc3) {
0030   ACTS_PYTHON_DECLARE_WRITER(HepMC3Writer, hepmc3, "HepMC3Writer", outputPath,
0031                              inputEvent, compression, maxEventsPending,
0032                              writeEventsInOrder);
0033
0034   // Declare the HepMC3Reader class first
0035   auto reader =
0036       py::class_<HepMC3Reader, IReader, std::shared_ptr<HepMC3Reader>>(
0037           hepmc3, "HepMC3Reader")
0038           .def(py::init<const HepMC3Reader::Config&, Acts::Logging::Level>(),
0039                py::arg("config"), py::arg("level"))
0040           .def_property_readonly("config", &HepMC3Reader::config);
0041
0042   // Expose Input struct as a nested class of HepMC3Reader
0043   py::class_<HepMC3Reader::Input>(reader, "Input")
0044       .def(py::init<>())
0045       .def(py::init([](const std::filesystem::path& path,
0046                        std::shared_ptr<const MultiplicityGenerator>
0047                            multiplicityGenerator) {
0048              HepMC3Reader::Input inp;
0049              inp.path = path;
0050              inp.multiplicityGenerator = multiplicityGenerator;
0051              return inp;
0052            }),
0053            py::arg("path"), py::arg("multiplicityGenerator"))
0054       .def_readwrite("path", &HepMC3Reader::Input::path)
0055       .def_readwrite("multiplicityGenerator",
0056                      &HepMC3Reader::Input::multiplicityGenerator)
0057       // Factory methods for convenience
0058       .def_static(
0059           "Fixed",
0060           [](const std::filesystem::path& path, std::size_t n) {
0061             HepMC3Reader::Input inp;
0062             inp.path = path;
0063             inp.multiplicityGenerator =
0064                 std::make_shared<FixedMultiplicityGenerator>(n);
0065             return inp;
0066           },
0067           py::arg("path"), py::arg("n") = 1,
0068           "Create Input with FixedMultiplicityGenerator")
0069       .def_static(
0070           "Poisson",
0071           [](const std::filesystem::path& path, double mean) {
0072             HepMC3Reader::Input inp;
0073             inp.path = path;
0074             inp.multiplicityGenerator =
0075                 std::make_shared<PoissonMultiplicityGenerator>(mean);
0076             return inp;
0077           },
0078           py::arg("path"), py::arg("mean"),
0079           "Create Input with PoissonMultiplicityGenerator");
0080
0081   auto config = py::class_<HepMC3Reader::Config>(reader, "Config")
0082                     .def(py::init<>(), "Default constructor");
0083   // Now configure the HepMC3Reader itself
0084   ACTS_PYTHON_STRUCT(config, inputs, inputPath, outputEvent, printListing,
0085                      numEvents, checkEventNumber, maxEventBufferSize,
0086                      vertexGenerator, randomNumbers);
0087
0088   ACTS_PYTHON_DECLARE_ALGORITHM(HepMC3OutputConverter, hepmc3,
0089                                 "HepMC3OutputConverter", inputParticles,
0090                                 inputVertices, outputEvent);
0091
0092   ACTS_PYTHON_DECLARE_ALGORITHM(
0093       HepMC3InputConverter, hepmc3, "HepMC3InputConverter", inputEvent,
0094       outputParticles, outputVertices, printListing, checkConsistency,
0095       mergePrimaries, primaryVertexSpatialThreshold, vertexSpatialThreshold,
0096       mergeSecondaries);
0097
0098   {
0099     using enum HepMC3Util::Compression;
0100     py::enum_<HepMC3Util::Compression>(hepmc3, "Compression")
0101         .value("none", none)
0102         .value("zlib", zlib)
0103         .value("lzma", lzma)
0104         .value("bzip2", bzip2)
0105         .value("zstd", zstd);
0106   }
0107
0108   {
0109     using enum HepMC3Util::Format;
0110     py::enum_<HepMC3Util::Format>(hepmc3, "Format")
0111         .value("ascii", ascii)
0112         .value("root", root);
0113   }
0114
0115   hepmc3.def("availableCompressionModes", []() {
0116     auto modes = HepMC3Util::availableCompressionModes();
0117     return std::vector(modes.begin(), modes.end());
0118   });
0119
0120   hepmc3.def("availableFormats", []() {
0121     auto formats = HepMC3Util::availableFormats();
0122     return std::vector(formats.begin(), formats.end());
0123   });
0124
0125   hepmc3.def("compressionExtension", &HepMC3Util::compressionExtension);
0126   hepmc3.def("compressionFromFilename", &HepMC3Util::compressionFromFilename,
0127              py::arg("filename"));
0128   hepmc3.def("formatFromFilename", &HepMC3Util::formatFromFilename,
0129              py::arg("filename"));
0130
0131   // HepMC3 normalize function and result
0132   {
0133     auto result =
0134         py::class_<HepMC3Util::NormalizeResult>(hepmc3, "NormalizeResult")
0135             .def(py::init<>())
0136             .def_readonly("numEvents", &HepMC3Util::NormalizeResult::numEvents)
0137             .def_readonly("outputFiles",
0138                           &HepMC3Util::NormalizeResult::outputFiles)
0139             .def_readonly("totalInputSize",
0140                           &HepMC3Util::NormalizeResult::totalInputSize)
0141             .def_readonly("totalOutputSize",
0142                           &HepMC3Util::NormalizeResult::totalOutputSize)
0143             .def_readonly("totalReadTime",
0144                           &HepMC3Util::NormalizeResult::totalReadTime)
0145             .def_readonly("totalWriteTime",
0146                           &HepMC3Util::NormalizeResult::totalWriteTime);
0147   }
0148
0149   hepmc3.def("normalizeFiles", &HepMC3Util::normalizeFiles,
0150              py::arg("inputFiles"), py::arg("singleOutputPath") = std::nullopt,
0151              py::arg("outputDir") = ".", py::arg("outputPrefix") = "events",
0152              py::arg("eventsPerFile") = 10000,
0153              py::arg("maxEvents") = std::nullopt,
0154              py::arg("format") = HepMC3Util::Format::ascii,
0155              py::arg("compression") = HepMC3Util::Compression::none,
0156              py::arg("compressionLevel") = 6, py::arg("verbose") = false);
0157 }