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