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