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