detray/options/track_generator_options.hpp

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 #pragma once
0010
0011 // Project include(s)
0012 #include "detray/definitions/algebra.hpp"
0013 #include "detray/utils/logging.hpp"
0014
0015 // Detray test include(s)
0016 #include "detray/options/options_handling.hpp"
0017 #include "detray/test/common/event_generator/random_track_generator_config.hpp"
0018 #include "detray/test/common/event_generator/uniform_track_generator_config.hpp"
0019
0020 // Boost
0021 #include "detray/options/boost_program_options.hpp"
0022
0023 // System include(s)
0024 #include <stdexcept>
0025 #include <vector>
0026
0027 namespace detray::options {
0028
0029 namespace detail {
0030
0031 /// Add options for detray event generation
0032 template <concepts::scalar scalar_t>
0033 void add_uniform_track_gen_options(
0034     boost::program_options::options_description &desc,
0035     const uniform_track_generator_config<scalar_t> &cfg) {
0036   desc.add_options()(
0037       "phi_steps",
0038       boost::program_options::value<std::size_t>()->default_value(
0039           cfg.phi_steps()),
0040       "No. phi steps for particle gun")(
0041       "eta_steps",
0042       boost::program_options::value<std::size_t>()->default_value(
0043           cfg.eta_steps()),
0044       "No. eta steps for particle gun")(
0045       "random_seed",
0046       boost::program_options::value<std::size_t>()->default_value(cfg.seed()),
0047       "Seed for the random number generator")(
0048       "eta_range",
0049       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0050       "Min, Max range of eta values for particle gun")(
0051       "randomize_charge", "Randomly flip charge sign per track")(
0052       "origin",
0053       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0054       "Coordinates for particle gun origin position [mm]")(
0055       "p_range",
0056       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0057       "Total momentum [range] of the test particles [GeV]")(
0058       "pT_range",
0059       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0060       "Transverse momentum [range] of the test particles [GeV]");
0061 }
0062
0063 /// Add options for detray event generation
0064 template <concepts::scalar scalar_t>
0065 void configure_uniform_track_gen_options(
0066     const boost::program_options::variables_map &vm,
0067     uniform_track_generator_config<scalar_t> &cfg) {
0068   cfg.phi_steps(vm["phi_steps"].as<std::size_t>());
0069   cfg.eta_steps(vm["eta_steps"].as<std::size_t>());
0070   cfg.seed(vm["random_seed"].as<std::size_t>());
0071   cfg.randomize_charge(vm.count("randomize_charge"));
0072
0073   if (vm.count("eta_range") != 0u) {
0074     const auto eta_range = vm["eta_range"].as<std::vector<scalar_t>>();
0075     if (eta_range.size() == 2u) {
0076       cfg.eta_range(eta_range[0], eta_range[1]);
0077     } else {
0078       throw std::invalid_argument("Eta range needs two arguments");
0079     }
0080   }
0081   if (vm.count("origin") != 0u) {
0082     const auto origin = vm["origin"].as<std::vector<scalar_t>>();
0083     if (origin.size() == 3u) {
0084       cfg.origin(origin[0] * unit<scalar_t>::mm, origin[1] * unit<scalar_t>::mm,
0085                  origin[2] * unit<scalar_t>::mm);
0086     } else {
0087       throw std::invalid_argument("Particle gun origin needs three arguments");
0088     }
0089   }
0090   if (vm.count("pT_range") != 0u && vm.count("p_range") != 0u) {
0091     throw std::invalid_argument(
0092         "Transverse and total momentum cannot be specified at the same "
0093         "time");
0094   }
0095   if (vm.count("pT_range") != 0u) {
0096     const auto pt_range = vm["pT_range"].as<std::vector<scalar_t>>();
0097
0098     // Default
0099     if (pt_range.empty()) {
0100       cfg.p_T(static_cast<scalar_t>(cfg.m_p_mag));
0101     } else if (pt_range.size() <= 2u) {
0102       if (pt_range.size() == 2u) {
0103         DETRAY_ERROR_HOST(
0104             "Momentum range not possible with uniform "
0105             "track generator: Using first value.");
0106       }
0107       cfg.p_T(pt_range[0] * unit<scalar_t>::GeV);
0108     } else {
0109       throw std::invalid_argument(
0110           "Wrong number of arguments for pT range: Need one argument or "
0111           "range");
0112     }
0113   } else {
0114     auto p_range = std::vector<scalar_t>{};
0115     if (vm.count("pT_range") != 0u) {
0116       p_range = vm["p_range"].as<std::vector<scalar_t>>();
0117     }
0118
0119     // Default
0120     if (p_range.empty()) {
0121       cfg.p_tot(static_cast<scalar_t>(cfg.m_p_mag));
0122     } else if (p_range.size() <= 2u) {
0123       if (p_range.size() == 2u) {
0124         DETRAY_ERROR_HOST(
0125             "Momentum range not possible with uniform "
0126             "track generator: Using first value.");
0127       }
0128       cfg.p_tot(p_range[0] * unit<scalar_t>::GeV);
0129     } else {
0130       std::string err_str{
0131           "Wrong number of arguments for p_tot range: Need one argument "
0132           "or range"};
0133       DETRAY_FATAL_HOST(err_str);
0134       throw std::invalid_argument(err_str);
0135     }
0136   }
0137 }
0138
0139 /// Add options for detray event generation
0140 template <concepts::scalar scalar_t>
0141 void add_rnd_track_gen_options(
0142     boost::program_options::options_description &desc,
0143     const random_track_generator_config<scalar_t> &cfg) {
0144   desc.add_options()(
0145       "n_tracks",
0146       boost::program_options::value<std::size_t>()->default_value(
0147           cfg.n_tracks()),
0148       "No. of tracks for particle gun")(
0149       "random_seed",
0150       boost::program_options::value<std::size_t>()->default_value(cfg.seed()),
0151       "Seed for the random number generator")(
0152       "theta_range",
0153       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0154       "Min, Max range of theta values for particle gun. Interval in [0, pi)")(
0155       "eta_range",
0156       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0157       "Min, Max range of eta values for particle gun")(
0158       "randomize_charge", "Randomly flip charge sign per track")(
0159       "origin",
0160       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0161       "Coordinates for particle gun origin position [mm]")(
0162       "p_range",
0163       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0164       "Total momentum [range] of the test particles [GeV]")(
0165       "pT_range",
0166       boost::program_options::value<std::vector<scalar_t>>()->multitoken(),
0167       "Transverse momentum [range] of the test particles [GeV]");
0168 }
0169
0170 /// Add options for detray event generation
0171 template <concepts::scalar scalar_t>
0172 void configure_rnd_track_gen_options(
0173     const boost::program_options::variables_map &vm,
0174     random_track_generator_config<scalar_t> &cfg) {
0175   cfg.n_tracks(vm["n_tracks"].as<std::size_t>());
0176   cfg.seed(vm["random_seed"].as<std::size_t>());
0177   cfg.randomize_charge(vm.count("randomize_charge"));
0178
0179   if (vm.count("eta_range") != 0u && vm.count("theta_range") != 0u) {
0180     throw std::invalid_argument(
0181         "Eta range and theta range cannot be specified at the same time");
0182   } else if (vm.count("eta_range") != 0u) {
0183     const auto eta_range = vm["eta_range"].as<std::vector<scalar_t>>();
0184     if (eta_range.size() == 2u) {
0185       scalar_t min_theta{2.f * std::atan(std::exp(-eta_range[0]))};
0186       scalar_t max_theta{2.f * std::atan(std::exp(-eta_range[1]))};
0187
0188       // Wrap around
0189       if (min_theta > max_theta) {
0190         scalar_t tmp{min_theta};
0191         min_theta = max_theta;
0192         max_theta = tmp;
0193       }
0194
0195       cfg.theta_range(min_theta, max_theta);
0196     } else {
0197       throw std::invalid_argument("Eta range needs two arguments");
0198     }
0199   } else if (vm.count("theta_range") != 0u) {
0200     const auto theta_range = vm["theta_range"].as<std::vector<scalar_t>>();
0201     if (theta_range.size() == 2u) {
0202       cfg.theta_range(theta_range[0], theta_range[1]);
0203     } else {
0204       throw std::invalid_argument("Theta range needs two arguments");
0205     }
0206   }
0207   if (vm.count("origin") != 0u) {
0208     const auto origin = vm["origin"].as<std::vector<scalar_t>>();
0209     if (origin.size() == 3u) {
0210       cfg.origin(origin[0] * unit<scalar_t>::mm, origin[1] * unit<scalar_t>::mm,
0211                  origin[2] * unit<scalar_t>::mm);
0212     } else {
0213       throw std::invalid_argument(
0214           "Particle gun origin needs three coordinates");
0215     }
0216   }
0217   if (vm.count("pT_range") != 0u && vm.count("p_range") != 0u) {
0218     throw std::invalid_argument(
0219         "Transverse and total momentum cannot be specified at the same "
0220         "time");
0221   }
0222   if (vm.count("pT_range") != 0u) {
0223     const auto pt_range = vm["pT_range"].as<std::vector<scalar_t>>();
0224
0225     // Default
0226     if (pt_range.empty()) {
0227       cfg.p_T(cfg.mom_range()[0]);
0228     } else if (pt_range.size() == 1u) {
0229       cfg.p_T(pt_range[0] * unit<scalar_t>::GeV);
0230     } else if (pt_range.size() == 2u) {
0231       cfg.pT_range(pt_range[0] * unit<scalar_t>::GeV,
0232                    pt_range[1] * unit<scalar_t>::GeV);
0233     } else {
0234       throw std::invalid_argument(
0235           "Wrong number of arguments for pT range: Need one argument or "
0236           "range");
0237     }
0238   } else {
0239     auto p_range = std::vector<scalar_t>();
0240     if (vm.count("pT_range") != 0u) {
0241       p_range = vm["p_range"].as<std::vector<scalar_t>>();
0242     }
0243
0244     // Default
0245     if (p_range.empty()) {
0246       cfg.p_tot(cfg.mom_range()[0]);
0247     } else if (p_range.size() == 1u) {
0248       cfg.p_tot(p_range[0] * unit<scalar_t>::GeV);
0249     } else if (p_range.size() == 2u) {
0250       cfg.mom_range(p_range[0] * unit<scalar_t>::GeV,
0251                     p_range[1] * unit<scalar_t>::GeV);
0252     } else {
0253       throw std::invalid_argument(
0254           "Wrong number of arguments for p_tot range: Need one argument "
0255           "or range");
0256     }
0257   }
0258 }
0259
0260 }  // namespace detail
0261
0262 /// Add options for the uniform track generator
0263 /// @{
0264 template <>
0265 void add_options<uniform_track_generator_config<float>>(
0266     boost::program_options::options_description &desc,
0267     const uniform_track_generator_config<float> &cfg) {
0268   detail::add_uniform_track_gen_options(desc, cfg);
0269 }
0270
0271 template <>
0272 void add_options<uniform_track_generator_config<double>>(
0273     boost::program_options::options_description &desc,
0274     const uniform_track_generator_config<double> &cfg) {
0275   detail::add_uniform_track_gen_options(desc, cfg);
0276 }
0277 /// @}
0278
0279 /// Configure the detray uniform track generator
0280 /// @{
0281 template <>
0282 void configure_options<uniform_track_generator_config<float>>(
0283     const boost::program_options::variables_map &vm,
0284     uniform_track_generator_config<float> &cfg) {
0285   detail::configure_uniform_track_gen_options(vm, cfg);
0286 }
0287
0288 template <>
0289 void configure_options<uniform_track_generator_config<double>>(
0290     const boost::program_options::variables_map &vm,
0291     uniform_track_generator_config<double> &cfg) {
0292   detail::configure_uniform_track_gen_options(vm, cfg);
0293 }
0294 /// @}
0295
0296 /// Add options for the random track generator
0297 /// @{
0298 template <>
0299 void add_options<random_track_generator_config<float>>(
0300     boost::program_options::options_description &desc,
0301     const random_track_generator_config<float> &cfg) {
0302   detail::add_rnd_track_gen_options(desc, cfg);
0303 }
0304
0305 template <>
0306 void add_options<random_track_generator_config<double>>(
0307     boost::program_options::options_description &desc,
0308     const random_track_generator_config<double> &cfg) {
0309   detail::add_rnd_track_gen_options(desc, cfg);
0310 }
0311 /// @}
0312
0313 /// Configure the detray random track generator
0314 /// @{
0315 template <>
0316 void configure_options<random_track_generator_config<float>>(
0317     const boost::program_options::variables_map &vm,
0318     random_track_generator_config<float> &cfg) {
0319   detail::configure_rnd_track_gen_options(vm, cfg);
0320 }
0321
0322 template <>
0323 void configure_options<random_track_generator_config<double>>(
0324     const boost::program_options::variables_map &vm,
0325     random_track_generator_config<double> &cfg) {
0326   detail::configure_rnd_track_gen_options(vm, cfg);
0327 }
0328 /// @}
0329
0330 }  // namespace detray::options