Axis.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file Axis.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 "Acts/Utilities/IAxis.hpp"
#include "Acts/Utilities/detail/AxisFwd.hpp"

#include <algorithm>
#include <cmath>
#include <vector>

namespace Acts {

namespace detail {

// This object can be iterated to produce up to two sequences of integer
// indices, corresponding to the half-open integer ranges [begin1, end1[ and
// [begin2, end2[.
//
// The goal is to emulate the effect of enumerating a range of neighbor
// indices on an axis (which may go out of bounds and wrap around since we
// have AxisBoundaryType::Closed), inserting them into an std::vector, and
// discarding duplicates, without paying the price of duplicate removal
// and dynamic memory allocation in hot magnetic field interpolation code.
//
class NeighborHoodIndices {
 public:
  NeighborHoodIndices() = default;

  NeighborHoodIndices(size_t begin, size_t end)
      : m_begin1(begin), m_end1(end), m_begin2(end), m_end2(end) {}

  NeighborHoodIndices(size_t begin1, size_t end1, size_t begin2, size_t end2)
      : m_begin1(begin1), m_end1(end1), m_begin2(begin2), m_end2(end2) {}

  class iterator {
   public:
    iterator() = default;

    // Specialized constructor for end() iterator
    iterator(size_t current) : m_current(current), m_wrapped(true) {}

    iterator(size_t begin1, size_t end1, size_t begin2)
        : m_current(begin1),
          m_end1(end1),
          m_begin2(begin2),
          m_wrapped(begin1 == begin2) {}

    size_t operator*() const { return m_current; }

    iterator& operator++() {
      ++m_current;
      if (m_current == m_end1) {
        m_current = m_begin2;
        m_wrapped = true;
      }
      return *this;
    }

    bool operator==(const iterator& it) const {
      return (m_current == it.m_current) && (m_wrapped == it.m_wrapped);
    }

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

   private:
    size_t m_current = 0, m_end1 = 0, m_begin2 = 0;
    bool m_wrapped = false;
  };

  iterator begin() const { return iterator(m_begin1, m_end1, m_begin2); }

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

  // Number of indices that will be produced if this sequence is iterated
  size_t size() const { return (m_end1 - m_begin1) + (m_end2 - m_begin2); }

  // 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:
  size_t m_begin1 = 0, m_end1 = 0, m_begin2 = 0, m_end2 = 0;
};

template <AxisBoundaryType bdt>
class Axis<AxisType::Equidistant, bdt> final : public IAxis {
 public:
  Axis(ActsScalar xmin, ActsScalar xmax, size_t nBins)
      : m_min(xmin),
        m_max(xmax),
        m_width((xmax - xmin) / nBins),
        m_bins(nBins) {}

  bool isEquidistant() const override { return true; }

  bool isVariable() const override { return false; }

  AxisBoundaryType getBoundaryType() const override { return bdt; }

  NeighborHoodIndices neighborHoodIndices(size_t idx, size_t size = 1) const {
    return neighborHoodIndices(idx, std::make_pair(-size, size));
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Open, int> = 0>
  NeighborHoodIndices neighborHoodIndices(size_t idx,
                                          std::pair<int, int> sizes = {
                                              -1, 1}) const {
    constexpr int min = 0;
    const int max = getNBins() + 1;
    const int itmin = std::clamp(static_cast<int>(idx + sizes.first), min, max);
    const int itmax =
        std::clamp(static_cast<int>(idx + sizes.second), min, max);
    return NeighborHoodIndices(itmin, itmax + 1);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Bound, int> = 0>
  NeighborHoodIndices neighborHoodIndices(size_t idx,
                                          std::pair<int, int> sizes = {
                                              -1, 1}) const {
    if (idx <= 0 || idx >= (getNBins() + 1)) {
      return NeighborHoodIndices();
    }
    constexpr int min = 1;
    const int max = getNBins();
    const int itmin = std::clamp(static_cast<int>(idx) + sizes.first, min, max);
    const int itmax =
        std::clamp(static_cast<int>(idx) + sizes.second, min, max);
    return NeighborHoodIndices(itmin, itmax + 1);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Closed, int> = 0>
  NeighborHoodIndices neighborHoodIndices(size_t idx,
                                          std::pair<int, int> sizes = {
                                              -1, 1}) const {
    // Handle invalid indices
    if (idx <= 0 || idx >= (getNBins() + 1)) {
      return NeighborHoodIndices();
    }

    // Handle corner case where user requests more neighbours than the number
    // of bins on the axis. All bins are returned in this case.

    const int max = getNBins();
    sizes.first = std::clamp(sizes.first, -max, max);
    sizes.second = std::clamp(sizes.second, -max, max);
    if (std::abs(sizes.first - sizes.second) >= max) {
      sizes.first = 1 - idx;
      sizes.second = max - idx;
    }

    // If the entire index range is not covered, we must wrap the range of
    // targeted neighbor indices into the range of valid bin indices. This may
    // split the range of neighbor indices in two parts:
    //
    // Before wraparound - [        XXXXX]XXX
    // After wraparound  - [ XXXX   XXXX ]
    //
    const int itmin = idx + sizes.first;
    const int itmax = idx + sizes.second;
    const size_t itfirst = wrapBin(itmin);
    const size_t itlast = wrapBin(itmax);
    if (itfirst <= itlast) {
      return NeighborHoodIndices(itfirst, itlast + 1);
    } else {
      return NeighborHoodIndices(itfirst, max + 1, 1, itlast + 1);
    }
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Open, int> = 0>
  size_t wrapBin(int bin) const {
    return std::max(std::min(bin, static_cast<int>(getNBins()) + 1), 0);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Bound, int> = 0>
  size_t wrapBin(int bin) const {
    return std::max(std::min(bin, static_cast<int>(getNBins())), 1);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Closed, int> = 0>
  size_t wrapBin(int bin) const {
    const int w = getNBins();
    return 1 + (w + ((bin - 1) % w)) % w;
    // return int(bin<1)*w - int(bin>w)*w + bin;
  }

  size_t getBin(ActsScalar x) const {
    return wrapBin(
        static_cast<int>(std::floor((x - getMin()) / getBinWidth()) + 1));
  }

  ActsScalar getBinWidth(size_t /*bin*/ = 0) const { return m_width; }

  ActsScalar getBinLowerBound(size_t bin) const {
    return getMin() + (bin - 1) * getBinWidth();
  }

  ActsScalar getBinUpperBound(size_t bin) const {
    return getMin() + bin * getBinWidth();
  }

  ActsScalar getBinCenter(size_t bin) const {
    return getMin() + (bin - 0.5) * getBinWidth();
  }

  ActsScalar getMax() const override { return m_max; }

  ActsScalar getMin() const override { return m_min; }

  size_t getNBins() const override { return m_bins; }

  bool isInside(ActsScalar x) const { return (m_min <= x) && (x < m_max); }

  std::vector<ActsScalar> getBinEdges() const override {
    std::vector<ActsScalar> binEdges;
    for (size_t i = 1; i <= m_bins; i++) {
      binEdges.push_back(getBinLowerBound(i));
    }
    binEdges.push_back(getBinUpperBound(m_bins));
    return binEdges;
  }

 private:
  ActsScalar m_min;
  ActsScalar m_max;
  ActsScalar m_width;
  size_t m_bins;
};

template <AxisBoundaryType bdt>
class Axis<AxisType::Variable, bdt> final : public IAxis {
 public:
  Axis(std::vector<ActsScalar> binEdges) : m_binEdges(std::move(binEdges)) {}

  bool isEquidistant() const override { return false; }

  bool isVariable() const override { return true; }

  AxisBoundaryType getBoundaryType() const override { return bdt; }

  NeighborHoodIndices neighborHoodIndices(size_t idx, size_t size = 1) const {
    return neighborHoodIndices(idx, std::make_pair(-size, size));
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Open, int> = 0>
  NeighborHoodIndices neighborHoodIndices(size_t idx,
                                          std::pair<int, int> sizes = {
                                              -1, 1}) const {
    constexpr int min = 0;
    const int max = getNBins() + 1;
    const int itmin = std::max(min, static_cast<int>(idx) + sizes.first);
    const int itmax = std::min(max, static_cast<int>(idx) + sizes.second);
    return NeighborHoodIndices(itmin, itmax + 1);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Bound, int> = 0>
  NeighborHoodIndices neighborHoodIndices(size_t idx,
                                          std::pair<int, int> sizes = {
                                              -1, 1}) const {
    if (idx <= 0 || idx >= (getNBins() + 1)) {
      return NeighborHoodIndices();
    }
    constexpr int min = 1;
    const int max = getNBins();
    const int itmin = std::max(min, static_cast<int>(idx) + sizes.first);
    const int itmax = std::min(max, static_cast<int>(idx) + sizes.second);
    return NeighborHoodIndices(itmin, itmax + 1);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Closed, int> = 0>
  NeighborHoodIndices neighborHoodIndices(size_t idx,
                                          std::pair<int, int> sizes = {
                                              -1, 1}) const {
    // Handle invalid indices
    if (idx <= 0 || idx >= (getNBins() + 1)) {
      return NeighborHoodIndices();
    }

    // Handle corner case where user requests more neighbours than the number
    // of bins on the axis. All bins are returned in this case

    const int max = getNBins();
    sizes.first = std::clamp(sizes.first, -max, max);
    sizes.second = std::clamp(sizes.second, -max, max);
    if (std::abs(sizes.first - sizes.second) >= max) {
      sizes.first = 1 - idx;
      sizes.second = max - idx;
    }

    // If the entire index range is not covered, we must wrap the range of
    // targeted neighbor indices into the range of valid bin indices. This may
    // split the range of neighbor indices in two parts:
    //
    // Before wraparound - [        XXXXX]XXX
    // After wraparound  - [ XXXX   XXXX ]
    //
    const int itmin = idx + sizes.first;
    const int itmax = idx + sizes.second;
    const size_t itfirst = wrapBin(itmin);
    const size_t itlast = wrapBin(itmax);
    if (itfirst <= itlast) {
      return NeighborHoodIndices(itfirst, itlast + 1);
    } else {
      return NeighborHoodIndices(itfirst, max + 1, 1, itlast + 1);
    }
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Open, int> = 0>
  size_t wrapBin(int bin) const {
    return std::max(std::min(bin, static_cast<int>(getNBins()) + 1), 0);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Bound, int> = 0>
  size_t wrapBin(int bin) const {
    return std::max(std::min(bin, static_cast<int>(getNBins())), 1);
  }

  template <AxisBoundaryType T = bdt,
            std::enable_if_t<T == AxisBoundaryType::Closed, int> = 0>
  size_t wrapBin(int bin) const {
    const int w = getNBins();
    return 1 + (w + ((bin - 1) % w)) % w;
    // return int(bin<1)*w - int(bin>w)*w + bin;
  }

  size_t getBin(ActsScalar x) const {
    const auto it =
        std::upper_bound(std::begin(m_binEdges), std::end(m_binEdges), x);
    return wrapBin(std::distance(std::begin(m_binEdges), it));
  }

  ActsScalar getBinWidth(size_t bin) const {
    return m_binEdges.at(bin) - m_binEdges.at(bin - 1);
  }

  ActsScalar getBinLowerBound(size_t bin) const {
    return m_binEdges.at(bin - 1);
  }

  ActsScalar getBinUpperBound(size_t bin) const { return m_binEdges.at(bin); }

  ActsScalar getBinCenter(size_t bin) const {
    return 0.5 * (getBinLowerBound(bin) + getBinUpperBound(bin));
  }

  ActsScalar getMax() const override { return m_binEdges.back(); }

  ActsScalar getMin() const override { return m_binEdges.front(); }

  size_t getNBins() const override { return m_binEdges.size() - 1; }

  bool isInside(ActsScalar x) const {
    return (m_binEdges.front() <= x) && (x < m_binEdges.back());
  }

  std::vector<ActsScalar> getBinEdges() const override { return m_binEdges; }

 private:
  std::vector<ActsScalar> m_binEdges;
};
}  // namespace detail

}  // namespace Acts