Python/src/Blueprint.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/Geometry/Blueprint.hpp"
0010
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "Acts/Geometry/BlueprintNode.hpp"
0013 #include "Acts/Geometry/CylinderContainerBlueprintNode.hpp"
0014 #include "Acts/Geometry/CylinderVolumeStack.hpp"
0015 #include "Acts/Geometry/LayerBlueprintNode.hpp"
0016 #include "Acts/Geometry/MaterialDesignatorBlueprintNode.hpp"
0017 #include "Acts/Geometry/StaticBlueprintNode.hpp"
0018 #include "Acts/Navigation/NavigationStream.hpp"
0019 #include "Acts/Plugins/Python/Utilities.hpp"
0020 #include "Acts/Utilities/AxisDefinitions.hpp"
0021 #include "Acts/Utilities/Logger.hpp"
0022
0023 #include <fstream>
0024 #include <random>
0025 #include <utility>
0026
0027 #include <pybind11/pybind11.h>
0028 #include <pybind11/pytypes.h>
0029 #include <pybind11/stl.h>
0030 #include <pybind11/stl/filesystem.h>
0031
0032 namespace py = pybind11;
0033 using namespace pybind11::literals;
0034
0035 namespace Acts::Python {
0036 namespace {
0037 using std::uniform_real_distribution;
0038
0039 // This is temporary!
0040 void pseudoNavigation(const TrackingGeometry& trackingGeometry,
0041                       const GeometryContext& gctx, std::filesystem::path& path,
0042                       std::size_t runs, std::size_t substepsPerCm,
0043                       std::pair<double, double> etaRange,
0044                       Logging::Level logLevel) {
0045   using namespace Acts::UnitLiterals;
0046
0047   ACTS_LOCAL_LOGGER(getDefaultLogger("pseudoNavigation", logLevel));
0048
0049   std::ofstream csv{path};
0050   csv << "x,y,z,volume,boundary,sensitive,material" << std::endl;
0051
0052   std::mt19937 rnd{42};
0053
0054   std::uniform_real_distribution<> dist{-1, 1};
0055   std::uniform_real_distribution<> subStepDist{0.01, 0.99};
0056
0057   double thetaMin = 2 * std::atan(std::exp(-etaRange.first));
0058   double thetaMax = 2 * std::atan(std::exp(-etaRange.second));
0059   std::uniform_real_distribution<> thetaDist{thetaMin, thetaMax};
0060
0061   using namespace Acts::UnitLiterals;
0062
0063   for (std::size_t run = 0; run < runs; run++) {
0064     Vector3 position = Vector3::Zero();
0065
0066     double theta = thetaDist(rnd);
0067     double phi = 2 * std::numbers::pi * dist(rnd);
0068
0069     Vector3 direction;
0070     direction[0] = std::sin(theta) * std::cos(phi);
0071     direction[1] = std::sin(theta) * std::sin(phi);
0072     direction[2] = std::cos(theta);
0073
0074     ACTS_VERBOSE("start navigation " << run);
0075     ACTS_VERBOSE("pos: " << position.transpose());
0076     ACTS_VERBOSE("dir: " << direction.transpose());
0077     ACTS_VERBOSE(direction.norm());
0078
0079     std::mt19937 rng{static_cast<unsigned int>(run)};
0080
0081     const auto* volume = trackingGeometry.lowestTrackingVolume(gctx, position);
0082     assert(volume != nullptr);
0083     ACTS_VERBOSE(volume->volumeName());
0084
0085     NavigationStream main;
0086     const TrackingVolume* currentVolume = volume;
0087
0088     csv << run << "," << position[0] << "," << position[1] << ","
0089         << position[2];
0090     csv << "," << volume->geometryId().volume();
0091     csv << "," << volume->geometryId().boundary();
0092     csv << "," << volume->geometryId().sensitive();
0093     csv << "," << 0;
0094     csv << std::endl;
0095
0096     ACTS_VERBOSE("start pseudo navigation");
0097
0098     auto writeIntersection = [&](const Vector3& pos, const Surface& surface) {
0099       csv << run << "," << pos[0] << "," << pos[1] << "," << pos[2];
0100       csv << "," << surface.geometryId().volume();
0101       csv << "," << surface.geometryId().boundary();
0102       csv << "," << surface.geometryId().sensitive();
0103       csv << "," << (surface.surfaceMaterial() != nullptr ? 1 : 0);
0104       csv << std::endl;
0105     };
0106
0107     for (std::size_t i = 0; i < 100; i++) {
0108       assert(currentVolume != nullptr);
0109       main = NavigationStream{};
0110
0111       AppendOnlyNavigationStream navStream{main};
0112       currentVolume->initializeNavigationCandidates(
0113           {.position = position, .direction = direction}, navStream, logger());
0114
0115       ACTS_VERBOSE(main.candidates().size() << " candidates");
0116
0117       for (const auto& candidate : main.candidates()) {
0118         ACTS_VERBOSE(" -> " << candidate.surface().geometryId());
0119         ACTS_VERBOSE("    " << candidate.surface().toStream(gctx));
0120       }
0121
0122       ACTS_VERBOSE("initializing candidates");
0123       main.initialize(gctx, {position, direction}, BoundaryTolerance::None());
0124
0125       ACTS_VERBOSE(main.candidates().size() << " candidates remaining");
0126
0127       for (const auto& candidate : main.candidates()) {
0128         ACTS_VERBOSE(" -> " << candidate.surface().geometryId());
0129         ACTS_VERBOSE("    " << candidate.surface().toStream(gctx));
0130       }
0131
0132       if (main.currentCandidate().surface().isOnSurface(gctx, position,
0133                                                         direction)) {
0134         ACTS_VERBOSE(
0135             "Already on surface at initialization, skipping candidate");
0136
0137         writeIntersection(position, main.currentCandidate().surface());
0138
0139         if (!main.switchToNextCandidate()) {
0140           ACTS_WARNING("candidates exhausted unexpectedly");
0141           break;
0142         }
0143       }
0144
0145       bool terminated = false;
0146       while (main.remainingCandidates() > 0) {
0147         const auto& candidate = main.currentCandidate();
0148
0149         ACTS_VERBOSE(candidate.portal);
0150         ACTS_VERBOSE(candidate.intersection.position().transpose());
0151
0152         ACTS_VERBOSE("moving to position: " << position.transpose() << " (r="
0153                                             << VectorHelpers::perp(position)
0154                                             << ")");
0155
0156         Vector3 delta = candidate.intersection.position() - position;
0157
0158         std::size_t substeps =
0159             std::max(1l, std::lround(delta.norm() / 10_cm * substepsPerCm));
0160
0161         for (std::size_t j = 0; j < substeps; j++) {
0162           Vector3 subpos = position + subStepDist(rng) * delta;
0163           csv << run << "," << subpos[0] << "," << subpos[1] << ","
0164               << subpos[2];
0165           csv << "," << currentVolume->geometryId().volume();
0166           csv << ",0,0,0";  // zero boundary and sensitive ids
0167           csv << std::endl;
0168         }
0169
0170         position = candidate.intersection.position();
0171         ACTS_VERBOSE("                 -> "
0172                      << position.transpose()
0173                      << " (r=" << VectorHelpers::perp(position) << ")");
0174
0175         writeIntersection(position, candidate.surface());
0176
0177         if (candidate.portal != nullptr) {
0178           ACTS_VERBOSE(
0179               "On portal: " << candidate.portal->surface().toStream(gctx));
0180           currentVolume =
0181               candidate.portal->resolveVolume(gctx, position, direction)
0182                   .value();
0183
0184           if (currentVolume == nullptr) {
0185             ACTS_VERBOSE("switched to nullptr -> we're done");
0186             terminated = true;
0187           }
0188           break;
0189
0190         } else {
0191           ACTS_VERBOSE("Not on portal");
0192         }
0193
0194         main.switchToNextCandidate();
0195       }
0196
0197       if (terminated) {
0198         ACTS_VERBOSE("Terminate pseudo navigation");
0199         break;
0200       }
0201
0202       ACTS_VERBOSE("switched to " << currentVolume->volumeName());
0203
0204       ACTS_VERBOSE("-----");
0205     }
0206   }
0207 }
0208
0209 }  // namespace
0210
0211 void addBlueprint(Context& ctx) {
0212   auto m = ctx.get("main");
0213
0214   auto blueprintNode =
0215       py::class_<BlueprintNode, std::shared_ptr<BlueprintNode>>(
0216           m, "BlueprintNode");
0217
0218   auto rootNode =
0219       py::class_<Blueprint, BlueprintNode, std::shared_ptr<Blueprint>>(
0220           m, "Blueprint");
0221
0222   rootNode
0223       .def(py::init<const Blueprint::Config&>())
0224       // Return value needs to be shared pointer because python otherwise
0225       // can't manage the lifetime
0226       .def(
0227           "construct",
0228           [](Blueprint& self, const BlueprintOptions& options,
0229              const GeometryContext& gctx,
0230              Logging::Level level) -> std::shared_ptr<TrackingGeometry> {
0231             return self.construct(options, gctx,
0232                                   *getDefaultLogger("Blueprint", level));
0233           },
0234           py::arg("options"), py::arg("gctx"),
0235           py::arg("level") = Logging::INFO);
0236
0237   {
0238     auto c = py::class_<Blueprint::Config>(rootNode, "Config").def(py::init());
0239     ACTS_PYTHON_STRUCT_BEGIN(c, Blueprint::Config);
0240     ACTS_PYTHON_MEMBER(envelope);
0241     ACTS_PYTHON_MEMBER(geometryIdentifierHook);
0242     ACTS_PYTHON_STRUCT_END();
0243   }
0244
0245   auto addContextManagerProtocol = []<typename class_>(class_& cls) {
0246     using type = typename class_::type;
0247     cls.def("__enter__", [](type& self) -> type& { return self; })
0248         .def("__exit__", [](type& /*self*/, const py::object& /*exc_type*/,
0249                             const py::object& /*exc_value*/,
0250                             const py::object& /*traceback*/) {
0251           // No action needed on exit
0252         });
0253   };
0254
0255   auto addNodeMethods = [&blueprintNode](const std::string& name,
0256                                          auto&& callable, auto&&... args) {
0257     blueprintNode.def(name.c_str(), callable, args...)
0258         .def(("add" + name).c_str(), callable, args...);
0259   };
0260
0261   blueprintNode
0262       .def("__str__",
0263            [](const BlueprintNode& self) {
0264              std::stringstream ss;
0265              ss << self;
0266              return ss.str();
0267            })
0268       .def("addChild", &BlueprintNode::addChild)
0269       .def_property_readonly("children",
0270                              py::overload_cast<>(&BlueprintNode::children))
0271       .def("clearChildren", &BlueprintNode::clearChildren)
0272       .def_property_readonly("name", &BlueprintNode::name)
0273       .def_property_readonly("depth", &BlueprintNode::depth)
0274       .def("graphviz", [](BlueprintNode& self, const py::object& fh) {
0275         std::stringstream ss;
0276         self.graphviz(ss);
0277         fh.attr("write")(ss.str());
0278       });
0279
0280   py::class_<BlueprintOptions>(m, "BlueprintOptions")
0281       .def(py::init<>())
0282       .def_readwrite("defaultNavigationPolicyFactory",
0283                      &BlueprintOptions::defaultNavigationPolicyFactory);
0284
0285   py::class_<BlueprintNode::MutableChildRange>(blueprintNode,
0286                                                "MutableChildRange")
0287       .def(
0288           "__iter__",
0289           [](BlueprintNode::MutableChildRange& self) {
0290             return py::make_iterator(self.begin(), self.end());
0291           },
0292           py::keep_alive<0, 1>())
0293       .def(
0294           "__getitem__",
0295           [](BlueprintNode::MutableChildRange& self,
0296              int i) -> Acts::BlueprintNode& {
0297             if (i < 0) {
0298               i += self.size();
0299             }
0300             return self.at(i);
0301           },
0302           py::return_value_policy::reference_internal)
0303       .def("__len__", [](const BlueprintNode::MutableChildRange& self) {
0304         return self.size();
0305       });
0306
0307   auto staticNode =
0308       py::class_<Acts::StaticBlueprintNode, Acts::BlueprintNode,
0309                  std::shared_ptr<Acts::StaticBlueprintNode>>(
0310           m, "StaticBlueprintNode")
0311           .def(py::init([](const Transform3& transform,
0312                            const std::shared_ptr<VolumeBounds>& bounds,
0313                            const std::string& name) {
0314                  return std::make_shared<Acts::StaticBlueprintNode>(
0315                      std::make_unique<Acts::TrackingVolume>(transform, bounds,
0316                                                             name));
0317                }),
0318                py::arg("transform"), py::arg("bounds"),
0319                py::arg("name") = "undefined")
0320           .def_property("navigationPolicyFactory",
0321                         &Acts::StaticBlueprintNode::navigationPolicyFactory,
0322                         &Acts::StaticBlueprintNode::setNavigationPolicyFactory);
0323
0324   addContextManagerProtocol(staticNode);
0325
0326   addNodeMethods(
0327       "StaticVolume",
0328       [](BlueprintNode& self, const Transform3& transform,
0329          const std::shared_ptr<VolumeBounds>& bounds, const std::string& name) {
0330         auto node = std::make_shared<Acts::StaticBlueprintNode>(
0331             std::make_unique<Acts::TrackingVolume>(transform, bounds, name));
0332         self.addChild(node);
0333         return node;
0334       },
0335       py::arg("transform"), py::arg("bounds"), py::arg("name") = "undefined");
0336
0337   auto cylNode =
0338       py::class_<Acts::CylinderContainerBlueprintNode, Acts::BlueprintNode,
0339                  std::shared_ptr<Acts::CylinderContainerBlueprintNode>>(
0340           m, "CylinderContainerBlueprintNode")
0341           .def(py::init<const std::string&, AxisDirection,
0342                         CylinderVolumeStack::AttachmentStrategy,
0343                         CylinderVolumeStack::ResizeStrategy>(),
0344                py::arg("name"), py::arg("direction"),
0345                py::arg("attachmentStrategy") =
0346                    CylinderVolumeStack::AttachmentStrategy::Gap,
0347                py::arg("resizeStrategy") =
0348                    CylinderVolumeStack::ResizeStrategy::Gap)
0349           .def_property(
0350               "attachmentStrategy",
0351               &Acts::CylinderContainerBlueprintNode::attachmentStrategy,
0352               &Acts::CylinderContainerBlueprintNode::setAttachmentStrategy)
0353           .def_property(
0354               "resizeStrategy",
0355               &Acts::CylinderContainerBlueprintNode::resizeStrategy,
0356               &Acts::CylinderContainerBlueprintNode::setResizeStrategy)
0357           .def_property("direction",
0358                         &Acts::CylinderContainerBlueprintNode::direction,
0359                         &Acts::CylinderContainerBlueprintNode::setDirection);
0360
0361   addContextManagerProtocol(cylNode);
0362
0363   addNodeMethods(
0364       "CylinderContainer",
0365       [](BlueprintNode& self, const std::string& name,
0366          AxisDirection direction) {
0367         auto cylinder =
0368             std::make_shared<CylinderContainerBlueprintNode>(name, direction);
0369         self.addChild(cylinder);
0370         return cylinder;
0371       },
0372       py::arg("name"), py::arg("direction"));
0373
0374   auto matNode =
0375       py::class_<MaterialDesignatorBlueprintNode, BlueprintNode,
0376                  std::shared_ptr<MaterialDesignatorBlueprintNode>>(
0377           m, "MaterialDesignatorBlueprintNode")
0378           .def(py::init<const std::string&>(), "name"_a)
0379           .def_property("binning", &MaterialDesignatorBlueprintNode::binning,
0380                         &MaterialDesignatorBlueprintNode::setBinning);
0381
0382   addContextManagerProtocol(matNode);
0383
0384   addNodeMethods(
0385       "Material",
0386       [](BlueprintNode& self, const std::string& name) {
0387         auto child = std::make_shared<MaterialDesignatorBlueprintNode>(name);
0388         self.addChild(child);
0389         return child;
0390       },
0391       "name"_a);
0392
0393   auto layerNode =
0394       py::class_<Acts::LayerBlueprintNode, Acts::StaticBlueprintNode,
0395                  std::shared_ptr<Acts::LayerBlueprintNode>>(
0396           m, "LayerBlueprintNode")
0397           .def(py::init<const std::string&>(), py::arg("name"))
0398           .def_property_readonly("name", &Acts::LayerBlueprintNode::name)
0399           .def_property("surfaces", &Acts::LayerBlueprintNode::surfaces,
0400                         &Acts::LayerBlueprintNode::setSurfaces)
0401           .def_property("transform", &Acts::LayerBlueprintNode::transform,
0402                         &Acts::LayerBlueprintNode::setTransform)
0403           .def_property("envelope", &Acts::LayerBlueprintNode::envelope,
0404                         &Acts::LayerBlueprintNode::setEnvelope)
0405           .def_property("layerType", &Acts::LayerBlueprintNode::layerType,
0406                         &Acts::LayerBlueprintNode::setLayerType)
0407           .def_property("navigationPolicyFactory",
0408                         &Acts::LayerBlueprintNode::navigationPolicyFactory,
0409                         &Acts::LayerBlueprintNode::setNavigationPolicyFactory);
0410
0411   py::enum_<Acts::LayerBlueprintNode::LayerType>(layerNode, "LayerType")
0412       .value("Cylinder", Acts::LayerBlueprintNode::LayerType::Cylinder)
0413       .value("Disc", Acts::LayerBlueprintNode::LayerType::Disc)
0414       .value("Plane", Acts::LayerBlueprintNode::LayerType::Plane);
0415
0416   addContextManagerProtocol(layerNode);
0417
0418   addNodeMethods(
0419       "Layer",
0420       [](BlueprintNode& self, const std::string& name) {
0421         auto child = std::make_shared<LayerBlueprintNode>(name);
0422         self.addChild(child);
0423         return child;
0424       },
0425       py::arg("name"));
0426
0427   // TEMPORARY
0428   m.def("pseudoNavigation", &pseudoNavigation, "trackingGeometry"_a, "gctx"_a,
0429         "path"_a, "runs"_a, "substepsPerCm"_a = 2,
0430         "etaRange"_a = std::pair{-4.5, 4.5}, "logLevel"_a = Logging::INFO);
0431 }
0432
0433 }  // namespace Acts::Python