grid_helper.hpp

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

// This file is part of the Acts project.
//
// Copyright (C) 2017-2018 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 <array>
#include <tuple>
#include <utility>
#include "Acts/Utilities/IAxis.hpp"
#include "Acts/Utilities/detail/Axis.hpp"

namespace Acts {

namespace detail {

// This object can be iterated to produce the (ordered) set of global indices
// associated with a neighborhood around a certain point on a grid.
//
// The goal is to emulate the effect of enumerating the global indices into
// an std::set (or into an std::vector that gets subsequently sorted), without
// paying the price of dynamic memory allocation in hot magnetic field
// interpolation code.
//
template <size_t DIM>
class GlobalNeighborHoodIndices {
 public:
  // You can get the local neighbor indices from
  // grid_helper_impl<DIM>::neighborHoodIndices and the number of bins in
  // each direction from grid_helper_impl<DIM>::getNBins.
  GlobalNeighborHoodIndices(
      std::array<NeighborHoodIndices, DIM>& neighborIndices,
      const std::array<size_t, DIM>& nBinsArray)
      : m_localIndices(neighborIndices) {
    if (DIM == 1)
      return;
    size_t globalStride = 1;
    for (long i = DIM - 2; i >= 0; --i) {
      globalStride *= (nBinsArray[i + 1] + 2);
      m_globalStrides[i] = globalStride;
    }
  }

  class iterator {
   public:
    iterator() = default;

    iterator(const GlobalNeighborHoodIndices& parent,
             std::array<NeighborHoodIndices::iterator, DIM>&& localIndicesIter)
        : m_localIndicesIter(std::move(localIndicesIter)), m_parent(&parent) {}

    size_t operator*() const {
      size_t globalIndex = *m_localIndicesIter[DIM - 1];
      if (DIM == 1)
        return globalIndex;
      for (size_t i = 0; i < DIM - 1; ++i) {
        globalIndex += m_parent->m_globalStrides[i] * (*m_localIndicesIter[i]);
      }
      return globalIndex;
    }

    iterator& operator++() {
      const auto& localIndices = m_parent->m_localIndices;

      // Go to the next global index via a lexicographic increment:
      // - Start by incrementing the last local index
      // - If it reaches the end, reset it and increment the previous one...
      for (long i = DIM - 1; i > 0; --i) {
        ++m_localIndicesIter[i];
        if (m_localIndicesIter[i] != localIndices[i].end())
          return *this;
        m_localIndicesIter[i] = localIndices[i].begin();
      }

      // The first index should stay at the end value when it reaches it, so
      // that we know when we've reached the end of iteration.
      ++m_localIndicesIter[0];
      return *this;
    }

    bool operator==(const iterator& it) {
      // We know when we've reached the end, so we don't need an end-iterator.
      // Sadly, in C++, there has to be one. Therefore, we special-case it
      // heavily so that it's super-efficient to create and compare to.
      if (it.m_parent == nullptr) {
        return m_localIndicesIter[0] == m_parent->m_localIndices[0].end();
      } else {
        return m_localIndicesIter == it.m_localIndicesIter;
      }
    }

    bool operator!=(const iterator& it) { return !(*this == it); }

   private:
    std::array<NeighborHoodIndices::iterator, DIM> m_localIndicesIter;
    const GlobalNeighborHoodIndices* m_parent = nullptr;
  };

  iterator begin() const {
    std::array<NeighborHoodIndices::iterator, DIM> localIndicesIter;
    for (size_t i = 0; i < DIM; ++i) {
      localIndicesIter[i] = m_localIndices[i].begin();
    }
    return iterator(*this, std::move(localIndicesIter));
  }

  iterator end() const { return iterator(); }

  // Number of indices that will be produced if this sequence is iterated
  size_t size() const {
    size_t result = m_localIndices[0].size();
    for (size_t i = 1; i < DIM; ++i) {
      result *= m_localIndices[i].size();
    }
    return result;
  }

  // Collect the sequence of indices into an std::vector
  std::vector<size_t> collect() const {
    std::vector<size_t> result;
    result.reserve(this->size());
    for (size_t idx : *this) {
      result.push_back(idx);
    }
    return result;
  }

 private:
  std::array<NeighborHoodIndices, DIM> m_localIndices;
  std::array<size_t, DIM - 1> m_globalStrides;
};

template <size_t N>
struct grid_helper_impl;

template <size_t N>
struct grid_helper_impl {
  template <class... Axes>
  static void getBinCenter(
      std::array<double, sizeof...(Axes)>& center,
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    center.at(N) = std::get<N>(axes).getBinCenter(localIndices.at(N));
    grid_helper_impl<N - 1>::getBinCenter(center, localIndices, axes);
  }

  template <class... Axes>
  static void getGlobalBin(const std::array<size_t, sizeof...(Axes)>& localBins,
                           const std::tuple<Axes...>& axes, size_t& bin,
                           size_t& area) {
    const auto& thisAxis = std::get<N>(axes);
    bin += area * localBins.at(N);
    // make sure to account for under-/overflow bins
    area *= (thisAxis.getNBins() + 2);
    grid_helper_impl<N - 1>::getGlobalBin(localBins, axes, bin, area);
  }

  template <class Point, class... Axes>
  static void getLocalBinIndices(const Point& point,
                                 const std::tuple<Axes...>& axes,
                                 std::array<size_t, sizeof...(Axes)>& indices) {
    const auto& thisAxis = std::get<N>(axes);
    indices.at(N) = thisAxis.getBin(point[N]);
    grid_helper_impl<N - 1>::getLocalBinIndices(point, axes, indices);
  }

  template <class... Axes>
  static void getLocalBinIndices(size_t& bin, const std::tuple<Axes...>& axes,
                                 size_t& area,
                                 std::array<size_t, sizeof...(Axes)>& indices) {
    const auto& thisAxis = std::get<N>(axes);
    // make sure to account for under-/overflow bins
    size_t new_area = area * (thisAxis.getNBins() + 2);
    grid_helper_impl<N - 1>::getLocalBinIndices(bin, axes, new_area, indices);
    indices.at(N) = bin / area;
    bin %= area;
  }

  template <class... Axes>
  static void getLowerLeftBinEdge(
      std::array<double, sizeof...(Axes)>& llEdge,
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    llEdge.at(N) = std::get<N>(axes).getBinLowerBound(localIndices.at(N));
    grid_helper_impl<N - 1>::getLowerLeftBinEdge(llEdge, localIndices, axes);
  }

  template <class... Axes>
  static void getLowerLeftBinIndices(
      std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    localIndices.at(N) = std::get<N>(axes).wrapBin(localIndices.at(N) - 1);
    grid_helper_impl<N - 1>::getLowerLeftBinIndices(localIndices, axes);
  }

  template <class... Axes>
  static void getNBins(const std::tuple<Axes...>& axes,
                       std::array<size_t, sizeof...(Axes)>& nBinsArray) {
    // by convention getNBins does not include under-/overflow bins
    nBinsArray[N] = std::get<N>(axes).getNBins();
    grid_helper_impl<N - 1>::getNBins(axes, nBinsArray);
  }

  template <class... Axes>
  static void getAxes(const std::tuple<Axes...>& axes,
                      std::array<const IAxis*, sizeof...(Axes)>& axesArr) {
    axesArr[N] = static_cast<const IAxis*>(&std::get<N>(axes));
    grid_helper_impl<N - 1>::getAxes(axes, axesArr);
  }

  template <class... Axes>
  static void getUpperRightBinEdge(
      std::array<double, sizeof...(Axes)>& urEdge,
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    urEdge.at(N) = std::get<N>(axes).getBinUpperBound(localIndices.at(N));
    grid_helper_impl<N - 1>::getUpperRightBinEdge(urEdge, localIndices, axes);
  }

  template <class... Axes>
  static void getUpperRightBinIndices(
      std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    localIndices.at(N) = std::get<N>(axes).wrapBin(localIndices.at(N) + 1);
    grid_helper_impl<N - 1>::getUpperRightBinIndices(localIndices, axes);
  }

  template <class... Axes>
  static void getMin(const std::tuple<Axes...>& axes,
                     std::array<double, sizeof...(Axes)>& minArray) {
    minArray[N] = std::get<N>(axes).getMin();
    grid_helper_impl<N - 1>::getMin(axes, minArray);
  }

  template <class... Axes>
  static void getMax(const std::tuple<Axes...>& axes,
                     std::array<double, sizeof...(Axes)>& maxArray) {
    maxArray[N] = std::get<N>(axes).getMax();
    grid_helper_impl<N - 1>::getMax(axes, maxArray);
  }

  template <class... Axes>
  static void getWidth(const std::tuple<Axes...>& axes,
                       std::array<double, sizeof...(Axes)>& widthArray) {
    widthArray[N] = std::get<N>(axes).getBinWidth();
    grid_helper_impl<N - 1>::getWidth(axes, widthArray);
  }

  template <class Point, class... Axes>
  static bool isInside(const Point& position, const std::tuple<Axes...>& axes) {
    bool insideThisAxis = std::get<N>(axes).isInside(position[N]);
    return insideThisAxis && grid_helper_impl<N - 1>::isInside(position, axes);
  }

  template <class... Axes>
  static void neighborHoodIndices(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      std::pair<size_t, size_t> sizes, const std::tuple<Axes...>& axes,
      std::array<NeighborHoodIndices, sizeof...(Axes)>& neighborIndices) {
    // ask n-th axis
    size_t locIdx = localIndices.at(N);
    NeighborHoodIndices locNeighbors =
        std::get<N>(axes).neighborHoodIndices(locIdx, sizes);
    neighborIndices.at(N) = locNeighbors;

    grid_helper_impl<N - 1>::neighborHoodIndices(localIndices, sizes, axes,
                                                 neighborIndices);
  }

  template <class... Axes>
  static void exteriorBinIndices(std::array<size_t, sizeof...(Axes)>& idx,
                                 std::array<bool, sizeof...(Axes)> isExterior,
                                 std::set<size_t>& combinations,
                                 const std::tuple<Axes...>& axes) {
    // iterate over this axis' bins, remembering which bins are exterior
    for (size_t i = 0; i < std::get<N>(axes).getNBins() + 2; ++i) {
      idx.at(N) = i;
      isExterior.at(N) = (i == 0) || (i == std::get<N>(axes).getNBins() + 1);
      // vary other axes recursively
      grid_helper_impl<N - 1>::exteriorBinIndices(idx, isExterior, combinations,
                                                  axes);
    }
  }
};

template <>
struct grid_helper_impl<0u> {
  template <class... Axes>
  static void getBinCenter(
      std::array<double, sizeof...(Axes)>& center,
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    center.at(0u) = std::get<0u>(axes).getBinCenter(localIndices.at(0u));
  }

  template <class... Axes>
  static void getGlobalBin(const std::array<size_t, sizeof...(Axes)>& localBins,
                           const std::tuple<Axes...>& /*axes*/, size_t& bin,
                           size_t& area) {
    bin += area * localBins.at(0u);
  }

  template <class Point, class... Axes>
  static void getLocalBinIndices(const Point& point,
                                 const std::tuple<Axes...>& axes,
                                 std::array<size_t, sizeof...(Axes)>& indices) {
    const auto& thisAxis = std::get<0u>(axes);
    indices.at(0u) = thisAxis.getBin(point[0u]);
  }

  template <class... Axes>
  static void getLocalBinIndices(size_t& bin,
                                 const std::tuple<Axes...>& /*axes*/,
                                 size_t& area,
                                 std::array<size_t, sizeof...(Axes)>& indices) {
    // make sure to account for under-/overflow bins
    indices.at(0u) = bin / area;
    bin %= area;
  }

  template <class... Axes>
  static void getLowerLeftBinEdge(
      std::array<double, sizeof...(Axes)>& llEdge,
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    llEdge.at(0u) = std::get<0u>(axes).getBinLowerBound(localIndices.at(0u));
  }

  template <class... Axes>
  static void getLowerLeftBinIndices(
      std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    localIndices.at(0u) = std::get<0u>(axes).wrapBin(localIndices.at(0u) - 1);
  }

  template <class... Axes>
  static void getNBins(const std::tuple<Axes...>& axes,
                       std::array<size_t, sizeof...(Axes)>& nBinsArray) {
    // by convention getNBins does not include under-/overflow bins
    nBinsArray[0u] = std::get<0u>(axes).getNBins();
  }

  template <class... Axes>
  static void getAxes(const std::tuple<Axes...>& axes,
                      std::array<const IAxis*, sizeof...(Axes)>& axesArr) {
    axesArr[0u] = static_cast<const IAxis*>(&std::get<0u>(axes));
  }

  template <class... Axes>
  static void getUpperRightBinEdge(
      std::array<double, sizeof...(Axes)>& urEdge,
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    urEdge.at(0u) = std::get<0u>(axes).getBinUpperBound(localIndices.at(0u));
  }

  template <class... Axes>
  static void getUpperRightBinIndices(
      std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    localIndices.at(0u) = std::get<0u>(axes).wrapBin(localIndices.at(0u) + 1);
  }

  template <class... Axes>
  static void getMin(const std::tuple<Axes...>& axes,
                     std::array<double, sizeof...(Axes)>& minArray) {
    minArray[0u] = std::get<0u>(axes).getMin();
  }

  template <class... Axes>
  static void getMax(const std::tuple<Axes...>& axes,
                     std::array<double, sizeof...(Axes)>& maxArray) {
    maxArray[0u] = std::get<0u>(axes).getMax();
  }

  template <class... Axes>
  static void getWidth(const std::tuple<Axes...>& axes,
                       std::array<double, sizeof...(Axes)>& widthArray) {
    widthArray[0u] = std::get<0u>(axes).getBinWidth();
  }

  template <class Point, class... Axes>
  static bool isInside(const Point& position, const std::tuple<Axes...>& axes) {
    return std::get<0u>(axes).isInside(position[0u]);
  }

  template <class... Axes>
  static void neighborHoodIndices(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      std::pair<size_t, size_t> sizes, const std::tuple<Axes...>& axes,
      std::array<NeighborHoodIndices, sizeof...(Axes)>& neighborIndices) {
    // ask 0-th axis
    size_t locIdx = localIndices.at(0u);
    NeighborHoodIndices locNeighbors =
        std::get<0u>(axes).neighborHoodIndices(locIdx, sizes);
    neighborIndices.at(0u) = locNeighbors;
  }

  template <class... Axes>
  static void exteriorBinIndices(std::array<size_t, sizeof...(Axes)>& idx,
                                 std::array<bool, sizeof...(Axes)> isExterior,
                                 std::set<size_t>& combinations,
                                 const std::tuple<Axes...>& axes) {
    // For each exterior bin on this axis, we will do this
    auto recordExteriorBin = [&](size_t i) {
      idx.at(0u) = i;
      // at this point, combinations are complete: save the global bin
      size_t bin = 0, area = 1;
      grid_helper_impl<sizeof...(Axes) - 1>::getGlobalBin(idx, axes, bin, area);
      combinations.insert(bin);
    };

    // The first and last bins on this axis are exterior by definition
    for (size_t i :
         {static_cast<size_t>(0), std::get<0u>(axes).getNBins() + 1}) {
      recordExteriorBin(i);
    }

    // If no other axis is on an exterior index, stop here
    bool otherAxisExterior = false;
    for (size_t N = 1; N < sizeof...(Axes); ++N) {
      otherAxisExterior = otherAxisExterior | isExterior[N];
    }
    if (!otherAxisExterior) {
      return;
    }

    // Otherwise, we're on a grid border: iterate over all the other indices
    for (size_t i = 1; i <= std::get<0u>(axes).getNBins(); ++i) {
      recordExteriorBin(i);
    }
  }
};

struct grid_helper {
  template <class... Axes>
  static GlobalNeighborHoodIndices<sizeof...(Axes)> closestPointsIndices(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    // get neighboring bins, but only increment.
    return neighborHoodIndices(localIndices, std::make_pair(0, 1), axes);
  }

  template <class... Axes>
  static std::array<double, sizeof...(Axes)> getBinCenter(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    std::array<double, sizeof...(Axes)> center;
    constexpr size_t MAX = sizeof...(Axes) - 1;
    grid_helper_impl<MAX>::getBinCenter(center, localIndices, axes);

    return center;
  }

  template <class... Axes>
  static size_t getGlobalBin(
      const std::array<size_t, sizeof...(Axes)>& localBins,
      const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;
    size_t area = 1;
    size_t bin = 0;

    grid_helper_impl<MAX>::getGlobalBin(localBins, axes, bin, area);

    return bin;
  }

  template <class Point, class... Axes>
  static std::array<size_t, sizeof...(Axes)> getLocalBinIndices(
      const Point& point, const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;
    std::array<size_t, sizeof...(Axes)> indices;

    grid_helper_impl<MAX>::getLocalBinIndices(point, axes, indices);

    return indices;
  }

  template <class... Axes>
  static std::array<size_t, sizeof...(Axes)> getLocalBinIndices(
      size_t bin, const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;
    size_t area = 1;
    std::array<size_t, sizeof...(Axes)> indices;

    grid_helper_impl<MAX>::getLocalBinIndices(bin, axes, area, indices);

    return indices;
  }

  template <class... Axes>
  static std::array<double, sizeof...(Axes)> getLowerLeftBinEdge(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    std::array<double, sizeof...(Axes)> llEdge;
    constexpr size_t MAX = sizeof...(Axes) - 1;
    grid_helper_impl<MAX>::getLowerLeftBinEdge(llEdge, localIndices, axes);

    return llEdge;
  }

  template <class... Axes>
  static std::array<size_t, sizeof...(Axes)> getLowerLeftBinIndices(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;
    auto llIndices = localIndices;
    grid_helper_impl<MAX>::getLowerLeftBinIndices(llIndices, axes);

    return llIndices;
  }

  template <class... Axes>
  static std::array<size_t, sizeof...(Axes)> getNBins(
      const std::tuple<Axes...>& axes) {
    std::array<size_t, sizeof...(Axes)> nBinsArray;
    grid_helper_impl<sizeof...(Axes) - 1>::getNBins(axes, nBinsArray);
    return nBinsArray;
  }

  template <class... Axes>
  static std::array<const IAxis*, sizeof...(Axes)> getAxes(
      const std::tuple<Axes...>& axes) {
    std::array<const IAxis*, sizeof...(Axes)> arr;
    grid_helper_impl<sizeof...(Axes) - 1>::getAxes(axes, arr);
    return arr;
  }

  template <class... Axes>
  static std::array<double, sizeof...(Axes)> getUpperRightBinEdge(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    std::array<double, sizeof...(Axes)> urEdge;
    constexpr size_t MAX = sizeof...(Axes) - 1;
    grid_helper_impl<MAX>::getUpperRightBinEdge(urEdge, localIndices, axes);

    return urEdge;
  }

  template <class... Axes>
  static std::array<size_t, sizeof...(Axes)> getUpperRightBinIndices(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;
    auto urIndices = localIndices;
    grid_helper_impl<MAX>::getUpperRightBinIndices(urIndices, axes);

    return urIndices;
  }

  template <class... Axes>
  static std::array<double, sizeof...(Axes)> getMin(
      const std::tuple<Axes...>& axes) {
    std::array<double, sizeof...(Axes)> minArray;
    grid_helper_impl<sizeof...(Axes) - 1>::getMin(axes, minArray);
    return minArray;
  }

  template <class... Axes>
  static std::array<double, sizeof...(Axes)> getMax(
      const std::tuple<Axes...>& axes) {
    std::array<double, sizeof...(Axes)> maxArray;
    grid_helper_impl<sizeof...(Axes) - 1>::getMax(axes, maxArray);
    return maxArray;
  }

  template <class... Axes>
  static std::array<double, sizeof...(Axes)> getWidth(
      const std::tuple<Axes...>& axes) {
    std::array<double, sizeof...(Axes)> widthArray;
    grid_helper_impl<sizeof...(Axes) - 1>::getWidth(axes, widthArray);
    return widthArray;
  }

  template <class... Axes>
  static GlobalNeighborHoodIndices<sizeof...(Axes)> neighborHoodIndices(
      const std::array<size_t, sizeof...(Axes)>& localIndices,
      std::pair<size_t, size_t> sizes, const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;

    // length N array which contains local neighbors based on size par
    std::array<NeighborHoodIndices, sizeof...(Axes)> neighborIndices;
    // get local bin indices for neighboring bins
    grid_helper_impl<MAX>::neighborHoodIndices(localIndices, sizes, axes,
                                               neighborIndices);

    // Query the number of bins
    std::array<size_t, sizeof...(Axes)> nBinsArray = getNBins(axes);

    // Produce iterator of global indices
    return GlobalNeighborHoodIndices(neighborIndices, nBinsArray);
  }

  template <class... Axes>
  static GlobalNeighborHoodIndices<sizeof...(Axes)> neighborHoodIndices(
      const std::array<size_t, sizeof...(Axes)>& localIndices, size_t size,
      const std::tuple<Axes...>& axes) {
    return neighborHoodIndices(localIndices, std::make_pair(size, size), axes);
  }

  template <class... Axes>
  static std::set<size_t> exteriorBinIndices(const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;

    std::array<size_t, sizeof...(Axes)> idx;
    std::array<bool, sizeof...(Axes)> isExterior;
    std::set<size_t> combinations;
    grid_helper_impl<MAX>::exteriorBinIndices(idx, isExterior, combinations,
                                              axes);

    return combinations;
  }

  template <class Point, class... Axes>
  static bool isInside(const Point& position, const std::tuple<Axes...>& axes) {
    constexpr size_t MAX = sizeof...(Axes) - 1;
    return grid_helper_impl<MAX>::isInside(position, axes);
  }
};

}  // namespace detail

}  // namespace Acts