Examples/src/Output.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/Utilities/Logger.hpp"
0010 #include "Acts/Visualization/IVisualization3D.hpp"
0011 #include "Acts/Visualization/ViewConfig.hpp"
0012 #include "ActsExamples/Io/Csv/CsvBFieldWriter.hpp"
0013 #include "ActsExamples/Io/Csv/CsvGnnGraphWriter.hpp"
0014 #include "ActsExamples/Io/Csv/CsvMeasurementWriter.hpp"
0015 #include "ActsExamples/Io/Csv/CsvParticleWriter.hpp"
0016 #include "ActsExamples/Io/Csv/CsvProtoTrackWriter.hpp"
0017 #include "ActsExamples/Io/Csv/CsvSeedWriter.hpp"
0018 #include "ActsExamples/Io/Csv/CsvSimHitWriter.hpp"
0019 #include "ActsExamples/Io/Csv/CsvSpacePointWriter.hpp"
0020 #include "ActsExamples/Io/Csv/CsvSpacePointsBucketWriter.hpp"
0021 #include "ActsExamples/Io/Csv/CsvTrackParameterWriter.hpp"
0022 #include "ActsExamples/Io/Csv/CsvTrackWriter.hpp"
0023 #include "ActsExamples/Io/Csv/CsvTrackingGeometryWriter.hpp"
0024 #include "ActsExamples/Io/Csv/CsvVertexWriter.hpp"
0025 #include "ActsExamples/Io/Obj/ObjPropagationStepsWriter.hpp"
0026 #include "ActsExamples/Io/Obj/ObjSimHitWriter.hpp"
0027 #include "ActsExamples/Io/Obj/ObjTrackingGeometryWriter.hpp"
0028 #include "ActsExamples/MaterialMapping/IMaterialWriter.hpp"
0029 #include "ActsExamples/TrackFinding/ITrackParamsLookupWriter.hpp"
0030 #include "ActsPython/Utilities/Helpers.hpp"
0031 #include "ActsPython/Utilities/Macros.hpp"
0032
0033 #include <memory>
0034 #include <string>
0035
0036 #include <pybind11/pybind11.h>
0037 #include <pybind11/stl.h>
0038 #include <pybind11/stl/filesystem.h>
0039
0040 namespace Acts {
0041 class TrackingGeometry;
0042 namespace detail {
0043 struct Step;
0044 }  // namespace detail
0045 }  // namespace Acts
0046 namespace ActsExamples {
0047 class IWriter;
0048 struct AlgorithmContext;
0049 }  // namespace ActsExamples
0050
0051 namespace py = pybind11;
0052 using namespace pybind11::literals;
0053
0054 using namespace Acts;
0055 using namespace ActsExamples;
0056
0057 namespace {
0058 template <CsvBFieldWriter::CoordinateType CType, bool Grid>
0059 void register_csv_bfield_writer_binding(pybind11::class_<CsvBFieldWriter>& w) {
0060   std::string name =
0061       std::string(CType == CsvBFieldWriter::CoordinateType::XYZ ? "Xyz"
0062                                                                 : "Rz") +
0063       std::string(Grid ? "Grid" : "Gridless");
0064
0065   using Config = CsvBFieldWriter::Config<CType, Grid>;
0066   w.def_static((std::string("run") + name).c_str(),
0067                [](const Config& config, Logging::Level level) {
0068                  CsvBFieldWriter::run(config, level);
0069                },
0070                py::arg("config"), py::arg("level"));
0071   auto c = py::class_<Config>(w, (std::string("Config") + name).c_str())
0072                .def(py::init<>());
0073   ACTS_PYTHON_STRUCT(c, fileName, bField, range, bins);
0074 }
0075 }  // namespace
0076
0077 namespace ActsPython {
0078
0079 void addOutput(py::module& mex) {
0080   {
0081     using Writer = ObjTrackingGeometryWriter;
0082     auto w =
0083         py::class_<Writer, std::shared_ptr<Writer>>(mex,
0084                                                     "ObjTrackingGeometryWriter")
0085             .def(py::init<const Writer::Config&, Logging::Level>(),
0086                  py::arg("config"), py::arg("level"))
0087             .def("write",
0088                  py::overload_cast<const AlgorithmContext&,
0089                                    const TrackingGeometry&>(&Writer::write));
0090
0091     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0092
0093     ACTS_PYTHON_STRUCT(c, outputScalor, outputPrecision, outputDir,
0094                        containerView, volumeView, sensitiveView, passiveView,
0095                        gridView);
0096   }
0097
0098   ACTS_PYTHON_DECLARE_WRITER(ObjPropagationStepsWriter, mex,
0099                              "ObjPropagationStepsWriter", collection, outputDir,
0100                              outputScalor, outputPrecision);
0101
0102   ACTS_PYTHON_DECLARE_WRITER(
0103       ObjSimHitWriter, mex, "ObjSimHitWriter", inputSimHits, outputDir,
0104       outputStem, outputPrecision, drawConnections, momentumThreshold,
0105       momentumThresholdTraj, nInterpolatedPoints, keepOriginalHits);
0106
0107   ACTS_PYTHON_DECLARE_WRITER(CsvParticleWriter, mex, "CsvParticleWriter",
0108                              inputParticles, outputDir, outputStem,
0109                              outputPrecision);
0110
0111   ACTS_PYTHON_DECLARE_WRITER(
0112       CsvMeasurementWriter, mex, "CsvMeasurementWriter", inputMeasurements,
0113       inputClusters, inputMeasurementSimHitsMap, outputDir, outputPrecision);
0114
0115   ACTS_PYTHON_DECLARE_WRITER(CsvSimHitWriter, mex, "CsvSimHitWriter",
0116                              inputSimHits, outputDir, outputStem,
0117                              outputPrecision);
0118
0119   ACTS_PYTHON_DECLARE_WRITER(CsvSpacePointWriter, mex, "CsvSpacePointWriter",
0120                              inputSpacepoints, outputDir, outputPrecision);
0121
0122   ACTS_PYTHON_DECLARE_WRITER(CsvSpacePointsBucketWriter, mex,
0123                              "CsvSpacePointsBucketWriter", inputBuckets,
0124                              outputDir, outputPrecision);
0125
0126   ACTS_PYTHON_DECLARE_WRITER(CsvTrackWriter, mex, "CsvTrackWriter", inputTracks,
0127                              outputDir, fileName, inputMeasurementParticlesMap,
0128                              outputPrecision, nMeasurementsMin,
0129                              truthMatchProbMin, ptMin);
0130
0131   ACTS_PYTHON_DECLARE_WRITER(CsvSeedWriter, mex, "CsvSeedWriter",
0132                              inputTrackParameters, inputSimSeeds, inputSimHits,
0133                              inputMeasurementParticlesMap,
0134                              inputMeasurementSimHitsMap, fileName, outputDir);
0135
0136   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvVertexWriter, mex,
0137                              "CsvVertexWriter", inputVertices, outputDir,
0138                              outputStem, outputPrecision);
0139
0140   ACTS_PYTHON_DECLARE_WRITER(
0141       CsvTrackingGeometryWriter, mex, "CsvTrackingGeometryWriter",
0142       trackingGeometry, outputDir, outputPrecision, writeSensitive,
0143       writeBoundary, writeSurfaceGrid, writeLayerVolume, writePerEvent);
0144
0145   ACTS_PYTHON_DECLARE_WRITER(CsvTrackParameterWriter, mex,
0146                              "CsvTrackParameterWriter", inputTracks, outputDir,
0147                              outputStem, outputPrecision);
0148
0149   ACTS_PYTHON_DECLARE_WRITER(CsvProtoTrackWriter, mex, "CsvProtoTrackWriter",
0150                              inputSpacepoints, inputPrototracks, outputDir);
0151
0152   {
0153     using Writer = CsvBFieldWriter;
0154
0155     auto w = py::class_<Writer>(mex, "CsvBFieldWriter");
0156
0157     py::enum_<Writer::CoordinateType>(w, "CoordinateType")
0158         .value("rz", Writer::CoordinateType::RZ)
0159         .value("xyz", Writer::CoordinateType::XYZ);
0160
0161     register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, true>(w);
0162     register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, false>(w);
0163     register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, true>(w);
0164     register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, false>(w);
0165   }
0166
0167   ACTS_PYTHON_DECLARE_WRITER(CsvGnnGraphWriter, mex, "CsvGnnGraphWriter",
0168                              inputGraph, outputDir, outputStem);
0169
0170   py::class_<IMaterialWriter, std::shared_ptr<IMaterialWriter>>(
0171       mex, "IMaterialWriter");
0172
0173   py::class_<ITrackParamsLookupWriter,
0174              std::shared_ptr<ITrackParamsLookupWriter>>(
0175       mex, "ITrackParamsLookupWriter");
0176 }
0177
0178 }  // namespace ActsPython