Merge pull request #85 from beomki-yeo/eigen-getter-size-type

Eigen getter size type template parameter
acts-project · Jan 24, 2023 · 39b6c0c · 39b6c0c
2 parents 4efbfdf + 20c6a47
commit 39b6c0c
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 25 deletions.
diff --git a/math/eigen/include/algebra/math/impl/eigen_getter.hpp b/math/eigen/include/algebra/math/impl/eigen_getter.hpp
@@ -96,52 +96,66 @@ ALGEBRA_HOST_DEVICE inline auto eta(
 /// Functor used to access elements of Eigen matrices
 struct element_getter {
   /// Get non-const access to a matrix element
-  template <
-      typename derived_type,
-      std::enable_if_t<std::is_base_of<Eigen::DenseCoeffsBase<
-                                           derived_type, Eigen::WriteAccessors>,
-                                       Eigen::MatrixBase<derived_type> >::value,
-                       bool> = true>
+  template <typename derived_type, typename size_type_1, typename size_type_2,
+            std::enable_if_t<
+                std::is_base_of<
+                    Eigen::DenseCoeffsBase<derived_type, Eigen::WriteAccessors>,
+                    Eigen::MatrixBase<derived_type> >::value &&
+                    std::is_convertible<size_type_1, Eigen::Index>::value &&
+                    std::is_convertible<size_type_2, Eigen::Index>::value,
+                bool> = true>
   ALGEBRA_HOST_DEVICE inline auto &operator()(
-      Eigen::MatrixBase<derived_type> &m, int row, int col) const {
+      Eigen::MatrixBase<derived_type> &m, size_type_1 row,
+      size_type_2 col) const {
 
     return m(row, col);
   }
   /// Get const access to a matrix element
-  template <typename derived_type>
+  template <typename derived_type, typename size_type_1, typename size_type_2,
+            std::enable_if_t<
+                std::is_convertible<size_type_1, Eigen::Index>::value &&
+                    std::is_convertible<size_type_2, Eigen::Index>::value,
+                bool> = true>
   ALGEBRA_HOST_DEVICE inline auto operator()(
-      const Eigen::MatrixBase<derived_type> &m, int row, int col) const {
+      const Eigen::MatrixBase<derived_type> &m, size_type_1 row,
+      size_type_2 col) const {
 
     return m(row, col);
   }
 };  // struct element_getter
 
 /// Function extracting an element from a matrix (const)
-template <typename derived_type>
+template <typename derived_type, typename size_type_1, typename size_type_2>
 ALGEBRA_HOST_DEVICE inline auto element(
-    const Eigen::MatrixBase<derived_type> &m, int row, int col) {
+    const Eigen::MatrixBase<derived_type> &m, size_type_1 row,
+    size_type_2 col) {
 
   return element_getter()(m, row, col);
 }
 
 /// Function extracting an element from a matrix (non-const)
 template <
-    typename derived_type,
+    typename derived_type, typename size_type_1, typename size_type_2,
     std::enable_if_t<std::is_base_of<Eigen::DenseCoeffsBase<
                                          derived_type, Eigen::WriteAccessors>,
                                      Eigen::MatrixBase<derived_type> >::value,
                      bool> = true>
 ALGEBRA_HOST_DEVICE inline auto &element(Eigen::MatrixBase<derived_type> &m,
-                                         int row, int col) {
+                                         size_type_1 row, size_type_2 col) {
 
   return element_getter()(m, row, col);
 }
 
 /// Functor used to extract a block from Eigen matrices
 struct block_getter {
-  template <int kROWS, int kCOLS, typename matrix_type>
-  ALGEBRA_HOST_DEVICE auto operator()(const matrix_type &m, int row,
-                                      int col) const {
+  template <int kROWS, int kCOLS, typename matrix_type, typename size_type_1,
+            typename size_type_2,
+            std::enable_if_t<
+                std::is_convertible<size_type_1, Eigen::Index>::value &&
+                    std::is_convertible<size_type_2, Eigen::Index>::value,
+                bool> = true>
+  ALGEBRA_HOST_DEVICE auto operator()(const matrix_type &m, size_type_1 row,
+                                      size_type_2 col) const {
 
     return m.template block<kROWS, kCOLS>(row, col);
   }

diff --git a/math/eigen/include/algebra/math/impl/eigen_matrix.hpp b/math/eigen/include/algebra/math/impl/eigen_matrix.hpp
@@ -44,31 +44,52 @@ struct actor {
   using vector3 = array_type<3>;
 
   /// Operator getting a reference to one element of a non-const matrix
-  template <int ROWS, int COLS>
+  template <int ROWS, int COLS, typename size_type_1, typename size_type_2,
+            std::enable_if_t<
+                std::is_convertible<size_type_1, Eigen::Index>::value &&
+                    std::is_convertible<size_type_2, Eigen::Index>::value,
+                bool> = true>
   ALGEBRA_HOST_DEVICE inline scalar_t &element(matrix_type<ROWS, COLS> &m,
-                                               int row, int col) const {
+                                               size_type_1 row,
+                                               size_type_2 col) const {
     return m(row, col);
   }
 
   /// Operator getting one value of a const matrix
-  template <int ROWS, int COLS>
+  template <int ROWS, int COLS, typename size_type_1, typename size_type_2,
+            std::enable_if_t<
+                std::is_convertible<size_type_1, Eigen::Index>::value &&
+                    std::is_convertible<size_type_2, Eigen::Index>::value,
+                bool> = true>
   ALGEBRA_HOST_DEVICE inline scalar_t element(const matrix_type<ROWS, COLS> &m,
-                                              int row, int col) const {
+                                              size_type_1 row,
+                                              size_type_2 col) const {
     return m(row, col);
   }
 
   /// Operator getting a block of a const matrix
-  template <int ROWS, int COLS, class input_matrix_type>
+  template <int ROWS, int COLS, class input_matrix_type, typename size_type_1,
+            typename size_type_2,
+            std::enable_if_t<
+                std::is_convertible<size_type_1, Eigen::Index>::value &&
+                    std::is_convertible<size_type_2, Eigen::Index>::value,
+                bool> = true>
   ALGEBRA_HOST_DEVICE matrix_type<ROWS, COLS> block(const input_matrix_type &m,
-                                                    int row, int col) {
+                                                    size_type_1 row,
+                                                    size_type_2 col) {
     return m.template block<ROWS, COLS>(row, col);
   }
 
   /// Operator setting a block
-  template <int ROWS, int COLS, class input_matrix_type>
+  template <int ROWS, int COLS, class input_matrix_type, typename size_type_1,
+            typename size_type_2,
+            std::enable_if_t<
+                std::is_convertible<size_type_1, Eigen::Index>::value &&
+                    std::is_convertible<size_type_2, Eigen::Index>::value,
+                bool> = true>
   ALGEBRA_HOST_DEVICE void set_block(input_matrix_type &m,
-                                     const matrix_type<ROWS, COLS> &b, int row,
-                                     int col) {
+                                     const matrix_type<ROWS, COLS> &b,
+                                     size_type_1 row, size_type_2 col) {
     m.template block<ROWS, COLS>(row, col) = b;
   }