add the possibility to solve for sparse rhs with Cholmod

2 files changed, 140 insertions, 9 deletions

diff --git a/unsupported/Eigen/src/SparseExtra/CholmodSupport.h b/unsupported/Eigen/src/SparseExtra/CholmodSupport.h
index 08f01a6e6..b0f1c8e42 100644
--- a/unsupported/Eigen/src/SparseExtra/CholmodSupport.h
+++ b/unsupported/Eigen/src/SparseExtra/CholmodSupport.h
@@ -26,7 +26,57 @@
 #define EIGEN_CHOLMODSUPPORT_H
 
 namespace internal {
+
+template<typename _DecompositionType, typename Rhs> struct sparse_solve_retval_base;
+template<typename _DecompositionType, typename Rhs> struct sparse_solve_retval;
   
+template<typename DecompositionType, typename Rhs>
+struct traits<sparse_solve_retval_base<DecompositionType, Rhs> >
+{
+  typedef typename DecompositionType::MatrixType MatrixType;
+  typedef SparseMatrix<typename Rhs::Scalar, Rhs::Options, typename Rhs::Index> ReturnType;
+};
+
+template<typename _DecompositionType, typename Rhs> struct sparse_solve_retval_base
+ : public ReturnByValue<sparse_solve_retval_base<_DecompositionType, Rhs> >
+{
+  typedef typename remove_all<typename Rhs::Nested>::type RhsNestedCleaned;
+  typedef _DecompositionType DecompositionType;
+  typedef ReturnByValue<sparse_solve_retval_base> Base;
+  typedef typename Base::Index Index;
+
+  sparse_solve_retval_base(const DecompositionType& dec, const Rhs& rhs)
+    : m_dec(dec), m_rhs(rhs)
+  {}
+
+  inline Index rows() const { return m_dec.cols(); }
+  inline Index cols() const { return m_rhs.cols(); }
+  inline const DecompositionType& dec() const { return m_dec; }
+  inline const RhsNestedCleaned& rhs() const { return m_rhs; }
+
+  template<typename Dest> inline void evalTo(Dest& dst) const
+  {
+    static_cast<const sparse_solve_retval<DecompositionType,Rhs>*>(this)->evalTo(dst);
+  }
+
+  protected:
+    const DecompositionType& m_dec;
+    const typename Rhs::Nested m_rhs;
+};
+
+#define EIGEN_MAKE_SPARSE_SOLVE_HELPERS(DecompositionType,Rhs) \
+  typedef typename DecompositionType::MatrixType MatrixType; \
+  typedef typename MatrixType::Scalar Scalar; \
+  typedef typename MatrixType::RealScalar RealScalar; \
+  typedef typename MatrixType::Index Index; \
+  typedef Eigen::internal::sparse_solve_retval_base<DecompositionType,Rhs> Base; \
+  using Base::dec; \
+  using Base::rhs; \
+  using Base::rows; \
+  using Base::cols; \
+  sparse_solve_retval(const DecompositionType& dec, const Rhs& rhs) \
+    : Base(dec, rhs) {}
+
 template<typename Scalar, typename CholmodType>
 void cholmod_configure_matrix(CholmodType& mat)
 {
@@ -94,6 +144,13 @@ cholmod_sparse viewAsCholmod(SparseMatrix<_Scalar,_Options,_Index>& mat)
   return res;
 }
 
+template<typename _Scalar, int _Options, typename _Index>
+const cholmod_sparse viewAsCholmod(const SparseMatrix<_Scalar,_Options,_Index>& mat)
+{
+  cholmod_sparse res = viewAsCholmod(mat.const_cast_derived());
+  return res;
+}
+
 /** Returns a view of the Eigen sparse matrix \a mat as Cholmod sparse matrix.
   * The data are not copied but shared. */
 template<typename _Scalar, int _Options, typename _Index, unsigned int UpLo>
@@ -247,6 +304,20 @@ class CholmodDecomposition
       return internal::solve_retval<CholmodDecomposition, Rhs>(*this, b.derived());
     }
     
+    /** \returns the solution x of \f$ A x = b \f$ using the current decomposition of A.
+      *
+      * \sa compute()
+      */
+    template<typename Rhs>
+    inline const internal::sparse_solve_retval<CholmodDecomposition, Rhs>
+    solve(const SparseMatrixBase<Rhs>& b) const
+    {
+      eigen_assert(m_isInitialized && "LLT is not initialized.");
+      eigen_assert(rows()==b.rows()
+                && "CholmodDecomposition::solve(): invalid number of rows of the right hand side matrix b");
+      return internal::sparse_solve_retval<CholmodDecomposition, Rhs>(*this, b.derived());
+    }
+    
     /** Performs a symbolic decomposition on the sparcity of \a matrix.
       *
       * This function is particularly useful when solving for several problems having the same structure.
@@ -305,10 +376,30 @@ class CholmodDecomposition
       {
         this->m_info = NumericalIssue;
       }
-      dest = Matrix<Scalar,Dynamic,1>::Map(reinterpret_cast<Scalar*>(x_cd->x),b.rows());
+      // TODO optimize this copy by swapping when possible (be carreful with alignment, etc.)
+      dest = Matrix<Scalar,Dest::RowsAtCompileTime,Dest::ColsAtCompileTime>::Map(reinterpret_cast<Scalar*>(x_cd->x),b.rows(),b.cols());
       cholmod_free_dense(&x_cd, &m_cholmod);
     }
     
+    /** \internal */
+    template<typename RhsScalar, int RhsOptions, typename RhsIndex, typename DestScalar, int DestOptions, typename DestIndex>
+    void _solve(const SparseMatrix<RhsScalar,RhsOptions,RhsIndex> &b, SparseMatrix<DestScalar,DestOptions,DestIndex> &dest) const
+    {
+      eigen_assert(m_factorizationIsOk && "The decomposition is not in a valid state for solving, you must first call either compute() or symbolic()/numeric()");
+      const Index size = m_cholmodFactor->n;
+      eigen_assert(size==b.rows());
+
+      // note: cs stands for Cholmod Sparse
+      cholmod_sparse b_cs = viewAsCholmod(b);
+      cholmod_sparse* x_cs = cholmod_spsolve(CHOLMOD_A, m_cholmodFactor, &b_cs, &m_cholmod);
+      if(!x_cs)
+      {
+        this->m_info = NumericalIssue;
+      }
+      // TODO optimize this copy by swapping when possible (be carreful with alignment, etc.)
+      dest = viewAsEigen<DestScalar,DestOptions,DestIndex>(*x_cs);
+      cholmod_free_sparse(&x_cs, &m_cholmod);
+    }
     #endif // EIGEN_PARSED_BY_DOXYGEN
 
   protected:
@@ -335,6 +426,19 @@ struct solve_retval<CholmodDecomposition<_MatrixType,_UpLo>, Rhs>
   }
 };
 
+template<typename _MatrixType, int _UpLo, typename Rhs>
+struct sparse_solve_retval<CholmodDecomposition<_MatrixType,_UpLo>, Rhs>
+  : sparse_solve_retval_base<CholmodDecomposition<_MatrixType,_UpLo>, Rhs>
+{
+  typedef CholmodDecomposition<_MatrixType,_UpLo> Dec;
+  EIGEN_MAKE_SPARSE_SOLVE_HELPERS(Dec,Rhs)
+
+  template<typename Dest> void evalTo(Dest& dst) const
+  {
+    dec()._solve(rhs(),dst);
+  }
+};
+
 }
 
 #endif // EIGEN_CHOLMODSUPPORT_H
diff --git a/unsupported/test/sparse_llt.cpp b/unsupported/test/sparse_llt.cpp
index 21bd36d35..7b1cad4c5 100644
--- a/unsupported/test/sparse_llt.cpp
+++ b/unsupported/test/sparse_llt.cpp
@@ -40,19 +40,21 @@ template<typename Scalar> void sparse_llt(int rows, int cols)
     DenseMatrix refMat2(rows, cols);
 
     DenseVector b = DenseVector::Random(cols);
-    DenseVector refX(cols), x(cols);
+    DenseVector ref_x(cols), x(cols);
+    DenseMatrix B = DenseMatrix::Random(rows,cols);
+    DenseMatrix ref_X(rows,cols), X(rows,cols);
 
     initSparse<Scalar>(density, refMat2, m2, ForceNonZeroDiag|MakeLowerTriangular, 0, 0);
 
     for(int i=0; i<rows; ++i)
       m2.coeffRef(i,i) = refMat2(i,i) = internal::abs(internal::real(refMat2(i,i)));
 
-    refX = refMat2.template selfadjointView<Lower>().llt().solve(b);
+    ref_x = refMat2.template selfadjointView<Lower>().llt().solve(b);
     if (!NumTraits<Scalar>::IsComplex)
     {
       x = b;
       SparseLLT<SparseMatrix<Scalar> > (m2).solveInPlace(x);
-      VERIFY(refX.isApprox(x,test_precision<Scalar>()) && "LLT: default");
+      VERIFY(ref_x.isApprox(x,test_precision<Scalar>()) && "LLT: default");
     }
         
 #ifdef EIGEN_CHOLMOD_SUPPORT
@@ -62,14 +64,14 @@ template<typename Scalar> void sparse_llt(int rows, int cols)
       SparseMatrix<Scalar> m3 = m2.adjoint()*m2;
       DenseMatrix refMat3 = refMat2.adjoint()*refMat2;
       
-      refX = refMat3.template selfadjointView<Lower>().llt().solve(b);
+      ref_x = refMat3.template selfadjointView<Lower>().llt().solve(b);
       
       x = b;
       SparseLLT<SparseMatrix<Scalar>, Cholmod>(m3).solveInPlace(x);
       VERIFY((m3*x).isApprox(b,test_precision<Scalar>()) && "LLT legacy: cholmod solveInPlace");
       
       x = SparseLLT<SparseMatrix<Scalar>, Cholmod>(m3).solve(b);
-      VERIFY(refX.isApprox(x,test_precision<Scalar>()) && "LLT legacy: cholmod solve");
+      VERIFY(ref_x.isApprox(x,test_precision<Scalar>()) && "LLT legacy: cholmod solve");
     }
     
     // new API
@@ -78,13 +80,38 @@ template<typename Scalar> void sparse_llt(int rows, int cols)
       SparseMatrix<Scalar> m3 = m2.adjoint()*m2;
       DenseMatrix refMat3 = refMat2.adjoint()*refMat2;
       
-      refX = refMat3.template selfadjointView<Lower>().llt().solve(b);
+      // with a single vector as the rhs
+      ref_x = refMat3.template selfadjointView<Lower>().llt().solve(b);
 
       x = CholmodDecomposition<SparseMatrix<Scalar>, Lower>(m3).solve(b);
-      VERIFY(refX.isApprox(x,test_precision<Scalar>()) && "LLT: cholmod solve");
+      VERIFY(ref_x.isApprox(x,test_precision<Scalar>()) && "LLT: cholmod solve, single dense rhs");
       
       x = CholmodDecomposition<SparseMatrix<Scalar>, Upper>(m3).solve(b);
-      VERIFY(refX.isApprox(x,test_precision<Scalar>()) && "LLT: cholmod solve");
+      VERIFY(ref_x.isApprox(x,test_precision<Scalar>()) && "LLT: cholmod solve, single dense rhs");
+      
+      
+      // with multiple rhs
+      ref_X = refMat3.template selfadjointView<Lower>().llt().solve(B);
+
+      X = CholmodDecomposition<SparseMatrix<Scalar>, Lower>(m3).solve(B);
+      VERIFY(ref_X.isApprox(X,test_precision<Scalar>()) && "LLT: cholmod solve, multiple dense rhs");
+      
+      X = CholmodDecomposition<SparseMatrix<Scalar>, Upper>(m3).solve(B);
+      VERIFY(ref_X.isApprox(X,test_precision<Scalar>()) && "LLT: cholmod solve, multiple dense rhs");
+      
+      
+      // with a sparse rhs
+      SparseMatrix<Scalar> spB(rows,cols), spX(rows,cols);
+      B.diagonal().array() += 1;
+      spB = B.sparseView(0.5,1);
+      
+      ref_X = refMat3.template selfadjointView<Lower>().llt().solve(DenseMatrix(spB));
+
+      spX = CholmodDecomposition<SparseMatrix<Scalar>, Lower>(m3).solve(spB);
+      VERIFY(ref_X.isApprox(spX.toDense(),test_precision<Scalar>()) && "LLT: cholmod solve, multiple sparse rhs");
+      
+      spX = CholmodDecomposition<SparseMatrix<Scalar>, Upper>(m3).solve(spB);
+      VERIFY(ref_X.isApprox(spX.toDense(),test_precision<Scalar>()) && "LLT: cholmod solve, multiple sparse rhs");
     }
 #endif