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/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(Context& ctx) {
0080   auto& mex = ctx.get("examples");
0081
0082   {
0083     using Writer = ObjTrackingGeometryWriter;
0084     auto w =
0085         py::class_<Writer, std::shared_ptr<Writer>>(mex,
0086                                                     "ObjTrackingGeometryWriter")
0087             .def(py::init<const Writer::Config&, Logging::Level>(),
0088                  py::arg("config"), py::arg("level"))
0089             .def("write",
0090                  py::overload_cast<const AlgorithmContext&,
0091                                    const TrackingGeometry&>(&Writer::write));
0092
0093     auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
0094
0095     ACTS_PYTHON_STRUCT(c, outputScalor, outputPrecision, outputDir,
0096                        containerView, volumeView, sensitiveView, passiveView,
0097                        gridView);
0098   }
0099
0100   ACTS_PYTHON_DECLARE_WRITER(ObjPropagationStepsWriter, mex,
0101                              "ObjPropagationStepsWriter", collection, outputDir,
0102                              outputScalor, outputPrecision);
0103
0104   ACTS_PYTHON_DECLARE_WRITER(
0105       ObjSimHitWriter, mex, "ObjSimHitWriter", inputSimHits, outputDir,
0106       outputStem, outputPrecision, drawConnections, momentumThreshold,
0107       momentumThresholdTraj, nInterpolatedPoints, keepOriginalHits);
0108
0109   ACTS_PYTHON_DECLARE_WRITER(CsvParticleWriter, mex, "CsvParticleWriter",
0110                              inputParticles, outputDir, outputStem,
0111                              outputPrecision);
0112
0113   ACTS_PYTHON_DECLARE_WRITER(
0114       CsvMeasurementWriter, mex, "CsvMeasurementWriter", inputMeasurements,
0115       inputClusters, inputMeasurementSimHitsMap, outputDir, outputPrecision);
0116
0117   ACTS_PYTHON_DECLARE_WRITER(CsvSimHitWriter, mex, "CsvSimHitWriter",
0118                              inputSimHits, outputDir, outputStem,
0119                              outputPrecision);
0120
0121   ACTS_PYTHON_DECLARE_WRITER(CsvSpacePointWriter, mex, "CsvSpacePointWriter",
0122                              inputSpacepoints, outputDir, outputPrecision);
0123
0124   ACTS_PYTHON_DECLARE_WRITER(CsvSpacePointsBucketWriter, mex,
0125                              "CsvSpacePointsBucketWriter", inputBuckets,
0126                              outputDir, outputPrecision);
0127
0128   ACTS_PYTHON_DECLARE_WRITER(CsvTrackWriter, mex, "CsvTrackWriter", inputTracks,
0129                              outputDir, fileName, inputMeasurementParticlesMap,
0130                              outputPrecision, nMeasurementsMin,
0131                              truthMatchProbMin, ptMin);
0132
0133   ACTS_PYTHON_DECLARE_WRITER(CsvSeedWriter, mex, "CsvSeedWriter",
0134                              inputTrackParameters, inputSimSeeds, inputSimHits,
0135                              inputMeasurementParticlesMap,
0136                              inputMeasurementSimHitsMap, fileName, outputDir);
0137
0138   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvVertexWriter, mex,
0139                              "CsvVertexWriter", inputVertices, outputDir,
0140                              outputStem, outputPrecision);
0141
0142   ACTS_PYTHON_DECLARE_WRITER(
0143       CsvTrackingGeometryWriter, mex, "CsvTrackingGeometryWriter",
0144       trackingGeometry, outputDir, outputPrecision, writeSensitive,
0145       writeBoundary, writeSurfaceGrid, writeLayerVolume, writePerEvent);
0146
0147   ACTS_PYTHON_DECLARE_WRITER(CsvTrackParameterWriter, mex,
0148                              "CsvTrackParameterWriter", inputTracks, outputDir,
0149                              outputStem, outputPrecision);
0150
0151   ACTS_PYTHON_DECLARE_WRITER(CsvProtoTrackWriter, mex, "CsvProtoTrackWriter",
0152                              inputSpacepoints, inputPrototracks, outputDir);
0153
0154   {
0155     using Writer = CsvBFieldWriter;
0156
0157     auto w = py::class_<Writer>(mex, "CsvBFieldWriter");
0158
0159     py::enum_<Writer::CoordinateType>(w, "CoordinateType")
0160         .value("rz", Writer::CoordinateType::RZ)
0161         .value("xyz", Writer::CoordinateType::XYZ);
0162
0163     register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, true>(w);
0164     register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, false>(w);
0165     register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, true>(w);
0166     register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, false>(w);
0167   }
0168
0169   ACTS_PYTHON_DECLARE_WRITER(CsvGnnGraphWriter, mex, "CsvGnnGraphWriter",
0170                              inputGraph, outputDir, outputStem);
0171
0172   py::class_<IMaterialWriter, std::shared_ptr<IMaterialWriter>>(
0173       mex, "IMaterialWriter");
0174
0175   py::class_<ITrackParamsLookupWriter,
0176              std::shared_ptr<ITrackParamsLookupWriter>>(
0177       mex, "ITrackParamsLookupWriter");
0178 }
0179
0180 }  // namespace ActsPython