edm4hep/utils/cov_matrix_utils.h

0001 #ifndef EDM4HEP_UTILS_COVMATRIXUTILS_H
0002 #define EDM4HEP_UTILS_COVMATRIXUTILS_H
0003
0004 #include <array>
0005 #include <stdexcept>
0006 #include <string>
0007 #include <type_traits>
0008
0009 namespace edm4hep {
0010
0011 namespace utils {
0012
0013   namespace detail {
0014     // From c++23 this is functionality offered by the STL
0015 #if __cpp_lib_to_underlying
0016     using to_index = std::to_underlying;
0017 #else
0018     // Otherwise it is simple enough to roll our own
0019     template <typename E>
0020     constexpr auto to_index(E e) {
0021       return static_cast<std::underlying_type_t<E>>(e);
0022     }
0023 #endif
0024
0025     /// Cast an index to an enum value
0026     template <typename DimEnum>
0027     constexpr DimEnum to_enum(DimEnum index) {
0028       return static_cast<DimEnum>(index);
0029     }
0030
0031     /// Need a constexpr swap for integers before c++20
0032     constexpr void swap(int& i, int& j) {
0033       int tmp = j;
0034       j = i;
0035       i = tmp;
0036     }
0037
0038     /**
0039      * Get the dimension of the covariance matrix from the size of the 1D array in
0040      * which it is stored
0041      *
0042      * **NOTE: to avoid having to do a square root operation and in order to keep
0043      * this constexpr this is currently only implemented for storage sizes up to
0044      * 21 (corresponding to covariance matrix dimensions of 6 x 6)**. This
0045      * function is intended to be called in constexpr or immediate contexts in
0046      * order to fail at compilation already for invalid values of N.
0047      *
0048      * @param N the size of the 1D storage
0049      *
0050      * @returns the dimension of the covariance matrix
0051      */
0052     consteval std::size_t get_cov_dim(std::size_t N) {
0053       switch (N) {
0054       case 21:
0055         return 6;
0056       case 15:
0057         return 5;
0058       case 10:
0059         return 4;
0060       case 6:
0061         return 3;
0062       case 3:
0063         return 2;
0064       case 1:
0065         return 1;
0066       }
0067
0068       // We simply use throwing an exception to make compilation fail in constexpr
0069       // cases.
0070       throw std::invalid_argument(
0071           "Not a valid size for a covariance matrix stored in lower triangular form (N = " + std::to_string(N) + ")");
0072     }
0073
0074     /**
0075      * Transform a 2D index to a 1D index in lower triangular storage.
0076      *
0077      * @param i row index
0078      * @param j column index
0079      *
0080      * @returns the index into the 1D storage
0081      */
0082     constexpr int to_lower_tri(int i, int j) {
0083       if (i < j) {
0084         detail::swap(i, j);
0085       }
0086       return i * (i + 1) / 2 + j;
0087     }
0088   } // namespace detail
0089
0090   template <typename DimEnum, typename Scalar, std::size_t N>
0091   constexpr Scalar get_cov_value(const std::array<Scalar, N>& cov, DimEnum dimI, DimEnum dimJ) {
0092     const auto i = detail::to_index(dimI);
0093     const auto j = detail::to_index(dimJ);
0094
0095     constexpr auto dim = detail::get_cov_dim(N);
0096     if (i < 0 || j < 0 || i >= dim || j >= dim) {
0097       throw std::invalid_argument("The covariance matrix dimensions (" + std::to_string(dim) +
0098                                   ") and the passed dimensions are not compatible (dimI = " + std::to_string(i) +
0099                                   ", dimJ = " + std::to_string(j) + ")");
0100     }
0101
0102     return cov[detail::to_lower_tri(i, j)];
0103   }
0104
0105   template <typename DimEnum, typename Scalar, std::size_t N>
0106   constexpr void set_cov_value(Scalar value, std::array<Scalar, N>& cov, DimEnum dimI, DimEnum dimJ) {
0107     auto i = detail::to_index(dimI);
0108     auto j = detail::to_index(dimJ);
0109
0110     constexpr auto dim = detail::get_cov_dim(N);
0111     if (i < 0 || j < 0 || i >= dim || j >= dim) {
0112       throw std::invalid_argument("The covariance matrix dimensions (" + std::to_string(dim) +
0113                                   ") and the passed dimensions are not compatible (dimI = " + std::to_string(i) +
0114                                   ", dimJ = " + std::to_string(j) + ")");
0115     }
0116
0117     // Covariance is in lower triangle this brings us from 2D indices to 1D
0118     cov[detail::to_lower_tri(i, j)] = value;
0119   }
0120 } // namespace utils
0121
0122 } // namespace edm4hep
0123
0124 #endif // EDM4HEP_UTILS_COVMATRIXUTILS_H