Python/src/MagneticField.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2021 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 http://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/NullBField.hpp"
0016 #include "Acts/MagneticField/SolenoidBField.hpp"
0017 #include "Acts/Plugins/Python/Utilities.hpp"
0018 #include "ActsExamples/MagneticField/FieldMapRootIo.hpp"
0019 #include "ActsExamples/MagneticField/FieldMapTextIo.hpp"
0020
0021 #include <array>
0022 #include <cstddef>
0023 #include <filesystem>
0024 #include <memory>
0025 #include <stdexcept>
0026 #include <string>
0027 #include <tuple>
0028 #include <type_traits>
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 void addMagneticField(Context& ctx) {
0040   auto [m, mex, prop] = ctx.get("main", "examples", "propagation");
0041
0042   py::class_<Acts::MagneticFieldProvider,
0043              std::shared_ptr<Acts::MagneticFieldProvider>>(
0044       m, "MagneticFieldProvider");
0045
0046   py::class_<Acts::InterpolatedMagneticField,
0047              std::shared_ptr<Acts::InterpolatedMagneticField>>(
0048       m, "InterpolatedMagneticField");
0049
0050   m.def("solenoidFieldMap", &Acts::solenoidFieldMap, py::arg("rlim"),
0051         py::arg("zlim"), py::arg("nbins"), py::arg("field"));
0052
0053   py::class_<Acts::ConstantBField, Acts::MagneticFieldProvider,
0054              std::shared_ptr<Acts::ConstantBField>>(m, "ConstantBField")
0055       .def(py::init<Acts::Vector3>());
0056
0057   py::class_<ActsExamples::detail::InterpolatedMagneticField2,
0058              Acts::InterpolatedMagneticField, Acts::MagneticFieldProvider,
0059              std::shared_ptr<ActsExamples::detail::InterpolatedMagneticField2>>(
0060       mex, "InterpolatedMagneticField2");
0061
0062   py::class_<ActsExamples::detail::InterpolatedMagneticField3,
0063              Acts::InterpolatedMagneticField, Acts::MagneticFieldProvider,
0064              std::shared_ptr<ActsExamples::detail::InterpolatedMagneticField3>>(
0065       mex, "InterpolatedMagneticField3");
0066
0067   py::class_<Acts::NullBField, Acts::MagneticFieldProvider,
0068              std::shared_ptr<Acts::NullBField>>(m, "NullBField")
0069       .def(py::init<>());
0070
0071   {
0072     using Config = Acts::SolenoidBField::Config;
0073
0074     auto sol =
0075         py::class_<Acts::SolenoidBField, Acts::MagneticFieldProvider,
0076                    std::shared_ptr<Acts::SolenoidBField>>(m, "SolenoidBField")
0077             .def(py::init<Config>())
0078             .def(py::init([](double radius, double length, std::size_t nCoils,
0079                              double bMagCenter) {
0080                    return Acts::SolenoidBField{
0081                        Config{radius, length, nCoils, bMagCenter}};
0082                  }),
0083                  py::arg("radius"), py::arg("length"), py::arg("nCoils"),
0084                  py::arg("bMagCenter"));
0085
0086     py::class_<Config>(sol, "Config")
0087         .def(py::init<>())
0088         .def_readwrite("radius", &Config::radius)
0089         .def_readwrite("length", &Config::length)
0090         .def_readwrite("nCoils", &Config::nCoils)
0091         .def_readwrite("bMagCenter", &Config::bMagCenter);
0092   }
0093
0094   mex.def(
0095       "MagneticFieldMapXyz",
0096       [](const std::string& filename, const std::string& tree,
0097          double lengthUnit, double BFieldUnit, bool firstOctant) {
0098         const std::filesystem::path file = filename;
0099
0100         auto mapBins = [](std::array<std::size_t, 3> bins,
0101                           std::array<std::size_t, 3> sizes) {
0102           return (bins[0] * (sizes[1] * sizes[2]) + bins[1] * sizes[2] +
0103                   bins[2]);
0104         };
0105
0106         if (file.extension() == ".root") {
0107           auto map = ActsExamples::makeMagneticFieldMapXyzFromRoot(
0108               std::move(mapBins), file.native(), tree, lengthUnit, BFieldUnit,
0109               firstOctant);
0110           return std::make_shared<
0111               ActsExamples::detail::InterpolatedMagneticField3>(std::move(map));
0112         } else if (file.extension() == ".txt") {
0113           auto map = ActsExamples::makeMagneticFieldMapXyzFromText(
0114               std::move(mapBins), file.native(), lengthUnit, BFieldUnit,
0115               firstOctant);
0116           return std::make_shared<
0117               ActsExamples::detail::InterpolatedMagneticField3>(std::move(map));
0118         } else {
0119           throw std::runtime_error("Unsupported magnetic field map file type");
0120         }
0121       },
0122       py::arg("file"), py::arg("tree") = "bField",
0123       py::arg("lengthUnit") = Acts::UnitConstants::mm,
0124       py::arg("BFieldUnit") = Acts::UnitConstants::T,
0125       py::arg("firstOctant") = false);
0126
0127   mex.def(
0128       "MagneticFieldMapRz",
0129       [](const std::string& filename, const std::string& tree,
0130          double lengthUnit, double BFieldUnit, bool firstQuadrant) {
0131         const std::filesystem::path file = filename;
0132
0133         auto mapBins = [](std::array<std::size_t, 2> bins,
0134                           std::array<std::size_t, 2> sizes) {
0135           return (bins[1] * sizes[0] + bins[0]);
0136         };
0137
0138         if (file.extension() == ".root") {
0139           auto map = ActsExamples::makeMagneticFieldMapRzFromRoot(
0140               std::move(mapBins), file.native(), tree, lengthUnit, BFieldUnit,
0141               firstQuadrant);
0142           return std::make_shared<
0143               ActsExamples::detail::InterpolatedMagneticField2>(std::move(map));
0144         } else if (file.extension() == ".txt") {
0145           auto map = ActsExamples::makeMagneticFieldMapRzFromText(
0146               std::move(mapBins), file.native(), lengthUnit, BFieldUnit,
0147               firstQuadrant);
0148           return std::make_shared<
0149               ActsExamples::detail::InterpolatedMagneticField2>(std::move(map));
0150         } else {
0151           throw std::runtime_error("Unsupported magnetic field map file type");
0152         }
0153       },
0154       py::arg("file"), py::arg("tree") = "bField",
0155       py::arg("lengthUnit") = Acts::UnitConstants::mm,
0156       py::arg("BFieldUnit") = Acts::UnitConstants::T,
0157       py::arg("firstQuadrant") = false);
0158 }
0159
0160 }  // namespace Acts::Python