feat: full_chain_itk.py uses eta-dependent cuts (acts-project#2934)

* `addCKFTracks` can now specify a `trackSelectorConfig = list(TrackSelectorConfig)` for eta-dependent cuts, with binning specified in `absEta[1]`. `addCKFTracks()` can be called with any of:
  1. `trackSelectorConfig = (TrackSelectorConfig(), TrackSelectorConfig(), ...)`. See `full_chain_itk.py` for an example.
  2. `trackSelectorConfig = TrackSelectorConfig()`. Can also be specified as `TrackSelectorConfig()` as an non-kwarg. See `full_chain_odd.py` for an example.
  3. `trackSelectorConfig = None` (default)
* `TrackSelectorConfig` includes new cuts `maxHoles`, `maxOutliers`, `maxSharedHits`, and `maxChi2`.
* `full_chain_itk.py` now uses Athena's ITk track selections.

Examples/Algorithms/TrackFinding/include/ActsExamples/TrackFinding/TrackFindingAlgorithm.hpp

@@ -34,7 +34,9 @@
 #include <functional>
 #include <limits>
 #include <memory>
+#include <optional>
 #include <string>
+#include <variant>
 #include <vector>
 
 #include <tbb/combinable.h>
@@ -94,7 +96,9 @@ class TrackFindingAlgorithm final : public IAlgorithm {
     /// Compute shared hit information
     bool computeSharedHits = false;
     /// Track selector config
-    std::optional<Acts::TrackSelector::Config> trackSelectorCfg = std::nullopt;
+    std::optional<std::variant<Acts::TrackSelector::Config,
+                               Acts::TrackSelector::EtaBinnedConfig>>
+        trackSelectorCfg = std::nullopt;
     /// Run backward finding
     bool backward = false;
     /// Maximum number of propagation steps

Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithm.cpp

@@ -44,13 +44,14 @@ ActsExamples::TrackFindingAlgorithm::TrackFindingAlgorithm(
     Config config, Acts::Logging::Level level)
     : ActsExamples::IAlgorithm("TrackFindingAlgorithm", level),
       m_cfg(std::move(config)),
-      m_trackSelector([this]() -> std::optional<Acts::TrackSelector> {
-        if (m_cfg.trackSelectorCfg.has_value()) {
-          return {m_cfg.trackSelectorCfg.value()};
-        } else {
-          return std::nullopt;
-        }
-      }()) {
+      m_trackSelector(
+          m_cfg.trackSelectorCfg.has_value()
+              ? std::visit(
+                    [](const auto& cfg) -> std::optional<Acts::TrackSelector> {
+                      return {cfg};
+                    },
+                    m_cfg.trackSelectorCfg.value())
+              : std::nullopt) {
   if (m_cfg.inputMeasurements.empty()) {
     throw std::invalid_argument("Missing measurements input collection");
   }

Examples/Python/python/acts/examples/__init__.py

@@ -159,6 +159,9 @@ def NamedTypeArgsWrapper(*args, **kwargs):
                     cls is not None
                     and v is not None
                     and type(v).__module__ == int.__module__  # is v a 'builtins'?
+                    and not (
+                        issubclass(type(v), Iterable) and all(type(e) is cls for e in v)
+                    )  # not [cls]
                 ):
                     if issubclass(cls, Iterable):
                         kwargs[k] = cls(*v)

Examples/Python/python/acts/examples/reconstruction.py

@@ -125,9 +125,13 @@
         "pt",
         "phi",
         "nMeasurementsMin",
+        "maxHoles",
+        "maxOutliers",
+        "maxSharedHits",
+        "maxChi2",
         "nMeasurementsGroupMin",
     ],
-    defaults=[(None, None)] * 7 + [None] * 2,
+    defaults=[(None, None)] * 7 + [None] * 6,
 )
 
 CkfConfig = namedtuple(
@@ -1125,7 +1129,9 @@ def addCKFTracks(
     s: acts.examples.Sequencer,
     trackingGeometry: acts.TrackingGeometry,
     field: acts.MagneticFieldProvider,
-    trackSelectorConfig: Optional[TrackSelectorConfig] = None,
+    trackSelectorConfig: Optional[
+        Union[TrackSelectorConfig, List[TrackSelectorConfig]]
+    ] = None,
     ckfConfig: CkfConfig = CkfConfig(),
     outputDirCsv: Optional[Union[Path, str]] = None,
     outputDirRoot: Optional[Union[Path, str]] = None,
@@ -1147,13 +1153,60 @@ def addCKFTracks(
         the output folder for the Root output, None triggers no output
     trackSelectorConfig : TrackSelectorConfig(loc0, loc1, time, eta, absEta, pt, phi, minMeasurements)
         TrackSelector configuration. Each range is specified as a tuple of (min,max).
+        Specify as a list(TrackSelectorConfig) for eta-dependent cuts, with binning specified by absEta[1].
         Defaults of no cuts specified in Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/TrackSelector.hpp
     writeTrajectories : bool, True
         write trackstates_ckf.root and tracksummary_ckf.root ntuples? These can be quite large.
     """
 
     customLogLevel = acts.examples.defaultLogging(s, logLevel)
 
+    tslist = (
+        []
+        if trackSelectorConfig is None
+        else (
+            [trackSelectorConfig]
+            if type(trackSelectorConfig) is TrackSelectorConfig
+            else trackSelectorConfig
+        )
+    )
+    cutSets = [
+        acts.TrackSelector.Config(
+            **acts.examples.defaultKWArgs(
+                loc0Min=c.loc0[0],
+                loc0Max=c.loc0[1],
+                loc1Min=c.loc1[0],
+                loc1Max=c.loc1[1],
+                timeMin=c.time[0],
+                timeMax=c.time[1],
+                phiMin=c.phi[0],
+                phiMax=c.phi[1],
+                etaMin=c.eta[0],
+                etaMax=c.eta[1],
+                absEtaMin=c.absEta[0],
+                absEtaMax=c.absEta[1] if len(tslist) == 1 else None,
+                ptMin=c.pt[0],
+                ptMax=c.pt[1],
+                minMeasurements=c.nMeasurementsMin,
+                maxHoles=c.maxHoles,
+                maxOutliers=c.maxOutliers,
+                maxSharedHits=c.maxSharedHits,
+                maxChi2=c.maxChi2,
+                measurementCounter=c.nMeasurementsGroupMin,
+            )
+        )
+        for c in tslist
+    ]
+    if len(tslist) == 0:
+        trkSelCfg = None
+    elif len(tslist) == 1:
+        trkSelCfg = cutSets[0]
+    else:
+        trkSelCfg = acts.TrackSelector.EtaBinnedConfig(
+            cutSets=cutSets,
+            absEtaEdges=[cutSets[0].absEtaMin] + [c.absEta[1] for c in tslist],
+        )
+
     # Setup the track finding algorithm with CKF
     # It takes all the source links created from truth hit smearing, seeds from
     # truth particle smearing and source link selection config
@@ -1171,28 +1224,6 @@ def addCKFTracks(
                 )
             ]
         ),
-        trackSelectorCfg=acts.TrackSelector.Config(
-            **acts.examples.defaultKWArgs(
-                loc0Min=trackSelectorConfig.loc0[0],
-                loc0Max=trackSelectorConfig.loc0[1],
-                loc1Min=trackSelectorConfig.loc1[0],
-                loc1Max=trackSelectorConfig.loc1[1],
-                timeMin=trackSelectorConfig.time[0],
-                timeMax=trackSelectorConfig.time[1],
-                phiMin=trackSelectorConfig.phi[0],
-                phiMax=trackSelectorConfig.phi[1],
-                etaMin=trackSelectorConfig.eta[0],
-                etaMax=trackSelectorConfig.eta[1],
-                absEtaMin=trackSelectorConfig.absEta[0],
-                absEtaMax=trackSelectorConfig.absEta[1],
-                ptMin=trackSelectorConfig.pt[0],
-                ptMax=trackSelectorConfig.pt[1],
-                minMeasurements=trackSelectorConfig.nMeasurementsMin,
-                measurementCounter=trackSelectorConfig.nMeasurementsGroupMin,
-            )
-        )
-        if trackSelectorConfig is not None
-        else None,
         inputMeasurements="measurements",
         inputSourceLinks="sourcelinks",
         inputInitialTrackParameters="estimatedparameters",
@@ -1201,6 +1232,7 @@ def addCKFTracks(
             trackingGeometry, field, customLogLevel()
         ),
         **acts.examples.defaultKWArgs(
+            trackSelectorCfg=trkSelCfg,
             maxSteps=ckfConfig.maxSteps,
         ),
     )

Examples/Python/src/ExampleAlgorithms.cpp

@@ -109,6 +109,10 @@ void addExampleAlgorithms(Context& ctx) {
       ACTS_PYTHON_MEMBER(ptMin);
       ACTS_PYTHON_MEMBER(ptMax);
       ACTS_PYTHON_MEMBER(minMeasurements);
+      ACTS_PYTHON_MEMBER(maxHoles);
+      ACTS_PYTHON_MEMBER(maxOutliers);
+      ACTS_PYTHON_MEMBER(maxSharedHits);
+      ACTS_PYTHON_MEMBER(maxChi2);
       ACTS_PYTHON_MEMBER(measurementCounter);
       ACTS_PYTHON_STRUCT_END();
 
@@ -126,6 +130,8 @@ void addExampleAlgorithms(Context& ctx) {
                    .def(py::init<>())
                    .def(py::init<const Config&>());
 
+      patchKwargsConstructor(c);
+
       c.def_property_readonly("nEtaBins", &EtaBinnedConfig::nEtaBins);
 
       ACTS_PYTHON_STRUCT_BEGIN(c, EtaBinnedConfig);

Examples/Scripts/Python/full_chain_itk.py

@@ -109,10 +109,12 @@
     s,
     trackingGeometry,
     field,
-    TrackSelectorConfig(
-        pt=(1.0 * u.GeV if ttbar_pu200 else 0.0, None),
-        absEta=(None, 4.0),
-        nMeasurementsMin=6,
+    trackSelectorConfig=(
+        # fmt: off
+        TrackSelectorConfig(absEta=(None, 2.0), pt=(0.9 * u.GeV, None), nMeasurementsMin=9, maxHoles=2, maxSharedHits=2),
+        TrackSelectorConfig(absEta=(None, 2.6), pt=(0.4 * u.GeV, None), nMeasurementsMin=8, maxHoles=2, maxSharedHits=2),
+        TrackSelectorConfig(absEta=(None, 4.0), pt=(0.4 * u.GeV, None), nMeasurementsMin=7, maxHoles=2, maxSharedHits=2),
+        # fmt: on
     ),
     outputDirRoot=outputDir,
 )