feat: add getValue for variable matrices

mimizh2418 · Nov 26, 2024 · cd821d4 · cd821d4
1 parent b5b1200
commit cd821d4
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 16 deletions.
diff --git a/include/suboptimal/autodiff/Variable.h b/include/suboptimal/autodiff/Variable.h
@@ -3,8 +3,10 @@
 #pragma once
 
 #include <memory>
+#include <vector>
 
 #include <Eigen/Core>
+#include <Eigen/SparseCore>
 
 #include "suboptimal/autodiff/Expression.h"
 #include "suboptimal/util/concepts.h"
@@ -247,18 +249,6 @@ Variable atan2(const Y& y, const X& x) {
     return {atan2(y.expr, x.expr)};
   }
 }
-
-// Eigen typedefs
-
-using VectorXv = Eigen::VectorX<Variable>;
-using Vector2v = Eigen::Vector2<Variable>;
-using Vector3v = Eigen::Vector3<Variable>;
-using Vector4v = Eigen::Vector4<Variable>;
-
-using MatrixXv = Eigen::MatrixX<Variable>;
-using Matrix2v = Eigen::Matrix2<Variable>;
-using Matrix3v = Eigen::Matrix3<Variable>;
-using Matrix4v = Eigen::Matrix4<Variable>;
 }  // namespace suboptimal
 
 namespace Eigen {
@@ -280,3 +270,63 @@ struct NumTraits<suboptimal::Variable> : NumTraits<double> {
   };
 };
 }  // namespace Eigen
+
+namespace suboptimal {
+// Eigen typedefs
+
+using VectorXv = Eigen::VectorX<Variable>;
+using Vector2v = Eigen::Vector2<Variable>;
+using Vector3v = Eigen::Vector3<Variable>;
+using Vector4v = Eigen::Vector4<Variable>;
+
+using MatrixXv = Eigen::MatrixX<Variable>;
+using Matrix2v = Eigen::Matrix2<Variable>;
+using Matrix3v = Eigen::Matrix3<Variable>;
+using Matrix4v = Eigen::Matrix4<Variable>;
+
+/**
+ * Returns a matrix holding the values of the Variables in the input matrix
+ * @tparam Derived the derived type of the input matrix
+ * @param var_mat the input matrix of Variables
+ * @return a double matrix holding the values of the Variables in the input matrix
+ */
+template <typename Derived>
+  requires std::same_as<typename Derived::Scalar, Variable>
+Eigen::Matrix<double, Derived::RowsAtCompileTime, Derived::ColsAtCompileTime> getValues(
+    const Eigen::MatrixBase<Derived>& var_mat) {
+  Eigen::Matrix<double, Derived::RowsAtCompileTime, Derived::ColsAtCompileTime> result(var_mat.rows(), var_mat.cols());
+  for (Eigen::Index i = 0; i < var_mat.rows(); i++) {
+    for (Eigen::Index j = 0; j < var_mat.cols(); j++) {
+      result(i, j) = var_mat(i, j).getValue();
+    }
+  }
+  return result;
+}
+
+
+/**
+* Returns a sparse matrix holding the values of the Variables in the sparse input matrix
+* @tparam Scalar the scalar type of the input matrix
+* @tparam Options the storage scheme of the input matrix
+* @tparam StorageIndex the storage index type of the input matrix
+* @param var_mat the sparse matrix of Variables
+* @return a sparse double matrix holding the values of the Variables in the input matrix
+*/
+template <typename Scalar, int Options, typename StorageIndex>
+  requires std::same_as<Scalar, Variable>
+Eigen::SparseMatrix<double> getValues(const Eigen::SparseMatrix<Scalar, Options, StorageIndex>& var_mat) {
+  using InputMatType = Eigen::SparseMatrix<Scalar, Options, StorageIndex>;
+  using OutputMatType = Eigen::SparseMatrix<double, Options, StorageIndex>;
+
+  OutputMatType result(var_mat.rows(), var_mat.cols());
+  std::vector<Eigen::Triplet<double>> triplets{};
+  triplets.reserve(var_mat.nonZeros());
+  for (Eigen::Index i = 0; i < var_mat.outerSize(); i++) {
+    for (typename InputMatType::InnerIterator it(var_mat, i); it; ++it) {
+      triplets.push_back(Eigen::Triplet<double>(it.row(), it.col(), it.value().getValue()));
+    }
+  }
+  result.setFromTriplets(triplets.begin(), triplets.end());
+  return result;
+}
+}  // namespace suboptimal
diff --git a/test/autodiff/Variable_test.cpp b/test/autodiff/Variable_test.cpp
@@ -93,7 +93,7 @@ TEST_CASE("Autodiff - Variable STL functions", "[autodiff]") {
   }
 }
 
-TEST_CASE("Autodiff - Variable matrix", "[autodiff]") {
+TEST_CASE("Autodiff - Eigen Variable support", "[autodiff]") {
   SECTION("Matrix operations") {
     const Eigen::Matrix2d x_val{{1, 2},  //
                                 {3, 4}};
@@ -106,10 +106,35 @@ TEST_CASE("Autodiff - Variable matrix", "[autodiff]") {
     const Matrix2v y{{Variable{5}, Variable{6}},  //
                      {Variable{7}, Variable{8}}};
     const Matrix2v f = x * y;
+    CHECK(suboptimal::getValues(f).isApprox(f_val));
+  }
 
-    for (Eigen::Index i = 0; i < f.reshaped().size(); i++) {
-      CHECK_THAT(f.reshaped()(i).getValue(), Catch::Matchers::WithinAbs(f_val(i), 1e-9));
-    }
+  SECTION("Sparse matrix operations") {
+    Eigen::SparseMatrix<double> x_val(2, 2);
+    x_val.insert(0, 0) = 1;
+    x_val.insert(0, 1) = 2;
+    x_val.insert(1, 0) = 3;
+    x_val.insert(1, 1) = 4;
+    Eigen::SparseMatrix<double> y_val(2, 2);
+    y_val.insert(0, 0) = 5;
+    y_val.insert(0, 1) = 6;
+    y_val.insert(1, 0) = 7;
+    y_val.insert(1, 1) = 8;
+    const Eigen::SparseMatrix<double> f_val = x_val * y_val;
+
+    Eigen::SparseMatrix<Variable> x(2, 2);
+    x.insert(0, 0) = Variable{1};
+    x.insert(0, 1) = Variable{2};
+    x.insert(1, 0) = Variable{3};
+    x.insert(1, 1) = Variable{4};
+    Eigen::SparseMatrix<Variable> y(2, 2);
+    y.insert(0, 0) = Variable{5};
+    y.insert(0, 1) = Variable{6};
+    y.insert(1, 0) = Variable{7};
+    y.insert(1, 1) = Variable{8};
+    const Eigen::SparseMatrix<Variable> f = x * y;
+
+    CHECK(suboptimal::getValues(f).isApprox(f_val));
   }
 
   SECTION("Vector operations") {