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