feat: Custom z bin search order (#1169)

This PR introduces a custom z looping feature, 
which allows us to change the order in which `BinnedSPGroup` search for the Space Points.

In the Main branch the groups are sorted in **phi** and **z** in ascending order. But in Athena **z** follows the order within the vector `{1, 2, 3, 4, 11, 10, 9, 8, 6, 5, 7}`

This PR introduces this same custom looping approach with the vector `zBinsCustomLooping`. If it is not defined, the previous behaviour is preserved.

@noemina @paulgessinger @robertlangenberg

Core/include/Acts/Seeding/BinnedSPGroup.hpp

@@ -147,17 +147,22 @@ class BinnedSPGroupIterator {
   BinnedSPGroupIterator& operator++() {
     if (zIndex < phiZbins[1]) {
       zIndex++;
-
     } else {
       zIndex = 1;
       phiIndex++;
     }
+
+    size_t this_zIndex = zIndex;
+    if (not customZorder.empty())
+      this_zIndex = customZorder.at(this_zIndex - 1);
+
     // set current & neighbor bins only if bin indices valid
     if (phiIndex <= phiZbins[0] && zIndex <= phiZbins[1]) {
-      currentBin =
-          NeighborhoodVector{grid->globalBinFromLocalBins({phiIndex, zIndex})};
-      bottomBinIndices = m_bottomBinFinder->findBins(phiIndex, zIndex, grid);
-      topBinIndices = m_topBinFinder->findBins(phiIndex, zIndex, grid);
+      currentBin = NeighborhoodVector{
+          grid->globalBinFromLocalBins({phiIndex, this_zIndex})};
+      bottomBinIndices =
+          m_bottomBinFinder->findBins(phiIndex, this_zIndex, grid);
+      topBinIndices = m_topBinFinder->findBins(phiIndex, this_zIndex, grid);
       outputIndex++;
       return *this;
     }
@@ -188,34 +193,50 @@ class BinnedSPGroupIterator {
 
   BinnedSPGroupIterator(const SpacePointGrid<external_spacepoint_t>* spgrid,
                         BinFinder<external_spacepoint_t>* botBinFinder,
-                        BinFinder<external_spacepoint_t>* tBinFinder)
-      : currentBin({spgrid->globalBinFromLocalBins({1, 1})}) {
+                        BinFinder<external_spacepoint_t>* tBinFinder,
+                        boost::container::small_vector<size_t, 20> bins = {}) {
     grid = spgrid;
     m_bottomBinFinder = botBinFinder;
     m_topBinFinder = tBinFinder;
     phiZbins = grid->numLocalBins();
     phiIndex = 1;
     zIndex = 1;
-    outputIndex = 0;
-    bottomBinIndices = m_bottomBinFinder->findBins(phiIndex, zIndex, grid);
-    topBinIndices = m_topBinFinder->findBins(phiIndex, zIndex, grid);
+    customZorder = bins;
+    // if m_bins vector was not defined, use z bin 1 (zIndex) to start the
+    // iterator, otherwise use the first value in m_bins vector
+    size_t this_zIndex = bins.empty() ? zIndex : bins.front();
+    outputIndex = grid->globalBinFromLocalBins({phiIndex, this_zIndex});
+    currentBin = NeighborhoodVector{
+        grid->globalBinFromLocalBins({phiIndex, this_zIndex})};
+    bottomBinIndices = m_bottomBinFinder->findBins(phiIndex, this_zIndex, grid);
+    topBinIndices = m_topBinFinder->findBins(phiIndex, this_zIndex, grid);
   }
 
   BinnedSPGroupIterator(const SpacePointGrid<external_spacepoint_t>* spgrid,
                         BinFinder<external_spacepoint_t>* botBinFinder,
                         BinFinder<external_spacepoint_t>* tBinFinder,
-                        size_t phiInd, size_t zInd)
-      : currentBin({spgrid->globalBinFromLocalBins({phiInd, zInd})}) {
+                        size_t phiInd, size_t zInd,
+                        boost::container::small_vector<size_t, 20> bins = {}) {
     m_bottomBinFinder = botBinFinder;
     m_topBinFinder = tBinFinder;
     grid = spgrid;
     phiIndex = phiInd;
     zIndex = zInd;
     phiZbins = grid->numLocalBins();
-    outputIndex = (phiInd - 1) * phiZbins[1] + zInd - 1;
+    customZorder = bins;
+    // if m_bins vector was not defined, use the next z bin (zInd), otherwise
+    // use the z bin value stored in m_bins vector for a custom order
+    size_t this_zIndex =
+        bins.empty()
+            ? zIndex
+            : (zIndex <= phiZbins[1] ? bins.at(zIndex - 1) : bins.back());
+    outputIndex = grid->globalBinFromLocalBins({phiIndex, this_zIndex});
+    currentBin =
+        NeighborhoodVector(grid->globalBinFromLocalBins({phiInd, this_zIndex}));
     if (phiIndex <= phiZbins[0] && zIndex <= phiZbins[1]) {
-      bottomBinIndices = m_bottomBinFinder->findBins(phiIndex, zIndex, grid);
-      topBinIndices = m_topBinFinder->findBins(phiIndex, zIndex, grid);
+      bottomBinIndices =
+          m_bottomBinFinder->findBins(phiIndex, this_zIndex, grid);
+      topBinIndices = m_topBinFinder->findBins(phiIndex, this_zIndex, grid);
     }
   }
 
@@ -231,6 +252,8 @@ class BinnedSPGroupIterator {
   std::array<long unsigned int, 2ul> phiZbins;
   BinFinder<external_spacepoint_t>* m_bottomBinFinder;
   BinFinder<external_spacepoint_t>* m_topBinFinder;
+  boost::container::small_vector<size_t, 20> customZorder;
+  // 	bool start = true;
 };
 
 /// @c BinnedSPGroup Provides access to begin and end BinnedSPGroupIterator
@@ -255,14 +278,15 @@ class BinnedSPGroup {
 
   BinnedSPGroupIterator<external_spacepoint_t> begin() {
     return BinnedSPGroupIterator<external_spacepoint_t>(
-        m_binnedSP.get(), m_bottomBinFinder.get(), m_topBinFinder.get());
+        m_binnedSP.get(), m_bottomBinFinder.get(), m_topBinFinder.get(),
+        m_bins);
   }
 
   BinnedSPGroupIterator<external_spacepoint_t> end() {
     auto phiZbins = m_binnedSP->numLocalBins();
     return BinnedSPGroupIterator<external_spacepoint_t>(
         m_binnedSP.get(), m_bottomBinFinder.get(), m_topBinFinder.get(),
-        phiZbins[0], phiZbins[1] + 1);
+        phiZbins[0], phiZbins[1] + 1, m_bins);
   }
 
  private:
@@ -273,6 +297,8 @@ class BinnedSPGroup {
   // each bin sorted in r (ascending)
   std::shared_ptr<BinFinder<external_spacepoint_t>> m_topBinFinder;
   std::shared_ptr<BinFinder<external_spacepoint_t>> m_bottomBinFinder;
+
+  boost::container::small_vector<size_t, 20> m_bins;
 };
 
 }  // namespace Acts

Core/include/Acts/Seeding/BinnedSPGroup.ipp

@@ -82,4 +82,6 @@ Acts::BinnedSPGroup<external_spacepoint_t>::BinnedSPGroup(
   m_binnedSP = std::move(grid);
   m_bottomBinFinder = botBinFinder;
   m_topBinFinder = tBinFinder;
+
+  m_bins = _config.zBinsCustomLooping;
 }

Core/include/Acts/Seeding/SeedfinderConfig.hpp

@@ -119,6 +119,8 @@ struct SeedfinderConfig {
   Acts::Vector2 beamPos{0 * Acts::UnitConstants::mm,
                         0 * Acts::UnitConstants::mm};
 
+  boost::container::small_vector<size_t, 20> zBinsCustomLooping = {};
+
   // average radiation lengths of material on the length of a seed. used for
   // scattering.
   // default is 5%

Examples/Algorithms/TrackFinding/src/SeedingAlgorithm.cpp

@@ -87,6 +87,20 @@ ActsExamples::SeedingAlgorithm::SeedingAlgorithm(
     throw std::invalid_argument("Inconsistent config zBinNeighborsBottom");
   }
 
+  if (m_cfg.seedFinderConfig.zBinsCustomLooping.size() != 0) {
+    // check if zBinsCustomLooping contains numbers from 1 to the total number
+    // of bin in zBinEdges
+    for (size_t i = 1; i != m_cfg.gridConfig.zBinEdges.size(); i++) {
+      if (std::find(m_cfg.seedFinderConfig.zBinsCustomLooping.begin(),
+                    m_cfg.seedFinderConfig.zBinsCustomLooping.end(),
+                    i) == m_cfg.seedFinderConfig.zBinsCustomLooping.end()) {
+        throw std::invalid_argument(
+            "Inconsistent config zBinsCustomLooping does not contain the same "
+            "bins as zBinEdges");
+      }
+    }
+  }
+
   m_cfg.seedFinderConfig.seedFilter =
       std::make_unique<Acts::SeedFilter<SimSpacePoint>>(m_cfg.seedFilterConfig);
 }

Examples/Python/src/TrackFinding.cpp

@@ -109,6 +109,7 @@ void addTrackFinding(Context& ctx) {
     ACTS_PYTHON_MEMBER(enableCutsForSortedSP);
     ACTS_PYTHON_MEMBER(interactionPointCut);
     ACTS_PYTHON_MEMBER(zBinEdges);
+    ACTS_PYTHON_MEMBER(zBinsCustomLooping);
     ACTS_PYTHON_MEMBER(rRangeMiddleSP);
     ACTS_PYTHON_MEMBER(useVariableMiddleSPRange);
     ACTS_PYTHON_MEMBER(deltaRMiddleSPRange);