TrackSelector.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file TrackSelector.hpp

// This file is part of the Acts project.
//
// Copyright (C) 2023 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#pragma once

#include <cmath>
#include <functional>
#include <limits>
#include <ostream>
#include <vector>

namespace Acts {

class TrackSelector {
  static constexpr double inf = std::numeric_limits<double>::infinity();

 public:
  struct Config {
    // Minimum/maximum local positions.
    double loc0Min = -inf;
    double loc0Max = inf;
    double loc1Min = -inf;
    double loc1Max = inf;
    // Minimum/maximum track time.
    double timeMin = -inf;
    double timeMax = inf;
    // Direction cuts.
    double phiMin = -inf;
    double phiMax = inf;
    double etaMin = -inf;
    double etaMax = inf;
    double absEtaMin = 0.0;
    double absEtaMax = inf;
    // Momentum cuts.
    double ptMin = 0.0;
    double ptMax = inf;

    std::size_t minMeasurements = 0;

    // Helper factory functions to produce a populated config object more
    // conveniently

    Config& loc0(double min, double max);

    Config& loc1(double min, double max);

    Config& time(double min, double max);

    Config& phi(double min, double max);

    Config& eta(double min, double max);

    Config& absEta(double min, double max);

    Config& pt(double min, double max);

    friend std::ostream& operator<<(std::ostream& os, const Config& cuts);
  };

  struct EtaBinnedConfig {
    std::vector<Config> cutSets = {};

    std::vector<double> absEtaEdges = {};

    std::size_t nEtaBins() const { return absEtaEdges.size() - 1; }

    EtaBinnedConfig() : cutSets{{}}, absEtaEdges{{0, inf}} {};

    EtaBinnedConfig(double etaMin) : cutSets{}, absEtaEdges{etaMin} {}

    EtaBinnedConfig(std::vector<double> absEtaEdgesIn)
        : absEtaEdges{std::move(absEtaEdgesIn)} {
      cutSets.resize(absEtaEdges.size() - 1);
    }

    EtaBinnedConfig(Config cutSet) : cutSets{cutSet}, absEtaEdges{{0, inf}} {}

    EtaBinnedConfig& addCuts(double etaMax,
                             const std::function<void(Config&)>& callback = {});

    EtaBinnedConfig& addCuts(const std::function<void(Config&)>& callback = {});

    friend std::ostream& operator<<(std::ostream& os,
                                    const EtaBinnedConfig& cfg);

    std::size_t binIndex(double eta) const;

    const Config& getCuts(double eta) const;
  };

  TrackSelector(const Config& config);

  TrackSelector(const EtaBinnedConfig& config);

  template <typename input_tracks_t, typename output_tracks_t>
  void selectTracks(const input_tracks_t& inputTracks,
                    output_tracks_t& outputTracks) const;

  template <typename track_proxy_t>
  bool isValidTrack(const track_proxy_t& track) const;

  const EtaBinnedConfig& config() const { return m_cfg; }

 private:
  EtaBinnedConfig m_cfg;
  bool m_isUnbinned;
  bool m_noEtaCuts;
};

inline TrackSelector::Config& TrackSelector::Config::loc0(double min,
                                                          double max) {
  loc0Min = min;
  loc0Max = max;
  return *this;
}

inline TrackSelector::Config& TrackSelector::Config::loc1(double min,
                                                          double max) {
  loc1Min = min;
  loc1Max = max;
  return *this;
}

inline TrackSelector::Config& TrackSelector::Config::time(double min,
                                                          double max) {
  timeMin = min;
  timeMax = max;
  return *this;
}

inline TrackSelector::Config& TrackSelector::Config::phi(double min,
                                                         double max) {
  phiMin = min;
  phiMax = max;
  return *this;
}

inline TrackSelector::Config& TrackSelector::Config::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::Config& TrackSelector::Config::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::Config& TrackSelector::Config::pt(double min,
                                                        double max) {
  ptMin = min;
  ptMax = max;
  return *this;
}

inline std::ostream& operator<<(std::ostream& os,
                                const TrackSelector::Config& 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::EtaBinnedConfig& TrackSelector::EtaBinnedConfig::addCuts(
    double etaMax, const std::function<void(Config&)>& 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::EtaBinnedConfig& TrackSelector::EtaBinnedConfig::addCuts(
    const std::function<void(Config&)>& callback) {
  return addCuts(inf, callback);
}

inline std::size_t TrackSelector::EtaBinnedConfig::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::Config& TrackSelector::EtaBinnedConfig::getCuts(
    double eta) const {
  return nEtaBins() == 1 ? cutSets[0] : cutSets[binIndex(eta)];
}

inline std::ostream& operator<<(std::ostream& os,
                                const TrackSelector::EtaBinnedConfig& cfg) {
  os << "TrackSelector::EtaBinnedConfig:\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 {
  for (auto track : inputTracks) {
    if (!isValidTrack(track)) {
      continue;
    }
    auto destProxy = outputTracks.getTrack(outputTracks.addTrack());
    destProxy.copyFrom(track, false);
    destProxy.tipIndex() = track.tipIndex();
  }
}

template <typename track_proxy_t>
bool TrackSelector::isValidTrack(const track_proxy_t& track) const {
  auto checkMin = [](auto x, auto min) { return min <= x; };
  auto within = [](double x, double min, double max) {
    return (min <= x) and (x < max);
  };

  const auto theta = track.theta();

  constexpr double kUnset = -std::numeric_limits<double>::infinity();

  double _eta = kUnset;
  double _absEta = kUnset;

  auto absEta = [&]() {
    if (_absEta == kUnset) {
      _eta = -std::log(std::tan(theta / 2));
      _absEta = std::abs(_eta);
    }
    return _absEta;
  };

  const Config* cutsPtr{nullptr};
  if (!m_isUnbinned) {
    if (absEta() < m_cfg.absEtaEdges.front() ||
        _absEta >= m_cfg.absEtaEdges.back()) {
      return false;
    }
    cutsPtr = &m_cfg.getCuts(_eta);
  } else {
    cutsPtr = &m_cfg.cutSets.front();
  }

  const Config& cuts = *cutsPtr;

  return within(track.transverseMomentum(), cuts.ptMin, cuts.ptMax) and
         (m_noEtaCuts || (within(absEta(), 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::EtaBinnedConfig& config)
    : m_cfg(config) {
  if (m_cfg.cutSets.size() != m_cfg.nEtaBins()) {
    throw std::invalid_argument{
        "TrackSelector cut / eta bin configuration is inconsistent"};
  }

  m_isUnbinned = false;
  if (m_cfg.nEtaBins() == 1) {
    static const std::vector<double> infVec = {0, inf};
    bool limitEta = m_cfg.absEtaEdges != infVec;
    m_isUnbinned = !limitEta;  // single bin, no eta edges given

    const Config& 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"};
    }
  }

  m_noEtaCuts = m_isUnbinned;
  for (const auto& cuts : m_cfg.cutSets) {
    if (cuts.etaMin != -inf || cuts.etaMax != inf || cuts.absEtaMin != 0.0 ||
        cuts.absEtaMax != inf) {
      m_noEtaCuts = false;
    }
  }
}

inline TrackSelector::TrackSelector(const Config& config)
    : TrackSelector{EtaBinnedConfig{config}} {}

}  // namespace Acts