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