Python/src/RootOutput.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/Root/RootBFieldWriter.hpp"
0011 #include "ActsExamples/Io/Root/RootMaterialTrackWriter.hpp"
0012 #include "ActsExamples/Io/Root/RootMaterialWriter.hpp"
0013 #include "ActsExamples/Io/Root/RootMeasurementWriter.hpp"
0014 #include "ActsExamples/Io/Root/RootNuclearInteractionParametersWriter.hpp"
0015 #include "ActsExamples/Io/Root/RootParticleWriter.hpp"
0016 #include "ActsExamples/Io/Root/RootPropagationStepsWriter.hpp"
0017 #include "ActsExamples/Io/Root/RootPropagationSummaryWriter.hpp"
0018 #include "ActsExamples/Io/Root/RootSeedWriter.hpp"
0019 #include "ActsExamples/Io/Root/RootSimHitWriter.hpp"
0020 #include "ActsExamples/Io/Root/RootSpacepointWriter.hpp"
0021 #include "ActsExamples/Io/Root/RootTrackParameterWriter.hpp"
0022 #include "ActsExamples/Io/Root/RootTrackStatesWriter.hpp"
0023 #include "ActsExamples/Io/Root/RootTrackSummaryWriter.hpp"
0024 #include "ActsExamples/Io/Root/RootVertexWriter.hpp"
0025 #include "ActsExamples/Io/Root/TrackFinderNTupleWriter.hpp"
0026 #include "ActsExamples/Io/Root/TrackFinderPerformanceWriter.hpp"
0027 #include "ActsExamples/Io/Root/TrackFitterPerformanceWriter.hpp"
0028 #include "ActsExamples/Io/Root/VertexNTupleWriter.hpp"
0029
0030 #include <pybind11/pybind11.h>
0031 #include <pybind11/stl.h>
0032 #include <pybind11/stl/filesystem.h>
0033
0034 namespace Acts {
0035 class TrackingGeometry;
0036 namespace detail {
0037 struct Step;
0038 }  // namespace detail
0039 }  // namespace Acts
0040 namespace ActsExamples {
0041 class IWriter;
0042 struct AlgorithmContext;
0043 }  // namespace ActsExamples
0044
0045 namespace py = pybind11;
0046 using namespace pybind11::literals;
0047
0048 using namespace ActsExamples;
0049
0050 namespace Acts::Python {
0051
0052 void addRootOutput(Context& ctx) {
0053   auto [m, mex] = ctx.get("main", "examples");
0054
0055   // Bindings for the binning in e.g., TrackFinderPerformanceWriter
0056   {
0057     py::class_<PlotHelpers::Binning>(mex, "Binning")
0058         .def(py::init<std::string, int, double, double>(), "title"_a, "bins"_a,
0059              "bMin"_a, "bMax"_a)
0060         .def(py::init<std::string, std::vector<double>>(), "title"_a, "bins"_a);
0061
0062     py::class_<EffPlotTool::Config>(mex, "EffPlotToolConfig")
0063         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0064              "varBinning"_a);
0065
0066     py::class_<FakePlotTool::Config>(mex, "FakePlotToolConfig")
0067         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0068              "varBinning"_a);
0069
0070     py::class_<DuplicationPlotTool::Config>(mex, "DuplicationPlotToolConfig")
0071         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0072              "varBinning"_a);
0073   }
0074
0075   // ROOT WRITERS
0076   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootPropagationStepsWriter, mex,
0077                              "RootPropagationStepsWriter", collection, filePath,
0078                              fileMode);
0079
0080   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootPropagationSummaryWriter, mex,
0081                              "RootPropagationSummaryWriter",
0082                              inputSummaryCollection, filePath, fileMode);
0083
0084   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootParticleWriter, mex,
0085                              "RootParticleWriter", inputParticles, filePath,
0086                              fileMode, treeName);
0087
0088   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootVertexWriter, mex,
0089                              "RootVertexWriter", inputVertices, filePath,
0090                              fileMode, treeName);
0091
0092   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::TrackFinderNTupleWriter, mex,
0093                              "TrackFinderNTupleWriter", inputTracks,
0094                              inputParticles, inputParticleMeasurementsMap,
0095                              inputTrackParticleMatching, filePath, fileMode,
0096                              treeNameTracks, treeNameParticles);
0097
0098   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::TrackFitterPerformanceWriter, mex,
0099                              "TrackFitterPerformanceWriter", inputTracks,
0100                              inputParticles, inputTrackParticleMatching,
0101                              filePath, resPlotToolConfig, effPlotToolConfig,
0102                              trackSummaryPlotToolConfig);
0103
0104   ACTS_PYTHON_DECLARE_WRITER(
0105       ActsExamples::RootTrackParameterWriter, mex, "RootTrackParameterWriter",
0106       inputTrackParameters, inputProtoTracks, inputParticles, inputSimHits,
0107       inputMeasurementParticlesMap, inputMeasurementSimHitsMap, filePath,
0108       treeName, fileMode);
0109
0110   ACTS_PYTHON_DECLARE_WRITER(
0111       ActsExamples::RootMaterialTrackWriter, mex, "RootMaterialTrackWriter",
0112       inputMaterialTracks, filePath, fileMode, treeName, recalculateTotals,
0113       prePostStep, storeSurface, storeVolume, collapseInteractions);
0114
0115   {
0116     using Writer = ActsExamples::RootBFieldWriter;
0117     auto w =
0118         py::class_<Writer>(mex, "RootBFieldWriter")
0119             .def_static(
0120                 "run",
0121                 [](const Writer::Config& config, Acts::Logging::Level level) {
0122                   Writer::run(config, Acts::getDefaultLogger("RootBFieldWriter",
0123                                                              level));
0124                 },
0125                 py::arg("config"), py::arg("level"));
0126
0127     py::enum_<Writer::GridType>(w, "GridType")
0128         .value("rz", Writer::GridType::rz)
0129         .value("xyz", Writer::GridType::xyz);
0130
0131     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0132     ACTS_PYTHON_STRUCT(c, treeName, fileName, fileMode, bField, gridType,
0133                        rBounds, zBounds, rBins, zBins, phiBins);
0134   }
0135
0136   {
0137     using Writer = ActsExamples::RootMeasurementWriter;
0138     auto w = py::class_<Writer, IWriter, std::shared_ptr<Writer>>(
0139                  mex, "RootMeasurementWriter")
0140                  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
0141                       py::arg("config"), py::arg("level"));
0142
0143     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0144
0145     ACTS_PYTHON_STRUCT(c, inputMeasurements, inputClusters, inputSimHits,
0146                        inputMeasurementSimHitsMap, filePath, fileMode,
0147                        surfaceByIdentifier);
0148   }
0149
0150   {
0151     using Writer = ActsExamples::RootMaterialWriter;
0152     auto w = py::class_<Writer, IMaterialWriter, std::shared_ptr<Writer>>(
0153                  mex, "RootMaterialWriter")
0154                  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
0155                       py::arg("config"), py::arg("level"))
0156                  .def("write", py::overload_cast<const Acts::TrackingGeometry&>(
0157                                    &Writer::write));
0158
0159     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0160
0161     ACTS_PYTHON_STRUCT(c, processSensitives, processApproaches,
0162                        processRepresenting, processBoundaries, processVolumes,
0163                        folderSurfaceNameBase, folderVolumeNameBase, voltag,
0164                        boutag, laytag, apptag, sentag, ntag, vtag, otag, mintag,
0165                        maxtag, ttag, x0tag, l0tag, atag, ztag, rhotag, filePath,
0166                        fileMode);
0167   }
0168
0169   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSeedWriter, mex,
0170                              "RootSeedWriter", inputSeeds, writingMode,
0171                              filePath, fileMode, treeName);
0172
0173   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSimHitWriter, mex,
0174                              "RootSimHitWriter", inputSimHits, filePath,
0175                              fileMode, treeName);
0176
0177   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSpacepointWriter, mex,
0178                              "RootSpacepointWriter", inputSpacepoints, filePath,
0179                              fileMode, treeName);
0180
0181   ACTS_PYTHON_DECLARE_WRITER(
0182       ActsExamples::RootTrackStatesWriter, mex, "RootTrackStatesWriter",
0183       inputTracks, inputParticles, inputTrackParticleMatching, inputSimHits,
0184       inputMeasurementSimHitsMap, filePath, treeName, fileMode);
0185
0186   ACTS_PYTHON_DECLARE_WRITER(
0187       ActsExamples::RootTrackSummaryWriter, mex, "RootTrackSummaryWriter",
0188       inputTracks, inputParticles, inputTrackParticleMatching, filePath,
0189       treeName, fileMode, writeCovMat, writeGsfSpecific, writeGx2fSpecific);
0190
0191   ACTS_PYTHON_DECLARE_WRITER(
0192       ActsExamples::VertexNTupleWriter, mex, "VertexNTupleWriter",
0193       inputVertices, inputTracks, inputTruthVertices, inputParticles,
0194       inputSelectedParticles, inputTrackParticleMatching, bField, filePath,
0195       treeName, fileMode, vertexMatchThreshold, trackMatchThreshold,
0196       writeTrackInfo);
0197
0198   ACTS_PYTHON_DECLARE_WRITER(
0199       ActsExamples::TrackFinderPerformanceWriter, mex,
0200       "TrackFinderPerformanceWriter", inputTracks, inputParticles,
0201       inputTrackParticleMatching, inputParticleTrackMatching,
0202       inputParticleMeasurementsMap, filePath, fileMode, effPlotToolConfig,
0203       fakePlotToolConfig, duplicationPlotToolConfig, trackSummaryPlotToolConfig,
0204       subDetectorTrackSummaryVolumes, writeMatchingDetails);
0205
0206   ACTS_PYTHON_DECLARE_WRITER(
0207       ActsExamples::RootNuclearInteractionParametersWriter, mex,
0208       "RootNuclearInteractionParametersWriter", inputSimulationProcesses,
0209       filePath, fileMode, interactionProbabilityBins, momentumBins,
0210       invariantMassBins, multiplicityMax, writeOptionalHistograms,
0211       nSimulatedEvents);
0212 }
0213
0214 }  // namespace Acts::Python