Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
Output.cpp
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file Output.cpp
1 // This file is part of the Acts project.
2 //
3 // Copyright (C) 2021-2023 CERN for the benefit of the Acts project
4 //
5 // This Source Code Form is subject to the terms of the Mozilla Public
6 // License, v. 2.0. If a copy of the MPL was not distributed with this
7 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
8 
9 #include "Acts/Definitions/TrackParametrization.hpp"
10 #include "Acts/Geometry/GeometryHierarchyMap.hpp"
11 #include "Acts/Plugins/Python/Utilities.hpp"
12 #include "Acts/Utilities/Logger.hpp"
13 #include "Acts/Visualization/ViewConfig.hpp"
14 #include "ActsExamples/Digitization/DigitizationConfig.hpp"
15 #include "ActsExamples/Framework/ProcessCode.hpp"
16 #include "ActsExamples/Io/Csv/CsvBFieldWriter.hpp"
17 #include "ActsExamples/Io/Csv/CsvMeasurementWriter.hpp"
18 #include "ActsExamples/Io/Csv/CsvMultiTrajectoryWriter.hpp"
19 #include "ActsExamples/Io/Csv/CsvParticleWriter.hpp"
20 #include "ActsExamples/Io/Csv/CsvPlanarClusterWriter.hpp"
21 #include "ActsExamples/Io/Csv/CsvProtoTrackWriter.hpp"
22 #include "ActsExamples/Io/Csv/CsvSimHitWriter.hpp"
23 #include "ActsExamples/Io/Csv/CsvSpacepointWriter.hpp"
24 #include "ActsExamples/Io/Csv/CsvTrackParameterWriter.hpp"
25 #include "ActsExamples/Io/Csv/CsvTrackingGeometryWriter.hpp"
26 #include "ActsExamples/Io/NuclearInteractions/RootNuclearInteractionParametersWriter.hpp"
27 #include "ActsExamples/Io/Performance/CKFPerformanceWriter.hpp"
28 #include "ActsExamples/Io/Performance/SeedingPerformanceWriter.hpp"
29 #include "ActsExamples/Io/Performance/TrackFinderPerformanceWriter.hpp"
30 #include "ActsExamples/Io/Performance/TrackFitterPerformanceWriter.hpp"
31 #include "ActsExamples/Io/Performance/VertexPerformanceWriter.hpp"
32 #include "ActsExamples/Io/Root/RootBFieldWriter.hpp"
33 #include "ActsExamples/Io/Root/RootMaterialTrackWriter.hpp"
34 #include "ActsExamples/Io/Root/RootMaterialWriter.hpp"
35 #include "ActsExamples/Io/Root/RootMeasurementWriter.hpp"
36 #include "ActsExamples/Io/Root/RootParticleWriter.hpp"
37 #include "ActsExamples/Io/Root/RootPlanarClusterWriter.hpp"
38 #include "ActsExamples/Io/Root/RootPropagationStepsWriter.hpp"
39 #include "ActsExamples/Io/Root/RootSimHitWriter.hpp"
40 #include "ActsExamples/Io/Root/RootSpacepointWriter.hpp"
41 #include "ActsExamples/Io/Root/RootTrackParameterWriter.hpp"
42 #include "ActsExamples/Io/Root/RootTrajectoryStatesWriter.hpp"
43 #include "ActsExamples/Io/Root/RootTrajectorySummaryWriter.hpp"
44 #include "ActsExamples/MaterialMapping/IMaterialWriter.hpp"
45 #include "ActsExamples/Plugins/Obj/ObjPropagationStepsWriter.hpp"
46 #include "ActsExamples/Plugins/Obj/ObjTrackingGeometryWriter.hpp"
47 
48 #include <array>
49 #include <memory>
50 #include <string>
51 #include <tuple>
52 #include <vector>
53 
54 #include <pybind11/pybind11.h>
55 #include <pybind11/stl.h>
56 
57 namespace Acts {
58 class TrackingGeometry;
59 namespace detail {
60 struct Step;
61 } // namespace detail
62 } // namespace Acts
63 namespace ActsExamples {
64 class IWriter;
65 struct AlgorithmContext;
66 } // namespace ActsExamples
67 
68 namespace py = pybind11;
69 using namespace pybind11::literals;
70 
71 using namespace ActsExamples;
72 
73 namespace {
74 template <ActsExamples::CsvBFieldWriter::CoordinateType CType, bool Grid>
75 void register_csv_bfield_writer_binding(
76  pybind11::class_<ActsExamples::CsvBFieldWriter>& w) {
77  std::string name =
78  std::string(CType == ActsExamples::CsvBFieldWriter::CoordinateType::XYZ
79  ? "Xyz"
80  : "Rz") +
81  std::string(Grid ? "Grid" : "Gridless");
82 
83  using Config = ActsExamples::CsvBFieldWriter::Config<CType, Grid>;
84  w.def_static((std::string("run") + name).c_str(),
85  [](const Config& config, Acts::Logging::Level level) {
86  ActsExamples::CsvBFieldWriter::run(config, level);
87  },
88  py::arg("config"), py::arg("level"));
89  auto c = py::class_<Config>(w, (std::string("Config") + name).c_str())
90  .def(py::init<>());
91  ACTS_PYTHON_STRUCT_BEGIN(c, Config);
92  ACTS_PYTHON_MEMBER(fileName);
93  ACTS_PYTHON_MEMBER(bField);
94  ACTS_PYTHON_MEMBER(range);
95  ACTS_PYTHON_MEMBER(bins);
96  ACTS_PYTHON_STRUCT_END();
97 }
98 } // namespace
99 
100 namespace Acts::Python {
101 void addOutput(Context& ctx) {
102  auto [m, mex] = ctx.get("main", "examples");
103 
104  ACTS_PYTHON_DECLARE_WRITER(
105  ActsExamples::ObjPropagationStepsWriter<Acts::detail::Step>, mex,
106  "ObjPropagationStepsWriter", collection, outputDir, outputScalor,
107  outputPrecision);
108 
109  {
110  auto c = py::class_<ViewConfig>(m, "ViewConfig").def(py::init<>());
111 
112  ACTS_PYTHON_STRUCT_BEGIN(c, ViewConfig);
113  ACTS_PYTHON_MEMBER(visible);
114  ACTS_PYTHON_MEMBER(color);
115  ACTS_PYTHON_MEMBER(offset);
116  ACTS_PYTHON_MEMBER(lineThickness);
117  ACTS_PYTHON_MEMBER(surfaceThickness);
118  ACTS_PYTHON_MEMBER(nSegments);
119  ACTS_PYTHON_MEMBER(triangulate);
120  ACTS_PYTHON_MEMBER(outputName);
121  ACTS_PYTHON_STRUCT_END();
122 
123  patchKwargsConstructor(c);
124  }
125 
126  {
127  using Writer = ActsExamples::ObjTrackingGeometryWriter;
128  auto w = py::class_<Writer, std::shared_ptr<Writer>>(
129  mex, "ObjTrackingGeometryWriter")
130  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
131  py::arg("config"), py::arg("level"))
132  .def("write", py::overload_cast<const AlgorithmContext&,
133  const Acts::TrackingGeometry&>(
134  &Writer::write));
135 
136  auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
137  ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
138  ACTS_PYTHON_MEMBER(outputScalor);
139  ACTS_PYTHON_MEMBER(outputPrecision);
140  ACTS_PYTHON_MEMBER(outputDir);
141  ACTS_PYTHON_MEMBER(containerView);
142  ACTS_PYTHON_MEMBER(volumeView);
143  ACTS_PYTHON_MEMBER(sensitiveView);
144  ACTS_PYTHON_MEMBER(passiveView);
145  ACTS_PYTHON_MEMBER(gridView);
146  ACTS_PYTHON_STRUCT_END();
147  }
148 
149  // Bindings for the binning in e.g., CKFPerformanceWriter
150  {
151  py::class_<PlotHelpers::Binning>(mex, "Binning")
152  .def(py::init<std::string, int, double, double>(), "title"_a, "bins"_a,
153  "bMin"_a, "bMax"_a)
154  .def(py::init<std::string, std::vector<double>>(), "title"_a, "bins"_a);
155 
156  py::class_<EffPlotTool::Config>(mex, "EffPlotToolConfig")
157  .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
158  "varBinning"_a);
159 
160  py::class_<FakeRatePlotTool::Config>(mex, "FakeRatePlotToolConfig")
161  .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
162  "varBinning"_a);
163 
164  py::class_<DuplicationPlotTool::Config>(mex, "DuplicationPlotToolConfig")
165  .def(py::init<std::map<std::string, PlotHelpers::Binning>>(),
166  "varBinning"_a);
167  }
168 
169  // ROOT WRITERS
170  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootPropagationStepsWriter, mex,
171  "RootPropagationStepsWriter", collection, filePath,
172  fileMode);
173 
174  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootParticleWriter, mex,
175  "RootParticleWriter", inputParticles, filePath,
176  fileMode, treeName);
177 
178  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::TrackFinderPerformanceWriter, mex,
179  "TrackFinderPerformanceWriter", inputProtoTracks,
180  inputMeasurementParticlesMap, inputParticles,
181  filePath, fileMode, treeNameTracks,
182  treeNameParticles);
183 
184  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::TrackFitterPerformanceWriter, mex,
185  "TrackFitterPerformanceWriter", inputTrajectories,
186  inputParticles, inputMeasurementParticlesMap,
187  filePath, resPlotToolConfig, effPlotToolConfig,
188  trackSummaryPlotToolConfig);
189 
190  ACTS_PYTHON_DECLARE_WRITER(
191  ActsExamples::SeedingPerformanceWriter, mex, "SeedingPerformanceWriter",
192  inputSeeds, inputMeasurementParticlesMap, inputParticles, filePath,
193  fileMode, effPlotToolConfig, duplicationPlotToolConfig);
194 
195  ACTS_PYTHON_DECLARE_WRITER(
196  ActsExamples::RootTrackParameterWriter, mex, "RootTrackParameterWriter",
197  inputTrackParameters, inputProtoTracks, inputParticles, inputSimHits,
198  inputMeasurementParticlesMap, inputMeasurementSimHitsMap, filePath,
199  treeName, fileMode);
200 
201  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootMaterialTrackWriter, mex,
202  "RootMaterialTrackWriter", collection, filePath,
203  fileMode, treeName, recalculateTotals, prePostStep,
204  storeSurface, storeVolume, collapseInteractions);
205 
206  {
207  using Writer = ActsExamples::RootBFieldWriter;
208  auto w =
209  py::class_<Writer>(mex, "RootBFieldWriter")
210  .def_static(
211  "run",
212  [](const Writer::Config& config, Acts::Logging::Level level) {
213  Writer::run(config, Acts::getDefaultLogger("RootBFieldWriter",
214  level));
215  },
216  py::arg("config"), py::arg("level"));
217 
218  py::enum_<Writer::GridType>(w, "GridType")
219  .value("rz", Writer::GridType::rz)
220  .value("xyz", Writer::GridType::xyz);
221 
222  auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
223  ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
224  ACTS_PYTHON_MEMBER(treeName);
225  ACTS_PYTHON_MEMBER(fileName);
226  ACTS_PYTHON_MEMBER(fileMode);
227  ACTS_PYTHON_MEMBER(bField);
228  ACTS_PYTHON_MEMBER(gridType);
229  ACTS_PYTHON_MEMBER(rBounds);
230  ACTS_PYTHON_MEMBER(zBounds);
231  ACTS_PYTHON_MEMBER(rBins);
232  ACTS_PYTHON_MEMBER(zBins);
233  ACTS_PYTHON_MEMBER(phiBins);
234  ACTS_PYTHON_STRUCT_END();
235  }
236 
237  {
238  using Writer = ActsExamples::RootMeasurementWriter;
239  auto w = py::class_<Writer, IWriter, std::shared_ptr<Writer>>(
240  mex, "RootMeasurementWriter")
241  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
242  py::arg("config"), py::arg("level"));
243 
244  auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
245 
246  c.def("addBoundIndicesFromDigiConfig",
247  [](Writer::Config& self, const DigitizationConfig& digiCfg) {
248  self.boundIndices =
249  Acts::GeometryHierarchyMap<std::vector<Acts::BoundIndices>>(
250  digiCfg.getBoundIndices());
251  });
252 
253  ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
254  ACTS_PYTHON_MEMBER(inputMeasurements);
255  ACTS_PYTHON_MEMBER(inputClusters);
256  ACTS_PYTHON_MEMBER(inputSimHits);
257  ACTS_PYTHON_MEMBER(inputMeasurementSimHitsMap);
258  ACTS_PYTHON_MEMBER(filePath);
259  ACTS_PYTHON_MEMBER(fileMode);
260  ACTS_PYTHON_MEMBER(boundIndices);
261  ACTS_PYTHON_MEMBER(trackingGeometry);
262  ACTS_PYTHON_STRUCT_END();
263  }
264 
265  py::class_<IMaterialWriter, std::shared_ptr<IMaterialWriter>>(
266  mex, "IMaterialWriter");
267 
268  {
269  using Writer = ActsExamples::RootMaterialWriter;
270  auto w = py::class_<Writer, IMaterialWriter, std::shared_ptr<Writer>>(
271  mex, "RootMaterialWriter")
272  .def(py::init<const Writer::Config&, Acts::Logging::Level>(),
273  py::arg("config"), py::arg("level"))
274  .def("write", py::overload_cast<const Acts::TrackingGeometry&>(
275  &Writer::write));
276 
277  auto c = py::class_<Writer::Config>(w, "Config").def(py::init<>());
278 
279  ACTS_PYTHON_STRUCT_BEGIN(c, Writer::Config);
280  ACTS_PYTHON_MEMBER(processSensitives);
281  ACTS_PYTHON_MEMBER(processApproaches);
282  ACTS_PYTHON_MEMBER(processRepresenting);
283  ACTS_PYTHON_MEMBER(processBoundaries);
284  ACTS_PYTHON_MEMBER(processVolumes);
285  ACTS_PYTHON_MEMBER(folderSurfaceNameBase);
286  ACTS_PYTHON_MEMBER(folderVolumeNameBase);
287  ACTS_PYTHON_MEMBER(voltag);
288  ACTS_PYTHON_MEMBER(boutag);
289  ACTS_PYTHON_MEMBER(laytag);
290  ACTS_PYTHON_MEMBER(apptag);
291  ACTS_PYTHON_MEMBER(sentag);
292  ACTS_PYTHON_MEMBER(ntag);
293  ACTS_PYTHON_MEMBER(vtag);
294  ACTS_PYTHON_MEMBER(otag);
295  ACTS_PYTHON_MEMBER(mintag);
296  ACTS_PYTHON_MEMBER(maxtag);
297  ACTS_PYTHON_MEMBER(ttag);
298  ACTS_PYTHON_MEMBER(x0tag);
299  ACTS_PYTHON_MEMBER(l0tag);
300  ACTS_PYTHON_MEMBER(atag);
301  ACTS_PYTHON_MEMBER(ztag);
302  ACTS_PYTHON_MEMBER(rhotag);
303  ACTS_PYTHON_MEMBER(filePath);
304  ACTS_PYTHON_MEMBER(fileMode);
305  ACTS_PYTHON_STRUCT_END();
306  }
307 
308  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootPlanarClusterWriter, mex,
309  "RootPlanarClusterWriter", inputClusters,
310  inputSimHits, filePath, fileMode, treeName,
311  trackingGeometry);
312 
313  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSimHitWriter, mex,
314  "RootSimHitWriter", inputSimHits, filePath,
315  fileMode, treeName);
316 
317  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootSpacepointWriter, mex,
318  "RootSpacepointWriter", inputSpacepoints, filePath,
319  fileMode, treeName);
320 
321  ACTS_PYTHON_DECLARE_WRITER(
322  ActsExamples::RootTrajectoryStatesWriter, mex,
323  "RootTrajectoryStatesWriter", inputTrajectories, inputParticles,
324  inputSimHits, inputMeasurementParticlesMap, inputMeasurementSimHitsMap,
325  filePath, treeName, fileMode);
326 
327  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootTrajectorySummaryWriter, mex,
328  "RootTrajectorySummaryWriter", inputTrajectories,
329  inputParticles, inputMeasurementParticlesMap,
330  filePath, treeName, fileMode, writeCovMat);
331 
332  ACTS_PYTHON_DECLARE_WRITER(
333  ActsExamples::VertexPerformanceWriter, mex, "VertexPerformanceWriter",
334  inputAllTruthParticles, inputSelectedTruthParticles, inputTrackParameters,
335  inputAssociatedTruthParticles, inputTrackParameters, inputTrajectories,
336  inputMeasurementParticlesMap, inputVertices, bField, filePath, treeName,
337  fileMode, minTrackVtxMatchFraction, truthMatchProbMin, useTracks);
338 
339  // CSV WRITERS
340  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvParticleWriter, mex,
341  "CsvParticleWriter", inputParticles, outputDir,
342  outputStem, outputPrecision);
343 
344  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvMeasurementWriter, mex,
345  "CsvMeasurementWriter", inputMeasurements,
346  inputClusters, inputMeasurementSimHitsMap,
347  outputDir, outputPrecision);
348 
349  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvPlanarClusterWriter, mex,
350  "CsvPlanarClusterWriter", inputClusters,
351  inputSimHits, outputDir, outputPrecision,
352  trackingGeometry);
353 
354  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvSimHitWriter, mex,
355  "CsvSimHitWriter", inputSimHits, outputDir,
356  outputStem, outputPrecision);
357 
358  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvSpacepointWriter, mex,
359  "CsvSpacepointWriter", inputSpacepoints, outputDir,
360  outputPrecision);
361 
362  ACTS_PYTHON_DECLARE_WRITER(
363  ActsExamples::CsvMultiTrajectoryWriter, mex, "CsvMultiTrajectoryWriter",
364  inputTrajectories, outputDir, fileName, inputMeasurementParticlesMap,
365  outputPrecision, nMeasurementsMin, truthMatchProbMin, ptMin);
366 
367  ACTS_PYTHON_DECLARE_WRITER(
368  ActsExamples::CsvTrackingGeometryWriter, mex, "CsvTrackingGeometryWriter",
369  trackingGeometry, outputDir, outputPrecision, writeSensitive,
370  writeBoundary, writeSurfaceGrid, writeLayerVolume, writePerEvent);
371 
372  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CKFPerformanceWriter, mex,
373  "CKFPerformanceWriter", inputTrajectories,
374  inputParticles, inputMeasurementParticlesMap,
375  filePath, fileMode, effPlotToolConfig,
376  fakeRatePlotToolConfig, duplicationPlotToolConfig,
377  trackSummaryPlotToolConfig, duplicatedPredictor);
378 
379  ACTS_PYTHON_DECLARE_WRITER(
380  ActsExamples::RootNuclearInteractionParametersWriter, mex,
381  "RootNuclearInteractionParametersWriter", inputSimulationProcesses,
382  filePath, fileMode, interactionProbabilityBins, momentumBins,
383  invariantMassBins, multiplicityMax, writeOptionalHistograms,
384  nSimulatedEvents);
385 
386  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvTrackParameterWriter, mex,
387  "CsvTrackParameterWriter", inputTrackParameters,
388  inputTrajectories, outputDir, outputStem,
389  outputPrecision);
390 
391  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::CsvProtoTrackWriter, mex,
392  "CsvProtoTrackWriter", inputSpacepoints,
393  inputPrototracks, outputDir);
394 
395  {
396  using Writer = ActsExamples::CsvBFieldWriter;
397 
398  auto w = py::class_<Writer>(mex, "CsvBFieldWriter");
399 
400  py::enum_<Writer::CoordinateType>(w, "CoordinateType")
401  .value("rz", Writer::CoordinateType::RZ)
402  .value("xyz", Writer::CoordinateType::XYZ);
403 
404  register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, true>(w);
405  register_csv_bfield_writer_binding<Writer::CoordinateType::XYZ, false>(w);
406  register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, true>(w);
407  register_csv_bfield_writer_binding<Writer::CoordinateType::RZ, false>(w);
408  }
409 }
410 } // namespace Acts::Python