feat: TrackSelector supports eta-dependent cuts (#2405)

This changes `TrackSelector` so that eta dependent cuts can be supplied. To this end, a new nested struct `CutConfig` is introduced. The existing `Config` struct is auto-convertible from that struct, and will produce a single cut set for all eta.

You can also supply `Config` with a vector of absolute eta bin edges and correspondingly sized (`N-1`) vector of `CutConfig` structs. The selector will then look up the associated cut value.

Note: the `CutConfig` properties `etaMin`, `etaMax`, `absEtaMin` and `absEtaMax` are only supported if there is a single cut configuration with an abs eta bin from `0.0` to `+infinity`. 

The `TrackSelector` constructor checks consistency of the configuration.

Core/include/Acts/TrackFinding/TrackSelector.hpp

@@ -9,41 +9,163 @@
 #pragma once
 
 #include <cmath>
+#include <functional>
 #include <limits>
+#include <ostream>
+#include <vector>
 
 namespace Acts {
 
 /// Class which performs filtering of tracks. It accepts an input and an output
 /// track container and uses the built-in copy facility to copy tracks into the
 /// output container.
 class TrackSelector {
+  static constexpr double inf = std::numeric_limits<double>::infinity();
+
  public:
-  struct Config {
+  /// Configuration of a set of cuts for a single eta bin
+  /// Default construction yields a set of cuts that accepts everything.
+  struct CutConfig {
     // Minimum/maximum local positions.
-    double loc0Min = -std::numeric_limits<double>::infinity();
-    double loc0Max = std::numeric_limits<double>::infinity();
-    double loc1Min = -std::numeric_limits<double>::infinity();
-    double loc1Max = std::numeric_limits<double>::infinity();
+    double loc0Min = -inf;
+    double loc0Max = inf;
+    double loc1Min = -inf;
+    double loc1Max = inf;
     // Minimum/maximum track time.
-    double timeMin = -std::numeric_limits<double>::infinity();
-    double timeMax = std::numeric_limits<double>::infinity();
+    double timeMin = -inf;
+    double timeMax = inf;
     // Direction cuts.
-    double phiMin = -std::numeric_limits<double>::infinity();
-    double phiMax = std::numeric_limits<double>::infinity();
-    double etaMin = -std::numeric_limits<double>::infinity();
-    double etaMax = std::numeric_limits<double>::infinity();
+    double phiMin = -inf;
+    double phiMax = inf;
+    double etaMin = -inf;
+    double etaMax = inf;
     double absEtaMin = 0.0;
-    double absEtaMax = std::numeric_limits<double>::infinity();
+    double absEtaMax = inf;
     // Momentum cuts.
     double ptMin = 0.0;
-    double ptMax = std::numeric_limits<double>::infinity();
+    double ptMax = inf;
 
     std::size_t minMeasurements = 0;
+
+    // Helper factory functions to produce a populated config object more
+    // conveniently
+
+    /// Set loc0 acceptance range
+    /// @param min Minimum value
+    /// @param max Maximum value
+    /// @return Reference to this object
+    CutConfig& loc0(double min, double max);
+
+    /// Set loc1 acceptance range
+    /// @param min Minimum value
+    /// @param max Maximum value
+    /// @return Reference to this object
+    CutConfig& loc1(double min, double max);
+
+    /// Set time acceptance range
+    /// @param min Minimum value
+    /// @param max Maximum value
+    /// @return Reference to this object
+    CutConfig& time(double min, double max);
+
+    /// Set phi acceptance range
+    /// @param min Minimum value
+    /// @param max Maximum value
+    /// @return Reference to this object
+    CutConfig& phi(double min, double max);
+
+    /// Set the eta acceptance range
+    /// @param min Minimum value
+    /// @param max Maximum value
+    /// @return Reference to this object
+    CutConfig& eta(double min, double max);
+
+    /// Set the absolute eta acceptance range
+    /// @param min Minimum value
+    /// @param max Maximum value
+    /// @return Reference to this object
+    CutConfig& absEta(double min, double max);
+
+    /// Set the pt acceptance range
+    /// @param min Minimum value
+    /// @param max Maximum value
+    /// @return Reference to this object
+    CutConfig& pt(double min, double max);
+
+    /// Print this set of cuts to an output stream
+    /// @param os Output stream
+    /// @param cuts Cuts to print
+    /// @return Reference to the output stream
+    friend std::ostream& operator<<(std::ostream& os, const CutConfig& cuts);
+  };
+
+  /// Main config object for the track selector. Combines a set of cut
+  /// configurations and corresponding eta bins
+  struct Config {
+    /// Cut sets for each eta bin
+    std::vector<CutConfig> cutSets = {};
+
+    /// Eta bin edges for varying cuts by eta
+    std::vector<double> absEtaEdges = {};
+
+    /// Get the number of eta bins
+    /// @return Number of eta bins
+    std::size_t nEtaBins() const { return absEtaEdges.size() - 1; }
+
+    /// Construct an empty (accepts everything) configuration.
+    /// Results in a single cut set and one abs eta bin from 0 to infinity.
+    Config() : cutSets{{}}, absEtaEdges{{0, inf}} {};
+
+    /// Constructor to create a config object that is not upper-bounded.
+    /// This is useful to use the "fluent" API to populate the configuration.
+    /// @param etaMin Minimum eta bin edge
+    Config(double etaMin) : cutSets{}, absEtaEdges{etaMin} {}
+
+    /// Constructor from a vector of eta bin edges. This automatically
+    /// initializes all the cuts to be the same for all eta and be essentially
+    /// no-op.
+    /// @param absEtaEdgesIn is the vector of eta bin edges
+    Config(std::vector<double> absEtaEdgesIn)
+        : absEtaEdges{std::move(absEtaEdgesIn)} {
+      cutSets.resize(absEtaEdges.size() - 1);
+    }
+
+    /// Auto-converting constructor from a single cut configuration.
+    /// Results in a single absolute eta bin from 0 to infinity.
+    Config(CutConfig cutSet) : cutSets{cutSet}, absEtaEdges{{0, inf}} {}
+
+    /// Add a new eta bin with the given upper bound.
+    /// @param etaMax Upper bound of the new eta bin
+    /// @param callback Callback to configure the cuts for this eta bin
+    /// @return Reference to this object
+    Config& addCuts(double etaMax,
+                    const std::function<void(CutConfig&)>& callback = {});
+
+    /// Add a new eta bin with an upper bound of +infinity.
+    /// @param callback Callback to configure the cuts for this eta bin
+    /// @return Reference to this object
+    Config& addCuts(const std::function<void(CutConfig&)>& callback = {});
+
+    /// Print this configuration to an output stream
+    /// @param os Output stream
+    /// @param cfg Configuration to print
+    /// @return Reference to the output stream
+    friend std::ostream& operator<<(std::ostream& os, const Config& cfg);
+
+    /// Get the index of the eta bin for a given eta
+    /// @param eta Eta value
+    /// @return Index of the eta bin
+    std::size_t binIndex(double eta) const;
+
+    /// Get the cuts for a given eta
+    /// @param eta Eta value
+    /// @return Cuts for the given eta
+    const CutConfig& getCuts(double eta) const;
   };
 
   /// Constructor from a config object
   /// @param config is the configuration object
-  TrackSelector(const Config& config) : m_cfg(config) {}
+  TrackSelector(const Config& config);
 
   /// Select tracks from an input container and copy them into an output
   /// container
@@ -70,6 +192,133 @@ class TrackSelector {
   Config m_cfg;
 };
 
+inline TrackSelector::CutConfig& TrackSelector::CutConfig::loc0(double min,
+                                                                double max) {
+  loc0Min = min;
+  loc0Max = max;
+  return *this;
+}
+
+inline TrackSelector::CutConfig& TrackSelector::CutConfig::loc1(double min,
+                                                                double max) {
+  loc1Min = min;
+  loc1Max = max;
+  return *this;
+}
+
+inline TrackSelector::CutConfig& TrackSelector::CutConfig::time(double min,
+                                                                double max) {
+  timeMin = min;
+  timeMax = max;
+  return *this;
+}
+
+inline TrackSelector::CutConfig& TrackSelector::CutConfig::phi(double min,
+                                                               double max) {
+  phiMin = min;
+  phiMax = max;
+  return *this;
+}
+
+inline TrackSelector::CutConfig& TrackSelector::CutConfig::eta(double min,
+                                                               double max) {
+  if (absEtaMin != 0.0 || absEtaMax != inf) {
+    throw std::invalid_argument(
+        "Cannot set both eta and absEta cuts in the same cut set");
+  }
+  etaMin = min;
+  etaMax = max;
+  return *this;
+}
+
+inline TrackSelector::CutConfig& TrackSelector::CutConfig::absEta(double min,
+                                                                  double max) {
+  if (etaMin != -inf || etaMax != inf) {
+    throw std::invalid_argument(
+        "Cannot set both eta and absEta cuts in the same cut set");
+  }
+  absEtaMin = min;
+  absEtaMax = max;
+  return *this;
+}
+
+inline TrackSelector::CutConfig& TrackSelector::CutConfig::pt(double min,
+                                                              double max) {
+  ptMin = min;
+  ptMax = max;
+  return *this;
+}
+
+inline std::ostream& operator<<(std::ostream& os,
+                                const TrackSelector::CutConfig& cuts) {
+  auto print = [&](const char* name, const auto& min, const auto& max) {
+    os << " - " << min << " <= " << name << " < " << max << "\n";
+  };
+
+  print("loc0", cuts.loc0Min, cuts.loc0Max);
+  print("loc1", cuts.loc1Min, cuts.loc1Max);
+  print("time", cuts.timeMin, cuts.timeMax);
+  print("phi", cuts.phiMin, cuts.phiMax);
+  print("eta", cuts.etaMin, cuts.etaMax);
+  print("absEta", cuts.absEtaMin, cuts.absEtaMax);
+  print("pt", cuts.ptMin, cuts.ptMax);
+  os << " - " << cuts.minMeasurements << " <= nMeasurements\n";
+
+  return os;
+}
+
+inline TrackSelector::Config& TrackSelector::Config::addCuts(
+    double etaMax, const std::function<void(CutConfig&)>& callback) {
+  if (etaMax <= absEtaEdges.back()) {
+    throw std::invalid_argument{
+        "Abs Eta bin edges must be in increasing order"};
+  }
+
+  if (etaMax < 0.0) {
+    throw std::invalid_argument{"Abs Eta bin edges must be positive"};
+  }
+
+  absEtaEdges.push_back(etaMax);
+  cutSets.emplace_back();
+  if (callback) {
+    callback(cutSets.back());
+  }
+  return *this;
+}
+
+inline TrackSelector::Config& TrackSelector::Config::addCuts(
+    const std::function<void(CutConfig&)>& callback) {
+  return addCuts(inf, callback);
+}
+
+inline std::size_t TrackSelector::Config::binIndex(double eta) const {
+  if (std::abs(eta) >= absEtaEdges.back()) {
+    throw std::invalid_argument{"Eta is outside the abs eta bin edges"};
+  }
+
+  auto binIt =
+      std::upper_bound(absEtaEdges.begin(), absEtaEdges.end(), std::abs(eta));
+  std::size_t index = std::distance(absEtaEdges.begin(), binIt) - 1;
+  return index;
+}
+
+inline const TrackSelector::CutConfig& TrackSelector::Config::getCuts(
+    double eta) const {
+  return nEtaBins() == 1 ? cutSets[0] : cutSets[binIndex(eta)];
+}
+
+inline std::ostream& operator<<(std::ostream& os,
+                                const TrackSelector::Config& cfg) {
+  os << "TrackSelector::Config:\n";
+
+  for (std::size_t i = 1; i < cfg.absEtaEdges.size(); i++) {
+    os << cfg.absEtaEdges[i - 1] << " <= eta < " << cfg.absEtaEdges[i] << "\n";
+    os << cfg.cutSets[i - 1];
+  }
+
+  return os;
+}
+
 template <typename input_tracks_t, typename output_tracks_t>
 void TrackSelector::selectTracks(const input_tracks_t& inputTracks,
                                  output_tracks_t& outputTracks) const {
@@ -89,16 +338,47 @@ bool TrackSelector::isValidTrack(const track_proxy_t& track) const {
   auto within = [](double x, double min, double max) {
     return (min <= x) and (x < max);
   };
+
   const auto theta = track.theta();
   const auto eta = -std::log(std::tan(theta / 2));
-  return within(track.transverseMomentum(), m_cfg.ptMin, m_cfg.ptMax) and
-         within(std::abs(eta), m_cfg.absEtaMin, m_cfg.absEtaMax) and
-         within(eta, m_cfg.etaMin, m_cfg.etaMax) and
-         within(track.phi(), m_cfg.phiMin, m_cfg.phiMax) and
-         within(track.loc0(), m_cfg.loc0Min, m_cfg.loc0Max) and
-         within(track.loc1(), m_cfg.loc1Min, m_cfg.loc1Max) and
-         within(track.time(), m_cfg.timeMin, m_cfg.timeMax) and
-         checkMin(track.nMeasurements(), m_cfg.minMeasurements);
+  const auto absEta = std::abs(eta);
+
+  if (absEta < m_cfg.absEtaEdges.front() ||
+      absEta >= m_cfg.absEtaEdges.back()) {
+    return false;
+  }
+
+  const CutConfig& cuts = m_cfg.getCuts(eta);
+
+  return within(track.transverseMomentum(), cuts.ptMin, cuts.ptMax) and
+         within(std::abs(eta), cuts.absEtaMin, cuts.absEtaMax) and
+         within(eta, cuts.etaMin, cuts.etaMax) and
+         within(track.phi(), cuts.phiMin, cuts.phiMax) and
+         within(track.loc0(), cuts.loc0Min, cuts.loc0Max) and
+         within(track.loc1(), cuts.loc1Min, cuts.loc1Max) and
+         within(track.time(), cuts.timeMin, cuts.timeMax) and
+         checkMin(track.nMeasurements(), cuts.minMeasurements);
+}
+
+inline TrackSelector::TrackSelector(const TrackSelector::Config& config)
+    : m_cfg(config) {
+  if (m_cfg.cutSets.size() != m_cfg.nEtaBins()) {
+    throw std::invalid_argument{
+        "TrackSelector cut / eta bin configuration is inconsistent"};
+  }
+
+  if (m_cfg.nEtaBins() == 1) {
+    static const std::vector<double> infVec = {0, inf};
+    bool limitEta = m_cfg.absEtaEdges != infVec;
+
+    const CutConfig& cuts = m_cfg.cutSets[0];
+
+    if (limitEta && (cuts.etaMin != -inf || cuts.etaMax != inf ||
+                     cuts.absEtaMin != 0.0 || cuts.absEtaMax != inf)) {
+      throw std::invalid_argument{
+          "Explicit eta cuts are only valid for single eta bin"};
+    }
+  }
 }
 
 }  // namespace Acts