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/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 #include "ActsPlugins/Root/RootMaterialMapIo.hpp"
0030 #include "ActsPython/Utilities/Helpers.hpp"
0031 #include "ActsPython/Utilities/Macros.hpp"
0032
0033 #include <pybind11/pybind11.h>
0034 #include <pybind11/stl.h>
0035 #include <pybind11/stl/filesystem.h>
0036
0037 namespace Acts {
0038 class TrackingGeometry;
0039 namespace detail {
0040 struct Step;
0041 }  // namespace detail
0042 }  // namespace Acts
0043 namespace ActsExamples {
0044 class IWriter;
0045 struct AlgorithmContext;
0046 }  // namespace ActsExamples
0047
0048 namespace py = pybind11;
0049 using namespace pybind11::literals;
0050
0051 using namespace Acts;
0052 using namespace ActsExamples;
0053
0054 namespace ActsPython {
0055
0056 void addRootOutput(Context& ctx) {
0057   auto& mex = ctx.get("examples");
0058
0059   // Bindings for the binning in e.g., TrackFinderPerformanceWriter
0060   {
0061     py::class_<PlotHelpers::Binning>(mex, "Binning")
0062         .def(py::init<std::string, int, double, double>(), "title"_a, "bins"_a,
0063              "bMin"_a, "bMax"_a)
0064         .def(py::init<std::string, std::vector<double>>(), "title"_a, "bins"_a);
0065
0066     py::class_<EffPlotTool::Config>(mex, "EffPlotToolConfig")
0067         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0068              "varBinning"_a);
0069
0070     py::class_<FakePlotTool::Config>(mex, "FakePlotToolConfig")
0071         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0072              "varBinning"_a);
0073
0074     py::class_<DuplicationPlotTool::Config>(mex, "DuplicationPlotToolConfig")
0075         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0076              "varBinning"_a);
0077   }
0078
0079   // ROOT WRITERS
0080   ACTS_PYTHON_DECLARE_WRITER(RootPropagationStepsWriter, mex,
0081                              "RootPropagationStepsWriter", collection, filePath,
0082                              fileMode);
0083
0084   ACTS_PYTHON_DECLARE_WRITER(RootPropagationSummaryWriter, mex,
0085                              "RootPropagationSummaryWriter",
0086                              inputSummaryCollection, filePath, fileMode);
0087
0088   ACTS_PYTHON_DECLARE_WRITER(RootParticleWriter, mex, "RootParticleWriter",
0089                              inputParticles, filePath, fileMode, treeName,
0090                              referencePoint, bField, writeHelixParameters);
0091
0092   ACTS_PYTHON_DECLARE_WRITER(RootVertexWriter, mex, "RootVertexWriter",
0093                              inputVertices, filePath, fileMode, treeName);
0094
0095   ACTS_PYTHON_DECLARE_WRITER(RootMuonSpacePointWriter, mex,
0096                              "RootMuonSpacePointWriter", inputSpacePoints,
0097                              filePath, fileMode, treeName, trackingGeometry,
0098                              writeGlobal);
0099
0100   ACTS_PYTHON_DECLARE_WRITER(
0101       TrackFinderNTupleWriter, mex, "TrackFinderNTupleWriter", inputTracks,
0102       inputParticles, inputParticleMeasurementsMap, inputTrackParticleMatching,
0103       filePath, fileMode, treeNameTracks, treeNameParticles);
0104
0105   ACTS_PYTHON_DECLARE_WRITER(
0106       TrackFitterPerformanceWriter, mex, "TrackFitterPerformanceWriter",
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 ac =
0166         py::class_<ActsPlugins::RootMaterialMapIo::Config>(w, "AccessorConfig")
0167             .def(py::init<>());
0168
0169     ACTS_PYTHON_STRUCT(ac, volumePrefix, portalPrefix, layerPrefix,
0170                        passivePrefix, sensitivePrefix, nBinsHistName,
0171                        axisDirHistName, axisBoundaryTypeHistName, indexHistName,
0172                        minRangeHistName, maxRangeHistName, thicknessHistName,
0173                        x0HistName, l0HistName, aHistName, zHistName,
0174                        rhoHistName);
0175
0176     auto ao = py::class_<ActsPlugins::RootMaterialMapIo::Options>(
0177                   w, "AccessorOptions")
0178                   .def(py::init<>());
0179
0180     ACTS_PYTHON_STRUCT(ao, homogeneousMaterialTreeName, indexedMaterialTreeName,
0181                        folderSurfaceNameBase, folderVolumeNameBase,
0182                        indexedMaterial);
0183
0184     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0185
0186     ACTS_PYTHON_STRUCT(c, processSensitives, processApproaches,
0187                        processRepresenting, processBoundaries, accessorConfig,
0188                        accessorOptions, filePath, fileMode);
0189   }
0190
0191   ACTS_PYTHON_DECLARE_WRITER(RootSeedWriter, mex, "RootSeedWriter", inputSeeds,
0192                              writingMode, filePath, fileMode, treeName);
0193
0194   ACTS_PYTHON_DECLARE_WRITER(RootSimHitWriter, mex, "RootSimHitWriter",
0195                              inputSimHits, filePath, fileMode, treeName);
0196
0197   ACTS_PYTHON_DECLARE_WRITER(RootSpacepointWriter, mex, "RootSpacepointWriter",
0198                              inputSpacepoints, inputMeasurementParticlesMap,
0199                              filePath, fileMode, treeName);
0200
0201   ACTS_PYTHON_DECLARE_WRITER(
0202       RootTrackStatesWriter, mex, "RootTrackStatesWriter", inputTracks,
0203       inputParticles, inputTrackParticleMatching, inputSimHits,
0204       inputMeasurementSimHitsMap, filePath, treeName, fileMode);
0205
0206   ACTS_PYTHON_DECLARE_WRITER(
0207       RootTrackSummaryWriter, mex, "RootTrackSummaryWriter", inputTracks,
0208       inputParticles, inputTrackParticleMatching, filePath, treeName, fileMode,
0209       writeCovMat, writeGsfSpecific, writeGx2fSpecific);
0210
0211   ACTS_PYTHON_DECLARE_WRITER(
0212       VertexNTupleWriter, mex, "VertexNTupleWriter", inputVertices, inputTracks,
0213       inputTruthVertices, inputParticles, inputSelectedParticles,
0214       inputTrackParticleMatching, bField, filePath, treeName, fileMode,
0215       vertexMatchThreshold, trackMatchThreshold, writeTrackInfo);
0216
0217   ACTS_PYTHON_DECLARE_WRITER(
0218       TrackFinderPerformanceWriter, mex, "TrackFinderPerformanceWriter",
0219       inputTracks, inputParticles, inputTrackParticleMatching,
0220       inputParticleTrackMatching, inputParticleMeasurementsMap, filePath,
0221       fileMode, effPlotToolConfig, fakePlotToolConfig,
0222       duplicationPlotToolConfig, trackSummaryPlotToolConfig,
0223       subDetectorTrackSummaryVolumes, writeMatchingDetails);
0224
0225   ACTS_PYTHON_DECLARE_WRITER(RootNuclearInteractionParametersWriter, mex,
0226                              "RootNuclearInteractionParametersWriter",
0227                              inputSimulationProcesses, filePath, fileMode,
0228                              interactionProbabilityBins, momentumBins,
0229                              invariantMassBins, multiplicityMax,
0230                              writeOptionalHistograms, nSimulatedEvents);
0231 }
0232
0233 }  // namespace ActsPython