EIC Software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
ParticleSelector.cpp
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file ParticleSelector.cpp
1 // This file is part of the Acts project.
2 //
3 // Copyright (C) 2019-2020 CERN for the benefit of the Acts project
4 //
5 // This Source Code Form is subject to the terms of the Mozilla Public
6 // License, v. 2.0. If a copy of the MPL was not distributed with this
7 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
8 
9 #include "ActsExamples/TruthTracking/ParticleSelector.hpp"
10 
11 #include "Acts/Utilities/Helpers.hpp"
12 #include "Acts/Utilities/Units.hpp"
13 #include "ActsExamples/EventData/SimParticle.hpp"
14 #include "ActsExamples/Framework/WhiteBoard.hpp"
15 #include "ActsExamples/Utilities/Options.hpp"
16 
17 #include <algorithm>
18 #include <stdexcept>
19 #include <vector>
20 
21 #include <boost/program_options.hpp>
22 
23 void ActsExamples::ParticleSelector::addOptions(Options::Description& desc) {
24  using boost::program_options::bool_switch;
25  using boost::program_options::value;
26  using Options::Interval;
27 
28  auto opt = desc.add_options();
29  opt("select-rho-mm", value<Interval>()->value_name("MIN:MAX"),
30  "Select particle transverse distance to the origin in mm");
31  opt("select-absz-mm", value<Interval>()->value_name("MIN:MAX"),
32  "Select particle absolute longitudinal distance to the origin in mm");
33  opt("select-time-ns", value<Interval>()->value_name("MIN:MAX"),
34  "Select particle time in ns");
35  opt("select-phi-degree", value<Interval>()->value_name("MIN:MAX"),
36  "Select particle direction angle in the transverse plane in degree");
37  opt("select-eta", value<Interval>()->value_name("MIN:MAX"),
38  "Select particle pseudo-rapidity");
39  opt("select-abseta", value<Interval>()->value_name("MIN:MAX"),
40  "Select particle absolute pseudo-rapidity");
41  opt("select-pt-gev", value<Interval>()->value_name("MIN:MAX"),
42  "Select particle transverse momentum in GeV");
43  opt("remove-charged", bool_switch(), "Remove charged particles");
44  opt("remove-neutral", bool_switch(), "Remove neutral particles");
45 }
46 
47 ActsExamples::ParticleSelector::Config
48 ActsExamples::ParticleSelector::readConfig(const Options::Variables& vars) {
49  using namespace Acts::UnitLiterals;
50 
51  // Set boundary values if the given config exists
52  auto extractInterval = [&](const char* name, auto unit, auto& lower,
53  auto& upper) {
54  if (vars[name].empty()) {
55  return;
56  }
57  auto interval = vars[name].as<Options::Interval>();
58  lower = interval.lower.value_or(lower) * unit;
59  upper = interval.upper.value_or(upper) * unit;
60  };
61 
62  Config cfg;
63  extractInterval("select-rho-mm", 1_mm, cfg.rhoMin, cfg.rhoMax);
64  extractInterval("select-absz-mm", 1_mm, cfg.absZMin, cfg.absZMax);
65  extractInterval("select-time-ns", 1_ns, cfg.timeMin, cfg.timeMax);
66  extractInterval("select-phi-degree", 1_degree, cfg.phiMin, cfg.phiMax);
67  extractInterval("select-eta", 1.0, cfg.etaMin, cfg.etaMax);
68  extractInterval("select-abseta", 1.0, cfg.absEtaMin, cfg.absEtaMax);
69  extractInterval("select-pt-gev", 1_GeV, cfg.ptMin, cfg.ptMax);
70  cfg.removeCharged = vars["remove-charged"].as<bool>();
71  cfg.removeNeutral = vars["remove-neutral"].as<bool>();
72  return cfg;
73 }
74 
75 ActsExamples::ParticleSelector::ParticleSelector(const Config& cfg,
76  Acts::Logging::Level lvl)
77  : BareAlgorithm("ParticleSelector", lvl), m_cfg(cfg) {
78  if (m_cfg.inputParticles.empty()) {
79  throw std::invalid_argument("Missing input particles collection");
80  }
81  if (m_cfg.outputParticles.empty()) {
82  throw std::invalid_argument("Missing output particles collection");
83  }
84  ACTS_DEBUG("selection particle rho [" << m_cfg.rhoMin << "," << m_cfg.rhoMax
85  << ")");
86  ACTS_DEBUG("selection particle |z| [" << m_cfg.absZMin << "," << m_cfg.absZMax
87  << ")");
88  ACTS_DEBUG("selection particle time [" << m_cfg.timeMin << ","
89  << m_cfg.timeMax << ")");
90  ACTS_DEBUG("selection particle phi [" << m_cfg.phiMin << "," << m_cfg.phiMax
91  << ")");
92  ACTS_DEBUG("selection particle eta [" << m_cfg.etaMin << "," << m_cfg.etaMax
93  << ")");
94  ACTS_DEBUG("selection particle |eta| [" << m_cfg.absEtaMin << ","
95  << m_cfg.absEtaMax << ")");
96  ACTS_DEBUG("selection particle pt [" << m_cfg.ptMin << "," << m_cfg.ptMax
97  << ")");
98  ACTS_DEBUG("remove charged particles " << m_cfg.removeCharged);
99  ACTS_DEBUG("remove neutral particles " << m_cfg.removeNeutral);
100 }
101 
102 ActsExamples::ProcessCode ActsExamples::ParticleSelector::execute(
103  const AlgorithmContext& ctx) const {
104  // helper functions to select tracks
105  auto within = [](double x, double min, double max) {
106  return (min <= x) and (x < max);
107  };
108  auto isValidParticle = [&](const ActsFatras::Particle& p) {
109  const auto eta = Acts::VectorHelpers::eta(p.unitDirection());
110  const auto phi = Acts::VectorHelpers::phi(p.unitDirection());
111  const auto rho = Acts::VectorHelpers::perp(p.position());
112  // define charge selection
113  const bool validNeutral = (p.charge() == 0) and not m_cfg.removeNeutral;
114  const bool validCharged = (p.charge() != 0) and not m_cfg.removeCharged;
115  const bool validCharge = validNeutral or validCharged;
116  return validCharge and
117  within(p.transverseMomentum(), m_cfg.ptMin, m_cfg.ptMax) and
118  within(std::abs(eta), m_cfg.absEtaMin, m_cfg.absEtaMax) and
119  within(eta, m_cfg.etaMin, m_cfg.etaMax) and
120  within(phi, m_cfg.phiMin, m_cfg.phiMax) and
121  within(std::abs(p.position()[Acts::ePos2]), m_cfg.absZMin,
122  m_cfg.absZMax) and
123  within(rho, m_cfg.rhoMin, m_cfg.rhoMax) and
124  within(p.time(), m_cfg.timeMin, m_cfg.timeMax);
125  };
126 
127  // prepare input/ output types
128  const auto& inputParticles =
129  ctx.eventStore.get<SimParticleContainer>(m_cfg.inputParticles);
130  SimParticleContainer outputParticles;
131  outputParticles.reserve(inputParticles.size());
132 
133  // copy selected particles
134  for (const auto& inputParticle : inputParticles) {
135  if (isValidParticle(inputParticle)) {
136  // the input parameters should already be
137  outputParticles.insert(outputParticles.end(), inputParticle);
138  }
139  }
140  outputParticles.shrink_to_fit();
141 
142  ACTS_DEBUG("event " << ctx.eventNumber << " selected "
143  << outputParticles.size() << " from "
144  << inputParticles.size() << " particles");
145 
146  ctx.eventStore.add(m_cfg.outputParticles, std::move(outputParticles));
147  return ProcessCode::SUCCESS;
148 }