src/plugins/Arrow.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/Framework/IAlgorithm.hpp"
0010 #include "ActsExamples/Io/Arrow/ArrowParticleOutputConverter.hpp"
0011 #include "ActsExamples/Io/Arrow/ArrowSimHitOutputConverter.hpp"
0012 #include "ActsExamples/Io/Arrow/ArrowTrackOutputConverter.hpp"
0013 #include "ActsExamples/Io/Parquet/ArrowOutputConverter.hpp"
0014 #include "ActsExamples/Io/Parquet/ParquetReader.hpp"
0015 #include "ActsExamples/Io/Parquet/ParquetWriter.hpp"
0016 #include "ActsPython/Utilities/Helpers.hpp"
0017 #include "ActsPython/Utilities/Macros.hpp"
0018
0019 #include <pybind11/functional.h>
0020 #include <pybind11/pybind11.h>
0021 #include <pybind11/stl.h>
0022 #include <pybind11/stl/filesystem.h>
0023
0024 namespace py = pybind11;
0025 using namespace pybind11::literals;
0026
0027 using namespace Acts;
0028 using namespace ActsPython;
0029 using namespace ActsExamples;
0030
0031 PYBIND11_MODULE(ActsExamplesPythonBindingsArrow, m) {
0032   // ArrowSchemaHandle (Python: `acts.arrow.ArrowSchema`) is registered by
0033   // the plugin-level binding `ActsPluginsPythonBindingsArrow`. The
0034   // generated `acts/examples/arrow.py` does `from acts import arrow`
0035   // before importing this module so that the type is in pybind's
0036   // registry by the time `expectedSchemas` is bound below.
0037   ACTS_PYTHON_DECLARE_READER(ParquetReader, m, "ParquetReader", inputDir,
0038                              collections, expectedSchemas);
0039
0040   ACTS_PYTHON_DECLARE_WRITER(ParquetWriter, m, "ParquetWriter", outputDir,
0041                              collections, expectedSchemas, eventsPerShard,
0042                              eventsPerRowGroup, maxOpenShards);
0043
0044   py::class_<ArrowOutputConverter, IAlgorithm,
0045              std::shared_ptr<ArrowOutputConverter>>(m, "ArrowOutputConverter")
0046       .def_property_readonly("collections", &ArrowOutputConverter::collections);
0047
0048   {
0049     auto [alg, c] =
0050         declareAlgorithm<ArrowParticleOutputConverter, ArrowOutputConverter>(
0051             m, "ArrowParticleOutputConverter");
0052     ACTS_PYTHON_STRUCT(c, inputParticles, outputTable);
0053   }
0054
0055   {
0056     auto [alg, c] =
0057         declareAlgorithm<ArrowTrackOutputConverter, ArrowOutputConverter>(
0058             m, "ArrowTrackOutputConverter");
0059     ACTS_PYTHON_STRUCT(c, inputTracks, inputTrackParticleMatching,
0060                        inputParticles, inputMeasurementSimHitsMap, outputTable,
0061                        writeTime);
0062   }
0063
0064   {
0065     auto [alg, c] =
0066         declareAlgorithm<ArrowSimHitOutputConverter, ArrowOutputConverter>(
0067             m, "ArrowSimHitOutputConverter");
0068     ACTS_PYTHON_STRUCT(c, inputSimHits, inputParticles, inputClusters,
0069                        inputSimHitMeasurementsMap, outputTable,
0070                        detectorResolver);
0071   }
0072
0073   m.def("makeVolumeIdDetectorResolver",
0074         &ArrowSimHitOutputConverter::makeVolumeIdDetectorResolver,
0075         "volumeToDetector"_a, "defaultValue"_a = static_cast<std::uint8_t>(255),
0076         R"doc(
0077 Create a C++ detector resolver from a volume-id lookup map.
0078
0079 This avoids per-hit Python callback overhead by baking the mapping into
0080 a C++ lambda once at configuration time.
0081 )doc");
0082 }