Support more types of linear systems

bindings/Modules/src/SofaPython3/SofaLinearSystem/Binding_LinearSystem.cpp

@@ -17,15 +17,12 @@
 *******************************************************************************
 * Contact information: [email protected]                             *
 ******************************************************************************/
-#include <SofaPython3/SofaLinearSystem/Binding_LinearSystem.h>
-#include <SofaPython3/SofaLinearSystem/Binding_LinearSystem_doc.h>
-#include <pybind11/pybind11.h>
+#include <SofaPython3/SofaLinearSystem/Binding_LinearSystem.inl>
 
-#include <SofaPython3/Sofa/Core/Binding_Base.h>
-#include <SofaPython3/PythonFactory.h>
 
-#include <sofa/component/linearsystem/TypedMatrixLinearSystem.h>
 #include <pybind11/eigen.h>
+#include <sofa/linearalgebra/BTDMatrix.h>
+#include <sofa/linearalgebra/BlockVector.h>
 #include <sofa/linearalgebra/CompressedRowSparseMatrix.h>
 
 
@@ -39,12 +36,6 @@ using EigenSparseMatrix = Eigen::SparseMatrix<Real, Eigen::RowMajor>;
 template<class Real>
 using EigenMatrixMap = Eigen::Map<Eigen::SparseMatrix<Real, Eigen::RowMajor> >;
 
-template<class Real>
-using Vector = Eigen::Matrix<Real,Eigen::Dynamic, 1>;
-
-template<class Real>
-using EigenVectorMap = Eigen::Map<Vector<Real>>;
-
 template<class TBlock>
 EigenSparseMatrix<typename sofa::linearalgebra::CompressedRowSparseMatrix<TBlock>::Real>
 toEigen(sofa::linearalgebra::CompressedRowSparseMatrix<TBlock>& matrix)
@@ -70,62 +61,104 @@ toEigen(sofa::linearalgebra::CompressedRowSparseMatrix<TBlock>& matrix)
     }
 }
 
-template<class TBlock>
-void bindLinearSystems(py::module &m)
+template<class TMatrix, class TVector>
+void bindMatrixAccessCRS(LinearSystemClass<TMatrix, TVector>& c)
 {
-    using CRS = sofa::linearalgebra::CompressedRowSparseMatrix<TBlock>;
-    using Real = typename CRS::Real;
-    using CRSLinearSystem = sofa::component::linearsystem::TypedMatrixLinearSystem<CRS, sofa::linearalgebra::FullVector<Real> >;
+    using Real = typename TMatrix::Real;
+    using CRSLinearSystem = sofa::component::linearsystem::TypedMatrixLinearSystem<TMatrix, TVector>;
 
-    const std::string typeName = CRSLinearSystem::GetClass()->className + CRSLinearSystem::GetCustomTemplateName();
-
-    py::class_<CRSLinearSystem,
-               sofa::core::objectmodel::BaseObject,
-               sofapython3::py_shared_ptr<CRSLinearSystem> > c(m, typeName.c_str(), sofapython3::doc::linearsystem::linearSystemClass);
-
-    c.def("A", [](CRSLinearSystem& self) -> EigenSparseMatrix<Real>
-    {
-        if (CRS* matrix = self.getSystemMatrix())
+    const auto matrixAccess =
+        [](CRSLinearSystem& self) -> EigenSparseMatrix<Real>
         {
-            if (matrix->colsValue.empty()) //null matrix: contains no entries
+            if (TMatrix* matrix = self.getSystemMatrix())
             {
-                return EigenSparseMatrix<Real>{matrix->rows(), matrix->cols()};
+                if (matrix->colsValue.empty()) //null matrix: contains no entries
+                {
+                    return EigenSparseMatrix<Real>{matrix->rows(), matrix->cols()};
+                }
+                return toEigen(*matrix);
             }
-            return toEigen(*matrix);
-        }
-        return {};
-    }, sofapython3::doc::linearsystem::linearSystem_A);
+            return {};
+        };
 
-    c.def("b", [](CRSLinearSystem& self) -> Vector<Real>
-    {
-        if (auto* vector = self.getRHSVector())
-        {
-            return EigenVectorMap<Real>(vector->ptr(), vector->size());
-        }
-        return {};
-    }, sofapython3::doc::linearsystem::linearSystem_b);
+    c.def("A", matrixAccess, sofapython3::doc::linearsystem::linearSystem_CRS_A);
+    c.def("get_system_matrix", matrixAccess, sofapython3::doc::linearsystem::linearSystem_CRS_get_system_matrix);
+}
 
-    c.def("x", [](CRSLinearSystem& self) -> Vector<Real>
-    {
-        if (auto* vector = self.getSolutionVector())
-        {
-            return EigenVectorMap<Real>(vector->ptr(), vector->size());
-        }
-        return {};
-    }, sofapython3::doc::linearsystem::linearSystem_x);
+template<>
+void bindMatrixAccess(LinearSystemClass<sofa::linearalgebra::CompressedRowSparseMatrix<SReal>, sofa::linearalgebra::FullVector<SReal>>& c)
+{
+    bindMatrixAccessCRS(c);
+}
 
+template<>
+void bindMatrixAccess(LinearSystemClass<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<2,2,SReal>>, sofa::linearalgebra::FullVector<SReal>>& c)
+{
+    bindMatrixAccessCRS(c);
+}
 
-    /// register the binding in the downcasting subsystem
-    PythonFactory::registerType<CRSLinearSystem>([](sofa::core::objectmodel::Base* object)
-    {
-        return py::cast(dynamic_cast<CRSLinearSystem*>(object));
-    });
+template<>
+void bindMatrixAccess(LinearSystemClass<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<3,3,SReal>>, sofa::linearalgebra::FullVector<SReal>>& c)
+{
+    bindMatrixAccessCRS(c);
+}
+
+template<>
+void bindMatrixAccess(LinearSystemClass<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<4,4,SReal>>, sofa::linearalgebra::FullVector<SReal>>& c)
+{
+    bindMatrixAccessCRS(c);
+}
+
+template<>
+void bindMatrixAccess(LinearSystemClass<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<6,6,SReal>>, sofa::linearalgebra::FullVector<SReal>>& c)
+{
+    bindMatrixAccessCRS(c);
+}
+
+template<>
+void bindMatrixAccess(LinearSystemClass<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<8,8,SReal>>, sofa::linearalgebra::FullVector<SReal>>& c)
+{
+    bindMatrixAccessCRS(c);
+}
+
+template<class Real>
+using DenseMatrix = Eigen::Matrix<Real,Eigen::Dynamic, Eigen::Dynamic>;
+
+template<class Real>
+using EigenDenseMatrixMap = Eigen::Map<DenseMatrix<Real>>;
+
+template<>
+void bindMatrixAccess(LinearSystemClass<sofa::linearalgebra::FullMatrix<SReal>, sofa::linearalgebra::FullVector<SReal>>& c)
+{
+    using CRSLinearSystem = sofa::component::linearsystem::TypedMatrixLinearSystem<sofa::linearalgebra::FullMatrix<SReal>, sofa::linearalgebra::FullVector<SReal>>;
+
+    const auto matrixAccess =
+        [](CRSLinearSystem& self) -> EigenDenseMatrixMap<SReal>
+        {
+            sofa::linearalgebra::FullMatrix<SReal>* matrix = self.getSystemMatrix();
+            const auto row = matrix ? matrix->rows() : 0;
+            const auto col = matrix ? matrix->cols() : 0;
+            return EigenDenseMatrixMap<SReal>(matrix ? matrix->ptr() : nullptr, row, col);
+        };
+    c.def("A", matrixAccess, sofapython3::doc::linearsystem::linearSystem_FullMatrix_A);
+    c.def("get_system_matrix", matrixAccess, sofapython3::doc::linearsystem::linearSystem_FullMatrix_get_system_matrix);
 }
 
 void moduleAddLinearSystem(py::module &m)
 {
-    bindLinearSystems<SReal>(m);
-    bindLinearSystems<sofa::type::Mat<3, 3, SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::FullMatrix<SReal>, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::SparseMatrix<SReal>, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::CompressedRowSparseMatrix<SReal>, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<2,2,SReal> >, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<3,3,SReal> >, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<4,4,SReal> >, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<6,6,SReal> >, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::CompressedRowSparseMatrix<sofa::type::Mat<8,8,SReal> >, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::DiagonalMatrix<SReal>, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::BlockDiagonalMatrix<3,SReal>, sofa::linearalgebra::FullVector<SReal> >(m);
+    bindLinearSystems<sofa::linearalgebra::RotationMatrix<SReal>, sofa::linearalgebra::FullVector<SReal> >(m);
+
+    bindLinearSystems<sofa::linearalgebra::BTDMatrix<6, SReal>, sofa::linearalgebra::BlockVector<6, SReal> >(m);
 }
 
 }

bindings/Modules/src/SofaPython3/SofaLinearSystem/Binding_LinearSystem.inl

@@ -0,0 +1,101 @@
+/******************************************************************************
+*                              SofaPython3 plugin                             *
+*                  (c) 2021 CNRS, University of Lille, INRIA                  *
+*                                                                             *
+* This program is free software; you can redistribute it and/or modify it     *
+* under the terms of the GNU Lesser General Public License as published by    *
+* the Free Software Foundation; either version 2.1 of the License, or (at     *
+* your option) any later version.                                             *
+*                                                                             *
+* This program is distributed in the hope that it will be useful, but WITHOUT *
+* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or       *
+* FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License *
+* for more details.                                                           *
+*                                                                             *
+* You should have received a copy of the GNU Lesser General Public License    *
+* along with this program. If not, see <http://www.gnu.org/licenses/>.        *
+*******************************************************************************
+* Contact information: [email protected]                             *
+******************************************************************************/
+#pragma once
+#include <SofaPython3/SofaLinearSystem/Binding_LinearSystem.h>
+#include <SofaPython3/SofaLinearSystem/Binding_LinearSystem_doc.h>
+#include <sofa/component/linearsystem/TypedMatrixLinearSystem.h>
+#include <pybind11/pybind11.h>
+
+#include <SofaPython3/PythonFactory.h>
+#include <SofaPython3/Sofa/Core/Binding_Base.h>
+#include <sofa/linearalgebra/EigenSparseMatrix.h>
+
+namespace py { using namespace pybind11; }
+
+namespace sofapython3
+{
+
+template<class Real>
+using Vector = Eigen::Matrix<Real,Eigen::Dynamic, 1>;
+
+template<class Real>
+using EigenVectorMap = Eigen::Map<Vector<Real>>;
+
+template<class TVector>
+Vector<typename TVector::Real>
+getVector(TVector* vector)
+{
+    using Real = typename TVector::Real;
+    if (vector)
+    {
+        return EigenVectorMap<Real>(vector->ptr(), vector->size());
+    }
+    return {};
+}
+
+template<class TMatrix, class TVector>
+Vector<typename TVector::Real> getRHSVector(sofa::component::linearsystem::TypedMatrixLinearSystem<TMatrix, TVector>& linearSystem)
+{
+    return getVector(linearSystem.getRHSVector());
+}
+
+template<class TMatrix, class TVector>
+Vector<typename TVector::Real> getSolutionVector(sofa::component::linearsystem::TypedMatrixLinearSystem<TMatrix, TVector>& linearSystem)
+{
+    return getVector(linearSystem.getSolutionVector());
+}
+
+template<class TMatrix, class TVector>
+using LinearSystemClass =
+    py::class_<sofa::component::linearsystem::TypedMatrixLinearSystem<TMatrix, TVector>,
+           sofa::core::objectmodel::BaseObject,
+           sofapython3::py_shared_ptr<sofa::component::linearsystem::TypedMatrixLinearSystem<TMatrix, TVector>> >;
+
+template<class TMatrix, class TVector>
+void bindMatrixAccess(LinearSystemClass<TMatrix, TVector>& c)
+{}
+
+template<class TMatrix, class TVector>
+void bindLinearSystems(py::module &m)
+{
+    using LinearSystem = sofa::component::linearsystem::TypedMatrixLinearSystem<TMatrix, TVector>;
+
+    const std::string typeName =
+        LinearSystem::GetClass()->className
+        + LinearSystem::GetCustomTemplateName();
+
+    LinearSystemClass<TMatrix, TVector> c(m, typeName.c_str(), sofapython3::doc::linearsystem::linearSystemClass);
+
+    c.def("b", getRHSVector<TMatrix, TVector>, sofapython3::doc::linearsystem::linearSystem_b);
+    c.def("get_rhs_vector", getRHSVector<TMatrix, TVector>, sofapython3::doc::linearsystem::linearSystem_b);
+
+    c.def("x", getSolutionVector<TMatrix, TVector>, sofapython3::doc::linearsystem::linearSystem_x);
+    c.def("get_solution_vector", getSolutionVector<TMatrix, TVector>, sofapython3::doc::linearsystem::linearSystem_x);
+
+    bindMatrixAccess<TMatrix, TVector>(c);
+
+    /// register the binding in the downcasting subsystem
+    PythonFactory::registerType<LinearSystem>([](sofa::core::objectmodel::Base* object)
+    {
+        return py::cast(dynamic_cast<LinearSystem*>(object));
+    });
+}
+
+}

bindings/Modules/src/SofaPython3/SofaLinearSystem/Binding_LinearSystem_doc.h

@@ -32,7 +32,7 @@ Linear system. Supports only CompressedRowSparseMatrix.
 linear_system = root.addObject('MatrixLinearSystem', template='CompressedRowSparseMatrixd')
 )";
 
-static auto linearSystem_A =
+static auto linearSystem_CRS_A =
 R"(
 Returns the global system matrix as a scipy sparse matrix
 
@@ -42,6 +42,36 @@ linear_system = root.addObject('MatrixLinearSystem', template='CompressedRowSpar
 matrix = linear_system.A()
 )";
 
+static auto linearSystem_CRS_get_system_matrix =
+R"(
+Returns the global system matrix as a scipy sparse matrix
+
+example:
+------------
+linear_system = root.addObject('MatrixLinearSystem', template='CompressedRowSparseMatrixd')
+matrix = linear_system.get_system_matrix()
+)";
+
+static auto linearSystem_FullMatrix_A =
+R"(
+Returns the global system matrix as a numpy array matrix
+
+example:
+------------
+linear_system = root.addObject('MatrixLinearSystem', template='FullMatrix')
+matrix = linear_system.A()
+)";
+
+static auto linearSystem_FullMatrix_get_system_matrix =
+R"(
+Returns the global system matrix as a numpy array matrix
+
+example:
+------------
+linear_system = root.addObject('MatrixLinearSystem', template='FullMatrix')
+matrix = linear_system.get_system_matrix()
+)";
+
 static auto linearSystem_b =
 R"(
 Returns the global system right hand side as a numpy array

bindings/Modules/src/SofaPython3/SofaLinearSystem/CMakeLists.txt

@@ -2,6 +2,7 @@ project(Bindings.Modules.SofaLinearSystem)
 
 set(HEADER_FILES
     ${CMAKE_CURRENT_SOURCE_DIR}/Binding_LinearSystem.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/Binding_LinearSystem.inl
     ${CMAKE_CURRENT_SOURCE_DIR}/Binding_LinearSystem_doc.h
 )