traccc/utils/subspace.hpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2023-2024 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 #pragma once
0009
0010 // Project include(s).
0011 #include "traccc/definitions/primitives.hpp"
0012 #include "traccc/definitions/qualifiers.hpp"
0013 #include "traccc/definitions/track_parametrization.hpp"
0014
0015 // System include(s).
0016 #include <array>
0017
0018 namespace traccc {
0019
0020 template <typename algebra_t, detray::dindex_type<algebra_t> kFullSize,
0021           detray::dindex_type<algebra_t> kSize = 2u>
0022 struct subspace {
0023   static_assert(1u <= kSize, "Subspace size must be at least 1");
0024   static_assert(kSize <= kFullSize,
0025                 "Subspace can only be as large as the full space");
0026
0027  public:
0028   // Type declarations
0029   using size_type = detray::dindex_type<algebra_t>;
0030
0031   // Define the number of bits we use to represent a single index;
0032   // currently, we hardcode to be three plus one sign bit so that the full
0033   // size of the subspace is limited to 8. Note that we also reserve one
0034   // element (all ones) as invalid.
0035   //
0036   // TODO: Find a nicer, automated solution to this.
0037   static constexpr size_type BITS_PER_INDEX = 4;
0038   static_assert(BITS_PER_INDEX >= 2);
0039
0040   // Compute the total number of bits necessary to represent a subspace.
0041   static constexpr size_type TOTAL_BITS = BITS_PER_INDEX * kSize;
0042   static_assert(TOTAL_BITS <= 64,
0043                 "Subspaces with more than 64 necessary bits are not supported");
0044   using axes_type = std::conditional_t<TOTAL_BITS <= 32, std::uint_least32_t,
0045                                        std::uint_least64_t>;
0046
0047   static constexpr axes_type SIGN_BIT_MASK = 1 << (BITS_PER_INDEX - 1);
0048   static constexpr axes_type VALUE_BITS_MASK = (1 << (BITS_PER_INDEX - 1)) - 1;
0049   static constexpr axes_type ELEMENT_BITS_MASK = (1 << BITS_PER_INDEX) - 1;
0050   static_assert(((1 << (BITS_PER_INDEX - 1)) - 1) >= kFullSize);
0051   static_assert(ELEMENT_BITS_MASK == (SIGN_BIT_MASK | VALUE_BITS_MASK));
0052
0053   template <size_type ROWS, size_type COLS>
0054   using matrix_type = detray::dmatrix<algebra_t, ROWS, COLS>;
0055
0056   using subspace_vector = matrix_type<kSize, 1u>;
0057   using fullspace_vector = matrix_type<kFullSize, 1u>;
0058   using projection_matrix = matrix_type<kSize, kFullSize>;
0059   using expansion_matrix = matrix_type<kFullSize, kSize>;
0060
0061   static constexpr size_type size = kSize;
0062   static constexpr size_type fullSize = kFullSize;
0063
0064   /// Construct from a container of axis indices.
0065   ///
0066   /// @param indices Unique, ordered indices
0067   template <typename SIZE_TYPE>
0068   TRACCC_HOST_DEVICE constexpr subspace(
0069       const std::array<SIZE_TYPE, kSize>& indices)
0070       : m_axes(0) {
0071     for (size_type i = 0u; i < kSize; ++i) {
0072       assert((indices[i] < kFullSize) and
0073              "Axis indices must be within the full space");
0074       set_index(i, static_cast<axes_type>(indices[i]));
0075     }
0076   }
0077
0078   /// Get the projection index for a given element of the subspace.
0079   TRACCC_HOST_DEVICE
0080   constexpr size_type get_index(const size_type& i) const {
0081     const size_type sr = i * BITS_PER_INDEX;
0082     return (m_axes >> sr) & VALUE_BITS_MASK;
0083   }
0084
0085   /// Check whether a given element of the subspace is valid or not.
0086   ///
0087   /// Invalid elements are not projected at all and result in an empty
0088   /// row or column in the projection and expansion matrix.
0089   TRACCC_HOST_DEVICE
0090   constexpr bool get_valid(const size_type& i) const {
0091     const size_type sr = i * BITS_PER_INDEX;
0092     return ((m_axes >> sr) & VALUE_BITS_MASK) != VALUE_BITS_MASK;
0093   }
0094
0095   /// Retrieve the sign of an element of the subspace; if a true value is
0096   /// returned, the element will be projected in the negative.
0097   TRACCC_HOST_DEVICE
0098   constexpr bool get_sign(const size_type& i) const {
0099     const size_type sr = i * BITS_PER_INDEX;
0100     return (m_axes >> sr) & SIGN_BIT_MASK;
0101   }
0102
0103   /// Set the projection index for a given element.
0104   TRACCC_HOST_DEVICE
0105   constexpr void set_index(const size_type& i, const axes_type& j) {
0106     assert(j == (j & ELEMENT_BITS_MASK));
0107     const size_type sl = i * BITS_PER_INDEX;
0108     m_axes = (m_axes & static_cast<axes_type>(~(VALUE_BITS_MASK << sl))) |
0109              static_cast<axes_type>(j << sl);
0110   }
0111
0112   /// Set an element to be invalid so that it is not projected.
0113   TRACCC_HOST_DEVICE
0114   constexpr void set_invalid(const size_type& i) {
0115     // HACK: For reasons not understood by the author, the
0116     // `VALUE_BITS_MASK` variable cannot be used here.
0117     set_index(i, ((1 << (BITS_PER_INDEX - 1)) - 1));
0118   }
0119
0120   /// Set the sign of an element, where trueish values indicate negative
0121   /// projection and falseish values indicate positive values.
0122   TRACCC_HOST_DEVICE
0123   constexpr void set_sign(const size_type& i, const bool s) {
0124     const size_type sl = i * BITS_PER_INDEX;
0125     m_axes = (m_axes & static_cast<axes_type>(~(SIGN_BIT_MASK << sl))) |
0126              static_cast<axes_type>((s ? SIGN_BIT_MASK : 0) << sl);
0127   }
0128
0129   /// Projection matrix that maps from the full space into the subspace.
0130   template <size_type D>
0131   TRACCC_HOST_DEVICE auto projector() const -> matrix_type<D, kFullSize> {
0132     static_assert(D <= kSize,
0133                   "The dimension of projection should be smaller than "
0134                   "subspace dimension");
0135
0136     auto proj = matrix::zero<matrix_type<D, kFullSize>>();
0137
0138     for (size_type i = 0u; i < D; ++i) {
0139       if (get_index(i) < kFullSize) {
0140         getter::element(proj, i, get_index(i)) = (get_sign(i) ? -1.f : 1.f);
0141       }
0142     }
0143     return proj;
0144   }
0145
0146   /// Expansion matrix that maps from the subspace into the full space.
0147   template <size_type D>
0148   TRACCC_HOST_DEVICE auto expander() const -> matrix_type<kFullSize, D> {
0149     static_assert(D <= kSize,
0150                   "The dimension of projection should be smaller than "
0151                   "subspace dimension");
0152
0153     auto expn = matrix::zero<matrix_type<kFullSize, D>>();
0154
0155     for (size_type i = 0u; i < kSize; ++i) {
0156       if (get_index(i) < kFullSize) {
0157         getter::element(expn, get_index(i), i) = (get_sign(i) ? -1.f : 1.f);
0158       }
0159     }
0160     return expn;
0161   }
0162
0163  private:
0164   axes_type m_axes;
0165 };
0166
0167 }  // namespace traccc