Python/src/MagneticField.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 "ActsExamples/MagneticField/MagneticField.hpp"
0010
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "Acts/MagneticField/BFieldMapUtils.hpp"
0013 #include "Acts/MagneticField/ConstantBField.hpp"
0014 #include "Acts/MagneticField/MagneticFieldProvider.hpp"
0015 #include "Acts/MagneticField/MultiRangeBField.hpp"
0016 #include "Acts/MagneticField/NullBField.hpp"
0017 #include "Acts/MagneticField/SolenoidBField.hpp"
0018 #include "Acts/Plugins/Python/Utilities.hpp"
0019 #include "ActsExamples/MagneticField/FieldMapRootIo.hpp"
0020 #include "ActsExamples/MagneticField/FieldMapTextIo.hpp"
0021
0022 #include <array>
0023 #include <cstddef>
0024 #include <filesystem>
0025 #include <memory>
0026 #include <stdexcept>
0027 #include <string>
0028 #include <tuple>
0029 #include <utility>
0030
0031 #include <pybind11/pybind11.h>
0032 #include <pybind11/stl.h>
0033
0034 namespace py = pybind11;
0035 using namespace pybind11::literals;
0036
0037 namespace Acts::Python {
0038
0039 /// @brief Get the value of a field, throwing an exception if the result is
0040 /// invalid.
0041 Acts::Vector3 getField(Acts::MagneticFieldProvider& self,
0042                        const Acts::Vector3& position,
0043                        Acts::MagneticFieldProvider::Cache& cache) {
0044   if (Result<Vector3> res = self.getField(position, cache); !res.ok()) {
0045     std::stringstream ss;
0046
0047     ss << "Field lookup failure with error: \"" << res.error() << "\"";
0048
0049     throw std::runtime_error{ss.str()};
0050   } else {
0051     return *res;
0052   }
0053 }
0054
0055 void addMagneticField(Context& ctx) {
0056   auto [m, mex, prop] = ctx.get("main", "examples", "propagation");
0057
0058   py::class_<Acts::MagneticFieldProvider,
0059              std::shared_ptr<Acts::MagneticFieldProvider>>(
0060       m, "MagneticFieldProvider")
0061       .def("getField", &getField)
0062       .def("makeCache", &Acts::MagneticFieldProvider::makeCache);
0063
0064   py::class_<Acts::InterpolatedMagneticField,
0065              std::shared_ptr<Acts::InterpolatedMagneticField>>(
0066       m, "InterpolatedMagneticField");
0067
0068   m.def("solenoidFieldMap", &Acts::solenoidFieldMap, py::arg("rlim"),
0069         py::arg("zlim"), py::arg("nbins"), py::arg("field"));
0070
0071   py::class_<Acts::ConstantBField, Acts::MagneticFieldProvider,
0072              std::shared_ptr<Acts::ConstantBField>>(m, "ConstantBField")
0073       .def(py::init<Acts::Vector3>());
0074
0075   py::class_<ActsExamples::detail::InterpolatedMagneticField2,
0076              Acts::InterpolatedMagneticField, Acts::MagneticFieldProvider,
0077              std::shared_ptr<ActsExamples::detail::InterpolatedMagneticField2>>(
0078       mex, "InterpolatedMagneticField2");
0079
0080   py::class_<ActsExamples::detail::InterpolatedMagneticField3,
0081              Acts::InterpolatedMagneticField, Acts::MagneticFieldProvider,
0082              std::shared_ptr<ActsExamples::detail::InterpolatedMagneticField3>>(
0083       mex, "InterpolatedMagneticField3");
0084
0085   py::class_<Acts::NullBField, Acts::MagneticFieldProvider,
0086              std::shared_ptr<Acts::NullBField>>(m, "NullBField")
0087       .def(py::init<>());
0088
0089   py::class_<Acts::MultiRangeBField, Acts::MagneticFieldProvider,
0090              std::shared_ptr<Acts::MultiRangeBField>>(m, "MultiRangeBField")
0091       .def(py::init<
0092            std::vector<std::pair<Acts::RangeXD<3, double>, Acts::Vector3>>>());
0093
0094   {
0095     using Config = Acts::SolenoidBField::Config;
0096
0097     auto sol =
0098         py::class_<Acts::SolenoidBField, Acts::MagneticFieldProvider,
0099                    std::shared_ptr<Acts::SolenoidBField>>(m, "SolenoidBField")
0100             .def(py::init<Config>())
0101             .def(py::init([](double radius, double length, std::size_t nCoils,
0102                              double bMagCenter) {
0103                    return Acts::SolenoidBField{
0104                        Config{radius, length, nCoils, bMagCenter}};
0105                  }),
0106                  py::arg("radius"), py::arg("length"), py::arg("nCoils"),
0107                  py::arg("bMagCenter"));
0108
0109     py::class_<Config>(sol, "Config")
0110         .def(py::init<>())
0111         .def_readwrite("radius", &Config::radius)
0112         .def_readwrite("length", &Config::length)
0113         .def_readwrite("nCoils", &Config::nCoils)
0114         .def_readwrite("bMagCenter", &Config::bMagCenter);
0115   }
0116
0117   mex.def(
0118       "MagneticFieldMapXyz",
0119       [](const std::string& filename, const std::string& tree,
0120          double lengthUnit, double BFieldUnit, bool firstOctant) {
0121         const std::filesystem::path file = filename;
0122
0123         auto mapBins = [](std::array<std::size_t, 3> bins,
0124                           std::array<std::size_t, 3> sizes) {
0125           return (bins[0] * (sizes[1] * sizes[2]) + bins[1] * sizes[2] +
0126                   bins[2]);
0127         };
0128
0129         if (file.extension() == ".root") {
0130           auto map = ActsExamples::makeMagneticFieldMapXyzFromRoot(
0131               std::move(mapBins), file.native(), tree, lengthUnit, BFieldUnit,
0132               firstOctant);
0133           return std::make_shared<
0134               ActsExamples::detail::InterpolatedMagneticField3>(std::move(map));
0135         } else if (file.extension() == ".txt") {
0136           auto map = ActsExamples::makeMagneticFieldMapXyzFromText(
0137               std::move(mapBins), file.native(), lengthUnit, BFieldUnit,
0138               firstOctant);
0139           return std::make_shared<
0140               ActsExamples::detail::InterpolatedMagneticField3>(std::move(map));
0141         } else {
0142           throw std::runtime_error("Unsupported magnetic field map file type");
0143         }
0144       },
0145       py::arg("file"), py::arg("tree") = "bField",
0146       py::arg("lengthUnit") = Acts::UnitConstants::mm,
0147       py::arg("BFieldUnit") = Acts::UnitConstants::T,
0148       py::arg("firstOctant") = false);
0149
0150   mex.def(
0151       "MagneticFieldMapRz",
0152       [](const std::string& filename, const std::string& tree,
0153          double lengthUnit, double BFieldUnit, bool firstQuadrant) {
0154         const std::filesystem::path file = filename;
0155
0156         auto mapBins = [](std::array<std::size_t, 2> bins,
0157                           std::array<std::size_t, 2> sizes) {
0158           return (bins[1] * sizes[0] + bins[0]);
0159         };
0160
0161         if (file.extension() == ".root") {
0162           auto map = ActsExamples::makeMagneticFieldMapRzFromRoot(
0163               std::move(mapBins), file.native(), tree, lengthUnit, BFieldUnit,
0164               firstQuadrant);
0165           return std::make_shared<
0166               ActsExamples::detail::InterpolatedMagneticField2>(std::move(map));
0167         } else if (file.extension() == ".txt") {
0168           auto map = ActsExamples::makeMagneticFieldMapRzFromText(
0169               std::move(mapBins), file.native(), lengthUnit, BFieldUnit,
0170               firstQuadrant);
0171           return std::make_shared<
0172               ActsExamples::detail::InterpolatedMagneticField2>(std::move(map));
0173         } else {
0174           throw std::runtime_error("Unsupported magnetic field map file type");
0175         }
0176       },
0177       py::arg("file"), py::arg("tree") = "bField",
0178       py::arg("lengthUnit") = Acts::UnitConstants::mm,
0179       py::arg("BFieldUnit") = Acts::UnitConstants::T,
0180       py::arg("firstQuadrant") = false);
0181 }
0182
0183 }  // namespace Acts::Python