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 "Acts/Plugins/Python/Utilities.hpp"
0010 #include "ActsExamples/Io/HepMC3/HepMC3InputConverter.hpp"
0011 #include "ActsExamples/Io/HepMC3/HepMC3OutputConverter.hpp"
0012 #include "ActsExamples/Io/HepMC3/HepMC3Reader.hpp"
0013 #include "ActsExamples/Io/HepMC3/HepMC3Util.hpp"
0014 #include "ActsExamples/Io/HepMC3/HepMC3Writer.hpp"
0015
0016 #include <pybind11/pybind11.h>
0017 #include <pybind11/stl.h>
0018 #include <pybind11/stl/filesystem.h>
0019
0020 namespace py = pybind11;
0021 using namespace pybind11::literals;
0022
0023 using namespace Acts;
0024
0025 namespace Acts::Python {
0026 void addHepMC3(Context& ctx) {
0027   auto [m, mex] = ctx.get("main", "examples");
0028
0029   auto hepmc3 = mex.def_submodule("_hepmc3");
0030
0031   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::HepMC3Writer, hepmc3, "HepMC3Writer",
0032                              outputPath, perEvent, inputEvent, compression,
0033                              maxEventsPending, writeEventsInOrder);
0034
0035   ACTS_PYTHON_DECLARE_READER(ActsExamples::HepMC3Reader, hepmc3, "HepMC3Reader",
0036                              inputPath, inputPaths, outputEvent, printListing,
0037                              numEvents, checkEventNumber, maxEventBufferSize,
0038                              vertexGenerator, randomNumbers);
0039
0040   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::HepMC3OutputConverter, hepmc3,
0041                                 "HepMC3OutputConverter", inputParticles,
0042                                 inputVertices, outputEvent);
0043
0044   ACTS_PYTHON_DECLARE_ALGORITHM(
0045       ActsExamples::HepMC3InputConverter, hepmc3, "HepMC3InputConverter",
0046       inputEvent, outputParticles, outputVertices, printListing,
0047       checkConsistency, mergePrimaries, primaryVertexSpatialThreshold,
0048       vertexSpatialThreshold, mergeSecondaries);
0049
0050   {
0051     using enum ActsExamples::HepMC3Util::Compression;
0052     py::enum_<ActsExamples::HepMC3Util::Compression>(hepmc3, "Compression")
0053         .value("none", none)
0054         .value("zlib", zlib)
0055         .value("lzma", lzma)
0056         .value("bzip2", bzip2)
0057         .value("zstd", zstd);
0058   }
0059
0060   hepmc3.def("availableCompressionModes", []() {
0061     auto modes = ActsExamples::HepMC3Util::availableCompressionModes();
0062     return std::vector(modes.begin(), modes.end());
0063   });
0064
0065   hepmc3.def("compressionExtension",
0066              &ActsExamples::HepMC3Util::compressionExtension);
0067 }
0068 }  // namespace Acts::Python