Caribou
EigenMatrix.h
1 #pragma once
2 #include <Caribou/config.h>
3 #include <Caribou/macros.h>
4 #include <Caribou/traits.h>
5 #include <sofa/defaulttype/BaseMatrix.h>
6 #include <Eigen/Core>
7 #include <Eigen/Sparse>
8 
9 namespace SofaCaribou::Algebra {
10 
49 template <typename Derived, typename Enable = void>
50 class EigenMatrix : public sofa::defaulttype::BaseMatrix
51 {
52  static_assert(
53  std::is_base_of_v<Eigen::EigenBase<std::decay_t<Derived> >, std::decay_t<Derived> >,
54  "The class template argument 'Derived' must be derived from Eigen::EigenBase<Derived>."
55  );
56 
57 public:
58  using EigenType = std::remove_cv_t<std::remove_reference_t<Derived>>;
59  using Base = sofa::defaulttype::BaseMatrix;
60  using Index = Base::Index;
61  using Real = SReal;
62 
69  explicit EigenMatrix(std::remove_reference_t<Derived> & eigen_matrix) : p_eigen_matrix(eigen_matrix) {}
70 
76  EigenMatrix(Eigen::Index rows, Eigen::Index cols) : p_eigen_matrix(rows, cols) {}
77 
78  // Abstract methods overrides
79  inline Index rowSize() const final { return p_eigen_matrix.rows(); }
80  inline Index colSize() const final { return p_eigen_matrix.cols(); }
81  inline Real element(Index i, Index j) const final { return p_eigen_matrix(i,j); }
82 
84  inline void resize(Index nbRow, Index nbCol) final { p_eigen_matrix.setConstant(nbRow, nbCol, 0); }
85 
87  inline void clear() override {p_eigen_matrix.setZero();}
88 
90  inline void set(Index i, Index j, double v) final {
91  p_eigen_matrix(i, j) = v;
92  }
93 
95  inline void add(Index i, Index j, double v) final {
96  p_eigen_matrix(i, j) += v;
97  }
98 
99  // Block operations on 3x3 and 2x2 sub-matrices
100  inline void add(Index i, Index j, const sofa::defaulttype::Mat3x3d & m) override { add_block<double, 3, 3>(i, j, m);}
101  inline void add(Index i, Index j, const sofa::defaulttype::Mat3x3f & m) override { add_block<float, 3, 3>(i, j, m);}
102  inline void add(Index i, Index j, const sofa::defaulttype::Mat2x2d & m) override { add_block<double, 2, 2>(i, j, m);}
103  inline void add(Index i, Index j, const sofa::defaulttype::Mat2x2f & m) override { add_block<float, 2, 2>(i, j, m);}
104 
106  inline void clearRow(Index i) final {
107  p_eigen_matrix.row(i).setZero();
108  }
109 
111  inline void clearRows(Index imin, Index imax) final {
112  p_eigen_matrix.middleRows(imin, (imax-imin)+1).setZero();
113  }
114 
116  inline void clearCol(Index i) final {
117  p_eigen_matrix.col(i).setZero();
118  }
119 
121  inline void clearCols(Index imin, Index imax) final {
122  p_eigen_matrix.middleCols(imin, (imax-imin)+1).setZero();
123  }
124 
126  const EigenType & matrix() const {return p_eigen_matrix;}
127 
128 private:
129 
130  template <typename Scalar, unsigned int nrows, unsigned int ncols>
131  void add_block(Index i, Index j, const sofa::defaulttype::Mat<nrows, ncols, Scalar> & m) {
132  const Eigen::Map<const Eigen::Matrix<Scalar, nrows, ncols, Eigen::RowMajor>> block(&(m[0][0]));
133  p_eigen_matrix.block(i, j, nrows, ncols) += block. template cast<typename EigenType::Scalar>();
134  }
135 
137  Derived p_eigen_matrix;
138 };
139 
143 template <typename Derived>
144 class EigenMatrix<Derived, CLASS_REQUIRES(std::is_base_of_v<Eigen::SparseMatrixBase<std::decay_t<Derived>>, std::decay_t<Derived>>)> : public sofa::defaulttype::BaseMatrix
145 {
146 
147 public:
148  using EigenType = std::remove_cv_t<std::remove_reference_t<Derived>>;
149  using Base = sofa::defaulttype::BaseMatrix;
150  using Index = Base::Index;
151  using Real = SReal;
152 
159  explicit EigenMatrix(std::remove_reference_t<Derived> & eigen_matrix) : p_eigen_matrix(eigen_matrix) {}
160 
166  EigenMatrix(Eigen::Index rows, Eigen::Index cols) : p_eigen_matrix(rows, cols) {}
167 
168  // Abstract methods overrides
169  inline Index rowSize() const final { return p_eigen_matrix.rows(); }
170  inline Index colSize() const final { return p_eigen_matrix.cols(); }
171 
175  inline bool symmetric() const {return p_is_symmetric;}
176 
180  inline void set_symmetric(bool is_symmetric) { p_is_symmetric = is_symmetric; }
181 
187  inline Real element(Index i, Index j) const final {
188  caribou_assert(p_initialized && "Accessing an element on an uninitialized matrix.");
189  return this->p_eigen_matrix.coeff(i,j);
190  }
191 
193  inline void resize(Index nbRow, Index nbCol) final {
194  p_triplets.clear();
195  this->p_eigen_matrix.resize(nbRow, nbCol);
196  p_initialized = false;
197  }
198 
200  inline void clear() final {
201  p_triplets.clear();
202  this->p_eigen_matrix.setZero();
203  p_initialized = false;
204  }
205 
214  inline void set(Index i, Index j, double v) final {
215  if (not p_initialized) {
216  initialize();
217  }
218 
219  this->p_eigen_matrix.coeffRef(i, j) = v;
220  }
221 
228  inline void add(Index i, Index j, double v) final {
229  // Note: this "if" condition should't slow down that much the addition of multiple entries
230  // because of branch predictions (the conditional branch will be same for the next
231  // X calls to add until compress is called).
232  if (not p_initialized) {
233  p_triplets.emplace_back(i, j, v);
234  } else {
235  p_eigen_matrix.coeffRef(i, j) += v;
236  }
237  }
238 
245  void compress() final {
246  if (not p_initialized) {
247  initialize();
248  } else {
249  p_eigen_matrix.makeCompressed();
250  }
251  }
252 
253  // Block operations on 3x3 and 2x2 sub-matrices
254  inline void add(Index i, Index j, const sofa::defaulttype::Mat3x3d & m) override { add_block<double, 3, 3>(i, j, m);}
255  inline void add(Index i, Index j, const sofa::defaulttype::Mat3x3f & m) override { add_block<float, 3, 3>(i, j, m);}
256  inline void add(Index i, Index j, const sofa::defaulttype::Mat2x2d & m) override { add_block<double, 2, 2>(i, j, m);}
257  inline void add(Index i, Index j, const sofa::defaulttype::Mat2x2f & m) override { add_block<float, 2, 2>(i, j, m);}
258 
260  inline void clearRow(Index row_id) final {
261  if (not p_initialized) {
262  initialize();
263  }
264 
265  using StorageIndex = typename EigenType::StorageIndex;
266  using Scalar = typename EigenType::Scalar;
267 
268  if constexpr (EigenType::IsRowMajor) {
269  const auto & start = p_eigen_matrix.outerIndexPtr()[row_id];
270  const auto & end = p_eigen_matrix.outerIndexPtr()[row_id+1];
271  memset(&(p_eigen_matrix.data().valuePtr()[start]), static_cast<const Scalar &>(0), (end-start)*sizeof(Scalar));
272  } else {
273  for (int col_id=0; col_id<p_eigen_matrix.outerSize(); ++col_id) {
274  const auto & start = p_eigen_matrix.outerIndexPtr()[col_id];
275  const auto & end = p_eigen_matrix.outerIndexPtr()[col_id+1];
276  const auto id = p_eigen_matrix.data().searchLowerIndex(start, end, static_cast<StorageIndex>(row_id));
277  if ((id<end) && (p_eigen_matrix.data().index(id)==row_id)) {
278  p_eigen_matrix.data().value(id) = static_cast<const Scalar &>(0);
279  }
280  }
281  }
282  }
283 
285  inline void clearRows(Index imin, Index imax) final {
286  if (not p_initialized) {
287  initialize();
288  }
289 
290  using StorageIndex = typename EigenType::StorageIndex;
291  using Scalar = typename EigenType::Scalar;
292 
293  if constexpr (EigenType::IsRowMajor) {
294  const auto & start = p_eigen_matrix.outerIndexPtr()[imin];
295  const auto & end = p_eigen_matrix.outerIndexPtr()[imax+1];
296  memset(&(p_eigen_matrix.data().valuePtr()[start]), static_cast<const Scalar &>(0), (end-start)*sizeof(Scalar));
297  } else {
298  for (int col_id=0; col_id<p_eigen_matrix.outerSize(); ++col_id) {
299  const auto & start = p_eigen_matrix.outerIndexPtr()[col_id];
300  const auto & end = p_eigen_matrix.outerIndexPtr()[col_id+1];
301  auto id = p_eigen_matrix.data().searchLowerIndex(start, end, static_cast<StorageIndex>(imin));
302  for (;id<end;++id) {
303  if (imin <= p_eigen_matrix.data().index(id) and p_eigen_matrix.data().index(id) <= imax) {
304  p_eigen_matrix.data().value(id) = static_cast<const Scalar &>(0);
305  } else if (p_eigen_matrix.data().index(id) > imax) {
306  break;
307  }
308  }
309  }
310  }
311  }
312 
314  inline void clearCol(Index col_id) final {
315  if (not p_initialized) {
316  initialize();
317  }
318 
319  using StorageIndex = typename EigenType::StorageIndex;
320  using Scalar = typename EigenType::Scalar;
321 
322  if constexpr (not EigenType::IsRowMajor) {
323  const auto & start = p_eigen_matrix.outerIndexPtr()[col_id];
324  const auto & end = p_eigen_matrix.outerIndexPtr()[col_id+1];
325  memset(&(p_eigen_matrix.data().valuePtr()[start]), static_cast<const Scalar &>(0), (end-start)*sizeof(Scalar));
326  } else {
327  for (int row_id=0; row_id<p_eigen_matrix.outerSize(); ++row_id) {
328  const auto & start = p_eigen_matrix.outerIndexPtr()[row_id];
329  const auto & end = p_eigen_matrix.outerIndexPtr()[row_id+1];
330  const auto id = p_eigen_matrix.data().searchLowerIndex(start, end, static_cast<StorageIndex>(col_id));
331  if ((id<end) && (p_eigen_matrix.data().index(id)==col_id)) {
332  p_eigen_matrix.data().value(id) = static_cast<const Scalar &>(0);
333  }
334  }
335  }
336  }
337 
339  inline void clearCols(Index imin, Index imax) final {
340  if (not p_initialized) {
341  initialize();
342  }
343 
344  using StorageIndex = typename EigenType::StorageIndex;
345  using Scalar = typename EigenType::Scalar;
346 
347  if constexpr (not EigenType::IsRowMajor) {
348  const auto & start = p_eigen_matrix.outerIndexPtr()[imin];
349  const auto & end = p_eigen_matrix.outerIndexPtr()[imax+1];
350  memset(&(p_eigen_matrix.data().valuePtr()[start]), static_cast<const Scalar &>(0), (end-start)*sizeof(Scalar));
351  } else {
352  for (int row_id=0; row_id<p_eigen_matrix.outerSize(); ++row_id) {
353  const auto & start = p_eigen_matrix.outerIndexPtr()[row_id];
354  const auto & end = p_eigen_matrix.outerIndexPtr()[row_id+1];
355  auto id = p_eigen_matrix.data().searchLowerIndex(start, end, static_cast<StorageIndex>(imin));
356  for (;id<end;++id) {
357  if (imin <= p_eigen_matrix.data().index(id) and p_eigen_matrix.data().index(id) <= imax) {
358  p_eigen_matrix.data().value(id) = static_cast<const Scalar &>(0);
359  } else if (p_eigen_matrix.data().index(id) > imax) {
360  break;
361  }
362  }
363  }
364  }
365  }
366 
368  inline void clearRowCol(Index i) final {
369  if (not p_initialized) {
370  initialize();
371  }
372 
373  using StorageIndex = typename EigenType::StorageIndex;
374  using Scalar = typename EigenType::Scalar;
375 
376  // Clear the complete inner vector i (the row i or column i if it is in row major or column major respectively)
377  const auto & start = p_eigen_matrix.outerIndexPtr()[i];
378  const auto & end = p_eigen_matrix.outerIndexPtr()[i+1];
379  memset(&(p_eigen_matrix.data().valuePtr()[start]), static_cast<const Scalar &>(0), (end-start)*sizeof(Scalar));
380 
381  // Next, to clear the ith inner coefficients of each row in row-major (each column in column-major)
382  if (symmetric() and p_eigen_matrix.rows() == p_eigen_matrix.cols()) {
383  // When the sparse matrix is symmetric, we can avoid doing the binary search on each inner vectors. Instead,
384  // We loop on each inner coefficient of the ith outer vector (eg each (i,j) of the ith row in row major)
385  // and do the binary search only on the jth outer vector
386  for (typename EigenType::InnerIterator it(p_eigen_matrix, i); it; ++it) {
387  const auto inner_index = EigenType::IsRowMajor ? it.col() : it.row();
388  const auto & start = p_eigen_matrix.outerIndexPtr()[inner_index];
389  const auto & end = p_eigen_matrix.outerIndexPtr()[inner_index+1];
390  const auto id = p_eigen_matrix.data().searchLowerIndex(start, end, static_cast<StorageIndex>(i));
391  if ((id<end) && (p_eigen_matrix.data().index(id)==i)) {
392  p_eigen_matrix.data().value(id) = static_cast<const Scalar &>(0);
393  }
394  }
395  } else {
396  for (int outer_id=0; outer_id<p_eigen_matrix.outerSize(); ++outer_id) {
397  const auto & start = p_eigen_matrix.outerIndexPtr()[outer_id];
398  const auto & end = p_eigen_matrix.outerIndexPtr()[outer_id+1];
399  const auto id = p_eigen_matrix.data().searchLowerIndex(start, end, static_cast<StorageIndex>(i));
400  if ((id<end) && (p_eigen_matrix.data().index(id)==i)) {
401  p_eigen_matrix.data().value(id) = static_cast<const Scalar &>(0);
402  }
403  }
404  }
405  }
406 
408  const Derived & matrix() const {return p_eigen_matrix;}
409 private:
410 
414  template <typename Scalar, unsigned int N, unsigned int C>
415  void add_block(Index i, Index j, const sofa::defaulttype::Mat<N, C, Scalar> & m) {
416  for (size_t k=0;k<N;++k) {
417  for (size_t l=0;l<C;++l) {
418  const auto value = static_cast<typename EigenType::Scalar>(m[k][l]);
419  if (not p_initialized) {
420  p_triplets.emplace_back(i+k, j+l, value);
421  } else {
422  p_eigen_matrix.coeffRef(i+k, j+l) += value;
423  }
424  }
425  }
426  }
427 
433  void initialize() {
434  p_eigen_matrix.setFromTriplets(p_triplets.begin(), p_triplets.end());
435  p_triplets.clear();
436  p_initialized = true;
437  }
438 
441  std::vector<Eigen::Triplet<typename EigenType::Scalar>> p_triplets;
442 
446  bool p_initialized = false;
447 
449  Derived p_eigen_matrix;
450 
453  bool p_is_symmetric = false;
454 };
455 
456 } // namespace SofaCaribou::Algebra
SofaCaribou::Algebra::EigenMatrix::matrix
const EigenType & matrix() const
Get a const reference to the underlying Eigen matrix
Definition: EigenMatrix.h:126
SofaCaribou::Algebra::EigenMatrix::clearCol
void clearCol(Index i) final
Sets the entire columns i to zero.
Definition: EigenMatrix.h:116
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::clear
void clear() final
Set all entries to zero.
Definition: EigenMatrix.h:200
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::clearCols
void clearCols(Index imin, Index imax) final
Sets the columns from index imin to index imax (inclusively) to zero.
Definition: EigenMatrix.h:339
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::set
void set(Index i, Index j, double v) final
Set this value of the matrix entry (i, j) to the value of v.
Definition: EigenMatrix.h:214
SofaCaribou::Algebra::EigenMatrix::add
void add(Index i, Index j, double v) final
Adds v to the value of the matrix entry (i, j).
Definition: EigenMatrix.h:95
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::clearRowCol
void clearRowCol(Index i) final
Sets the entire row i and column i to zero.
Definition: EigenMatrix.h:368
SofaCaribou::Algebra::EigenMatrix::resize
void resize(Index nbRow, Index nbCol) final
Resize the matrix to nbRow x nbCol dimensions.
Definition: EigenMatrix.h:84
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::set_symmetric
void set_symmetric(bool is_symmetric)
Explicitly states if this matrix is symmetric.
Definition: EigenMatrix.h:180
SofaCaribou::Algebra::EigenMatrix::EigenMatrix
EigenMatrix(Eigen::Index rows, Eigen::Index cols)
Construct a new rows x cols Eigen matrix of type Derived.
Definition: EigenMatrix.h:76
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::add
void add(Index i, Index j, double v) final
Adds v to the value of the matrix entry (i, j).
Definition: EigenMatrix.h:228
SofaCaribou::Algebra::EigenMatrix::clearRow
void clearRow(Index i) final
Sets the entire row i to zero.
Definition: EigenMatrix.h:106
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::symmetric
bool symmetric() const
States if the matrix is symmetric.
Definition: EigenMatrix.h:175
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::clearRow
void clearRow(Index row_id) final
Sets the entire row i to zero.
Definition: EigenMatrix.h:260
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::element
Real element(Index i, Index j) const final
Return the matrix entry (i,j).
Definition: EigenMatrix.h:187
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::clearRows
void clearRows(Index imin, Index imax) final
Sets the rows from index imin to index imax (inclusively) to zero.
Definition: EigenMatrix.h:285
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::EigenMatrix
EigenMatrix(Eigen::Index rows, Eigen::Index cols)
Construct a new rows x cols Eigen matrix of type Derived.
Definition: EigenMatrix.h:166
SofaCaribou::Algebra::EigenMatrix::set
void set(Index i, Index j, double v) final
Set this value of the matrix entry (i, j) to the value of v.
Definition: EigenMatrix.h:90
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::clearCol
void clearCol(Index col_id) final
Sets the entire columns i to zero.
Definition: EigenMatrix.h:314
SofaCaribou::Algebra::EigenMatrix::clearRows
void clearRows(Index imin, Index imax) final
Sets the rows from index imin to index imax (inclusively) to zero.
Definition: EigenMatrix.h:111
SofaCaribou::Algebra::EigenMatrix::clearCols
void clearCols(Index imin, Index imax) final
Sets the columns from index imin to index imax (inclusively) to zero.
Definition: EigenMatrix.h:121
SofaCaribou::Algebra::EigenMatrix::clear
void clear() override
Set all entries to zero.
Definition: EigenMatrix.h:87
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::compress
void compress() final
Compress the matrix.
Definition: EigenMatrix.h:245
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::EigenMatrix
EigenMatrix(std::remove_reference_t< Derived > &eigen_matrix)
Construct the class using another Eigen matrix.
Definition: EigenMatrix.h:159
SofaCaribou::Algebra::EigenMatrix
This class can be use to represent any Eigen matrix within SOFA.
Definition: EigenMatrix.h:51
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::matrix
const Derived & matrix() const
Get a const reference to the underlying Eigen matrix
Definition: EigenMatrix.h:408
SofaCaribou::Algebra::EigenMatrix::EigenMatrix
EigenMatrix(std::remove_reference_t< Derived > &eigen_matrix)
Construct the class using another Eigen matrix.
Definition: EigenMatrix.h:69
SofaCaribou::Algebra::EigenMatrix< Derived, CLASS_REQUIRES(std::is_base_of_v< Eigen::SparseMatrixBase< std::decay_t< Derived >>, std::decay_t< Derived >>)>::resize
void resize(Index nbRow, Index nbCol) final
Resize the matrix to nbRow x nbCol dimensions.
Definition: EigenMatrix.h:193