Python/src/Output.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2021-2023 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 http://mozilla.org/MPL/2.0/.
0008
0009 #include "Acts/Definitions/TrackParametrization.hpp"
0010 #include "Acts/Geometry/GeometryHierarchyMap.hpp"
0011 #include "Acts/Plugins/Python/Utilities.hpp"
0012 #include "Acts/Utilities/Logger.hpp"
0013 #include "Acts/Visualization/ViewConfig.hpp"
0014 #include "ActsExamples/Digitization/DigitizationConfig.hpp"
0015 #include "ActsExamples/Framework/ProcessCode.hpp"
0016 #include "ActsExamples/Io/Csv/CsvBFieldWriter.hpp"
0017 #include "ActsExamples/Io/Csv/CsvExaTrkXGraphWriter.hpp"
0018 #include "ActsExamples/Io/Csv/CsvMeasurementWriter.hpp"
0019 #include "ActsExamples/Io/Csv/CsvParticleWriter.hpp"
0020 #include "ActsExamples/Io/Csv/CsvPlanarClusterWriter.hpp"
0021 #include "ActsExamples/Io/Csv/CsvProtoTrackWriter.hpp"
0022 #include "ActsExamples/Io/Csv/CsvSeedWriter.hpp"
0023 #include "ActsExamples/Io/Csv/CsvSimHitWriter.hpp"
0024 #include "ActsExamples/Io/Csv/CsvSpacepointWriter.hpp"
0025 #include "ActsExamples/Io/Csv/CsvTrackParameterWriter.hpp"
0026 #include "ActsExamples/Io/Csv/CsvTrackWriter.hpp"
0027 #include "ActsExamples/Io/Csv/CsvTrackingGeometryWriter.hpp"
0028 #include "ActsExamples/Io/NuclearInteractions/RootNuclearInteractionParametersWriter.hpp"
0029 #include "ActsExamples/Io/Performance/CKFPerformanceWriter.hpp"
0030 #include "ActsExamples/Io/Performance/SeedingPerformanceWriter.hpp"
0031 #include "ActsExamples/Io/Performance/TrackFinderPerformanceWriter.hpp"
0032 #include "ActsExamples/Io/Performance/TrackFitterPerformanceWriter.hpp"
0033 #include "ActsExamples/Io/Performance/VertexPerformanceWriter.hpp"
0034 #include "ActsExamples/Io/Root/RootBFieldWriter.hpp"
0035 #include "ActsExamples/Io/Root/RootMaterialTrackWriter.hpp"
0036 #include "ActsExamples/Io/Root/RootMaterialWriter.hpp"
0037 #include "ActsExamples/Io/Root/RootMeasurementWriter.hpp"
0038 #include "ActsExamples/Io/Root/RootParticleWriter.hpp"
0039 #include "ActsExamples/Io/Root/RootPlanarClusterWriter.hpp"
0040 #include "ActsExamples/Io/Root/RootPropagationStepsWriter.hpp"
0041 #include "ActsExamples/Io/Root/RootSeedWriter.hpp"
0042 #include "ActsExamples/Io/Root/RootSimHitWriter.hpp"
0043 #include "ActsExamples/Io/Root/RootSpacepointWriter.hpp"
0044 #include "ActsExamples/Io/Root/RootTrackParameterWriter.hpp"
0045 #include "ActsExamples/Io/Root/RootTrackStatesWriter.hpp"
0046 #include "ActsExamples/Io/Root/RootTrackSummaryWriter.hpp"
0047 #include "ActsExamples/Io/Root/RootVertexWriter.hpp"
0048 #include "ActsExamples/MaterialMapping/IMaterialWriter.hpp"
0049 #include "ActsExamples/Plugins/Obj/ObjPropagationStepsWriter.hpp"
0050 #include "ActsExamples/Plugins/Obj/ObjTrackingGeometryWriter.hpp"
0051
0052 #include <array>
0053 #include <memory>
0054 #include <string>
0055 #include <tuple>
0056 #include <vector>
0057
0058 #include <pybind11/pybind11.h>
0059 #include <pybind11/stl.h>
0060
0061 namespace Acts {
0062 class TrackingGeometry;
0063 namespace detail {
0064 struct Step;
0065 }  // namespace detail
0066 }  // namespace Acts
0067 namespace ActsExamples {
0068 class IWriter;
0069 struct AlgorithmContext;
0070 }  // namespace ActsExamples
0071
0072 namespace py = pybind11;
0073 using namespace pybind11::literals;
0074
0075 using namespace ActsExamples;
0076
0077 namespace {
0078 template <ActsExamples::CsvBFieldWriter::CoordinateType CType, bool Grid>
0079 void register_csv_bfield_writer_binding(
0080     pybind11::class_<ActsExamples::CsvBFieldWriter>& w) {
0081   std::string name =
0082       std::string(CType == ActsExamples::CsvBFieldWriter::CoordinateType::XYZ
0083                       ? "Xyz"
0084                       : "Rz") +
0085       std::string(Grid ? "Grid" : "Gridless");
0086
0087   using Config = ActsExamples::CsvBFieldWriter::Config<CType, Grid>;
0088   w.def_static((std::string("run") + name).c_str(),
0089                [](const Config& config, Acts::Logging::Level level) {
0090                  ActsExamples::CsvBFieldWriter::run(config, level);
0091                },
0092                py::arg("config"), py::arg("level"));
0093   auto c = py::class_<Config>(w, (std::string("Config") + name).c_str())
0094                .def(py::init<>());
0095   ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0096   ACTS_PYTHON_MEMBER(fileName);
0097   ACTS_PYTHON_MEMBER(bField);
0098   ACTS_PYTHON_MEMBER(range);
0099   ACTS_PYTHON_MEMBER(bins);
0100   ACTS_PYTHON_STRUCT_END();
0101 }
0102 }  // namespace
0103
0104 namespace Acts::Python {
0105
0106 void addOutput(Context& ctx) {
0107   auto [m, mex] = ctx.get("main", "examples");
0108
0109   ACTS_PYTHON_DECLARE_WRITER(
0110       ActsExamples::ObjPropagationStepsWriter<Acts::detail::Step>, mex,
0111       "ObjPropagationStepsWriter", collection, outputDir, outputScalor,
0112       outputPrecision);
0113
0114   {
0115     auto c = py::class_<ViewConfig>(m, "ViewConfig").def(py::init<>());
0116
0117     ACTS_PYTHON_STRUCT_BEGIN(c, ViewConfig);
0118     ACTS_PYTHON_MEMBER(visible);
0119     ACTS_PYTHON_MEMBER(color);
0120     ACTS_PYTHON_MEMBER(offset);
0121     ACTS_PYTHON_MEMBER(lineThickness);
0122     ACTS_PYTHON_MEMBER(surfaceThickness);
0123     ACTS_PYTHON_MEMBER(nSegments);
0124     ACTS_PYTHON_MEMBER(triangulate);
0125     ACTS_PYTHON_MEMBER(outputName);
0126     ACTS_PYTHON_STRUCT_END();
0127
0128     patchKwargsConstructor(c);
0129   }
0130
0131   {
0132     using Writer = ActsExamples::ObjTrackingGeometryWriter;
0133     auto w = py::class_<Writer, std::shared_ptr<Writer>>(
0134                  mex, "ObjTrackingGeometryWriter")
0135                  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
0136                       py::arg("config"), py::arg("level"))
0137                  .def("write", py::overload_cast<const AlgorithmContext&,
0138                                                  const Acts::TrackingGeometry&>(
0139                                    &Writer::write));
0140
0141     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0142     ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
0143     ACTS_PYTHON_MEMBER(outputScalor);
0144     ACTS_PYTHON_MEMBER(outputPrecision);
0145     ACTS_PYTHON_MEMBER(outputDir);
0146     ACTS_PYTHON_MEMBER(containerView);
0147     ACTS_PYTHON_MEMBER(volumeView);
0148     ACTS_PYTHON_MEMBER(sensitiveView);
0149     ACTS_PYTHON_MEMBER(passiveView);
0150     ACTS_PYTHON_MEMBER(gridView);
0151     ACTS_PYTHON_STRUCT_END();
0152   }
0153
0154   // Bindings for the binning in e.g., CKFPerformanceWriter
0155   {
0156     py::class_<PlotHelpers::Binning>(mex, "Binning")
0157         .def(py::init<std::string, int, double, double>(), "title"_a, "bins"_a,
0158              "bMin"_a, "bMax"_a)
0159         .def(py::init<std::string, std::vector<double>>(), "title"_a, "bins"_a);
0160
0161     py::class_<EffPlotTool::Config>(mex, "EffPlotToolConfig")
0162         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0163              "varBinning"_a);
0164
0165     py::class_<FakeRatePlotTool::Config>(mex, "FakeRatePlotToolConfig")
0166         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0167              "varBinning"_a);
0168
0169     py::class_<DuplicationPlotTool::Config>(mex, "DuplicationPlotToolConfig")
0170         .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
0171              "varBinning"_a);
0172   }
0173
0174   // ROOT WRITERS
0175   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootPropagationStepsWriter, mex,
0176                              "RootPropagationStepsWriter", collection, filePath,
0177                              fileMode);
0178
0179   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootParticleWriter, mex,
0180                              "RootParticleWriter", inputParticles,
0181                              inputFinalParticles, filePath, fileMode, treeName);
0182
0183   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootVertexWriter, mex,
0184                              "RootVertexWriter", inputVertices, filePath,
0185                              fileMode, treeName);
0186
0187   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::TrackFinderPerformanceWriter, mex,
0188                              "TrackFinderPerformanceWriter", inputProtoTracks,
0189                              inputParticles, inputMeasurementParticlesMap,
0190                              inputProtoTrackParticleMatching, filePath,
0191                              fileMode, treeNameTracks, treeNameParticles);
0192
0193   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::TrackFitterPerformanceWriter, mex,
0194                              "TrackFitterPerformanceWriter", inputTracks,
0195                              inputParticles, inputTrackParticleMatching,
0196                              filePath, resPlotToolConfig, effPlotToolConfig,
0197                              trackSummaryPlotToolConfig);
0198
0199   ACTS_PYTHON_DECLARE_WRITER(
0200       ActsExamples::SeedingPerformanceWriter, mex, "SeedingPerformanceWriter",
0201       inputSeeds, inputMeasurementParticlesMap, inputParticles, filePath,
0202       fileMode, effPlotToolConfig, duplicationPlotToolConfig);
0203
0204   ACTS_PYTHON_DECLARE_WRITER(
0205       ActsExamples::RootTrackParameterWriter, mex, "RootTrackParameterWriter",
0206       inputTrackParameters, inputProtoTracks, inputParticles, inputSimHits,
0207       inputMeasurementParticlesMap, inputMeasurementSimHitsMap, filePath,
0208       treeName, fileMode);
0209
0210   ACTS_PYTHON_DECLARE_WRITER(
0211       ActsExamples::RootMaterialTrackWriter, mex, "RootMaterialTrackWriter",
0212       inputMaterialTracks, filePath, fileMode, treeName, recalculateTotals,
0213       prePostStep, storeSurface, storeVolume, collapseInteractions);
0214
0215   {
0216     using Writer = ActsExamples::RootBFieldWriter;
0217     auto w =
0218         py::class_<Writer>(mex, "RootBFieldWriter")
0219             .def_static(
0220                 "run",
0221                 [](const Writer::Config& config, Acts::Logging::Level level) {
0222                   Writer::run(config, Acts::getDefaultLogger("RootBFieldWriter",
0223                                                              level));
0224                 },
0225                 py::arg("config"), py::arg("level"));
0226
0227     py::enum_<Writer::GridType>(w, "GridType")
0228         .value("rz", Writer::GridType::rz)
0229         .value("xyz", Writer::GridType::xyz);
0230
0231     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0232     ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
0233     ACTS_PYTHON_MEMBER(treeName);
0234     ACTS_PYTHON_MEMBER(fileName);
0235     ACTS_PYTHON_MEMBER(fileMode);
0236     ACTS_PYTHON_MEMBER(bField);
0237     ACTS_PYTHON_MEMBER(gridType);
0238     ACTS_PYTHON_MEMBER(rBounds);
0239     ACTS_PYTHON_MEMBER(zBounds);
0240     ACTS_PYTHON_MEMBER(rBins);
0241     ACTS_PYTHON_MEMBER(zBins);
0242     ACTS_PYTHON_MEMBER(phiBins);
0243     ACTS_PYTHON_STRUCT_END();
0244   }
0245
0246   {
0247     using Writer = ActsExamples::RootMeasurementWriter;
0248     auto w = py::class_<Writer, IWriter, std::shared_ptr<Writer>>(
0249                  mex, "RootMeasurementWriter")
0250                  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
0251                       py::arg("config"), py::arg("level"));
0252
0253     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0254
0255     c.def("addBoundIndicesFromDigiConfig",
0256           [](Writer::Config& self, const DigitizationConfig& digiCfg) {
0257             self.boundIndices =
0258                 Acts::GeometryHierarchyMap<std::vector<Acts::BoundIndices>>(
0259                     digiCfg.getBoundIndices());
0260           });
0261
0262     ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
0263     ACTS_PYTHON_MEMBER(inputMeasurements);
0264     ACTS_PYTHON_MEMBER(inputClusters);
0265     ACTS_PYTHON_MEMBER(inputSimHits);
0266     ACTS_PYTHON_MEMBER(inputMeasurementSimHitsMap);
0267     ACTS_PYTHON_MEMBER(filePath);
0268     ACTS_PYTHON_MEMBER(fileMode);
0269     ACTS_PYTHON_MEMBER(boundIndices);
0270     ACTS_PYTHON_MEMBER(trackingGeometry);
0271     ACTS_PYTHON_STRUCT_END();
0272   }
0273
0274   py::class_<IMaterialWriter, std::shared_ptr<IMaterialWriter>>(
0275       mex, "IMaterialWriter");
0276
0277   {
0278     using Writer = ActsExamples::RootMaterialWriter;
0279     auto w = py::class_<Writer, IMaterialWriter, std::shared_ptr<Writer>>(
0280                  mex, "RootMaterialWriter")
0281                  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
0282                       py::arg("config"), py::arg("level"))
0283                  .def("write", py::overload_cast<const Acts::TrackingGeometry&>(
0284                                    &Writer::write));
0285
0286     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0287
0288     ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
0289     ACTS_PYTHON_MEMBER(processSensitives);
0290     ACTS_PYTHON_MEMBER(processApproaches);
0291     ACTS_PYTHON_MEMBER(processRepresenting);
0292     ACTS_PYTHON_MEMBER(processBoundaries);
0293     ACTS_PYTHON_MEMBER(processVolumes);
0294     ACTS_PYTHON_MEMBER(folderSurfaceNameBase);
0295     ACTS_PYTHON_MEMBER(folderVolumeNameBase);
0296     ACTS_PYTHON_MEMBER(voltag);
0297     ACTS_PYTHON_MEMBER(boutag);
0298     ACTS_PYTHON_MEMBER(laytag);
0299     ACTS_PYTHON_MEMBER(apptag);
0300     ACTS_PYTHON_MEMBER(sentag);
0301     ACTS_PYTHON_MEMBER(ntag);
0302     ACTS_PYTHON_MEMBER(vtag);
0303     ACTS_PYTHON_MEMBER(otag);
0304     ACTS_PYTHON_MEMBER(mintag);
0305     ACTS_PYTHON_MEMBER(maxtag);
0306     ACTS_PYTHON_MEMBER(ttag);
0307     ACTS_PYTHON_MEMBER(x0tag);
0308     ACTS_PYTHON_MEMBER(l0tag);
0309     ACTS_PYTHON_MEMBER(atag);
0310     ACTS_PYTHON_MEMBER(ztag);
0311     ACTS_PYTHON_MEMBER(rhotag);
0312     ACTS_PYTHON_MEMBER(filePath);
0313     ACTS_PYTHON_MEMBER(fileMode);
0314     ACTS_PYTHON_STRUCT_END();
0315   }
0316
0317   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootPlanarClusterWriter, mex,
0318                              "RootPlanarClusterWriter", inputClusters,
0319                              inputSimHits, filePath, fileMode, treeName,
0320                              trackingGeometry);
0321
0322   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSeedWriter, mex,
0323                              "RootSeedWriter", inputSeeds, writingMode,
0324                              filePath, fileMode, treeName);
0325
0326   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSimHitWriter, mex,
0327                              "RootSimHitWriter", inputSimHits, filePath,
0328                              fileMode, treeName);
0329
0330   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSpacepointWriter, mex,
0331                              "RootSpacepointWriter", inputSpacepoints, filePath,
0332                              fileMode, treeName);
0333
0334   ACTS_PYTHON_DECLARE_WRITER(
0335       ActsExamples::RootTrackStatesWriter, mex, "RootTrackStatesWriter",
0336       inputTracks, inputParticles, inputTrackParticleMatching, inputSimHits,
0337       inputMeasurementSimHitsMap, filePath, treeName, fileMode);
0338
0339   ACTS_PYTHON_DECLARE_WRITER(
0340       ActsExamples::RootTrackSummaryWriter, mex, "RootTrackSummaryWriter",
0341       inputTracks, inputParticles, inputTrackParticleMatching, filePath,
0342       treeName, fileMode, writeCovMat, writeGsfSpecific, writeGx2fSpecific);
0343
0344   ACTS_PYTHON_DECLARE_WRITER(
0345       ActsExamples::VertexPerformanceWriter, mex, "VertexPerformanceWriter",
0346       inputVertices, inputTracks, inputTruthVertices, inputParticles,
0347       inputSelectedParticles, inputTrackParticleMatching, bField, filePath,
0348       treeName, fileMode, vertexMatchThreshold, trackMatchThreshold, useTracks);
0349
0350   // CSV WRITERS
0351   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvParticleWriter, mex,
0352                              "CsvParticleWriter", inputParticles, outputDir,
0353                              outputStem, outputPrecision);
0354
0355   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvMeasurementWriter, mex,
0356                              "CsvMeasurementWriter", inputMeasurements,
0357                              inputClusters, inputMeasurementSimHitsMap,
0358                              outputDir, outputPrecision);
0359
0360   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvPlanarClusterWriter, mex,
0361                              "CsvPlanarClusterWriter", inputClusters,
0362                              inputSimHits, outputDir, outputPrecision,
0363                              trackingGeometry);
0364
0365   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvSimHitWriter, mex,
0366                              "CsvSimHitWriter", inputSimHits, outputDir,
0367                              outputStem, outputPrecision);
0368
0369   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvSpacepointWriter, mex,
0370                              "CsvSpacepointWriter", inputSpacepoints, outputDir,
0371                              outputPrecision);
0372
0373   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvTrackWriter, mex,
0374                              "CsvTrackWriter", inputTracks, outputDir, fileName,
0375                              inputMeasurementParticlesMap, outputPrecision,
0376                              nMeasurementsMin, truthMatchProbMin, ptMin);
0377
0378   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvSeedWriter, mex, "CsvSeedWriter",
0379                              inputTrackParameters, inputSimSeeds, inputSimHits,
0380                              inputMeasurementParticlesMap,
0381                              inputMeasurementSimHitsMap, fileName, outputDir);
0382
0383   ACTS_PYTHON_DECLARE_WRITER(
0384       ActsExamples::CsvTrackingGeometryWriter, mex, "CsvTrackingGeometryWriter",
0385       trackingGeometry, outputDir, outputPrecision, writeSensitive,
0386       writeBoundary, writeSurfaceGrid, writeLayerVolume, writePerEvent);
0387
0388   ACTS_PYTHON_DECLARE_WRITER(
0389       ActsExamples::CKFPerformanceWriter, mex, "CKFPerformanceWriter",
0390       inputTracks, inputParticles, inputTrackParticleMatching,
0391       inputParticleTrackMatching, filePath, fileMode, effPlotToolConfig,
0392       fakeRatePlotToolConfig, duplicationPlotToolConfig,
0393       trackSummaryPlotToolConfig, writeMatchingDetails);
0394
0395   ACTS_PYTHON_DECLARE_WRITER(
0396       ActsExamples::RootNuclearInteractionParametersWriter, mex,
0397       "RootNuclearInteractionParametersWriter", inputSimulationProcesses,
0398       filePath, fileMode, interactionProbabilityBins, momentumBins,
0399       invariantMassBins, multiplicityMax, writeOptionalHistograms,
0400       nSimulatedEvents);
0401
0402   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvTrackParameterWriter, mex,
0403                              "CsvTrackParameterWriter", inputTrackParameters,
0404                              inputTracks, outputDir, outputStem,
0405                              outputPrecision);
0406
0407   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvProtoTrackWriter, mex,
0408                              "CsvProtoTrackWriter", inputSpacepoints,
0409                              inputPrototracks, outputDir);
0410
0411   {
0412     using Writer = ActsExamples::CsvBFieldWriter;
0413
0414     auto w = py::class_<Writer>(mex, "CsvBFieldWriter");
0415
0416     py::enum_<Writer::CoordinateType>(w, "CoordinateType")
0417         .value("rz", Writer::CoordinateType::RZ)
0418         .value("xyz", Writer::CoordinateType::XYZ);
0419
0420     register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, true>(w);
0421     register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, false>(w);
0422     register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, true>(w);
0423     register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, false>(w);
0424   }
0425
0426   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvExaTrkXGraphWriter, mex,
0427                              "CsvExaTrkXGraphWriter", inputGraph, outputDir,
0428                              outputStem);
0429 }
0430
0431 }  // namespace Acts::Python