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