Introduce bindings to assemble components matrices

Author: alxbilger

bindings/Sofa/src/SofaPython3/Sofa/Core/Binding_ForceField.cpp

@@ -28,7 +28,10 @@
 #include <sofa/core/behavior/MechanicalState.h>
 #include <sofa/core/behavior/ForceField.h>
 #include <sofa/core/MechanicalParams.h>
-#include <sofa/core/behavior/MultiMatrixAccessor.h>
+#include <sofa/core/behavior/DefaultMultiMatrixAccessor.h>
+#include <sofa/linearalgebra/CompressedRowSparseMatrix.h>
+
+#include <pybind11/eigen.h>
 
 #include <SofaPython3/PythonEnvironment.h>
 using sofapython3::PythonEnvironment;
@@ -177,10 +180,10 @@ namespace sofapython3
 
 
     template<class TDOFType>
-    void declare_forcefield(py::module &m, std::string typestr) {
-        std::string pyclass_name = std::string("ForceField") + typestr;
+    void declare_forcefield(py::module &m) {
+        const std::string pyclass_name = std::string("ForceField") + TDOFType::Name();
         py::class_<ForceField<TDOFType>, BaseObject, ForceField_Trampoline<TDOFType>, py_shared_ptr<ForceField<TDOFType>>> f(m, pyclass_name.c_str(), py::dynamic_attr(), py::multiple_inheritance(), sofapython3::doc::forceField::forceFieldClass);
-        
+
         f.def(py::init([](py::args &args, py::kwargs &kwargs) {
             auto ff = sofa::core::sptr<ForceField_Trampoline<TDOFType>> (new ForceField_Trampoline<TDOFType>());
 
@@ -204,15 +207,50 @@ namespace sofapython3
             }
             return ff;
         }));
+
+        using Real = typename TDOFType::Real;
+        using EigenSparseMatrix = Eigen::SparseMatrix<typename TDOFType::Real, Eigen::RowMajor>;
+        using EigenMatrixMap = Eigen::Map<EigenSparseMatrix>;
+
+        f.def("assembleKMatrix", [](ForceField<TDOFType>& self) -> EigenSparseMatrix
+        {
+            sofa::linearalgebra::CompressedRowSparseMatrix<Real> matrix;
+
+            if (const auto* mstate = self.getMState())
+            {
+                const auto matrixSize = static_cast<sofa::linearalgebra::BaseMatrix::Index>(mstate->getMatrixSize());
+                matrix.resize(matrixSize, matrixSize);
+
+                sofa::core::behavior::DefaultMultiMatrixAccessor accessor;
+                accessor.addMechanicalState(mstate);
+                accessor.setGlobalMatrix(&matrix);
+
+                auto mparams = *MechanicalParams::defaultInstance();
+                mparams.setKFactor(1.).setMFactor(0.).setBFactor(0.);
+
+                self.addKToMatrix(&mparams, &accessor);
+            }
+            matrix.compress();
+
+            return EigenMatrixMap(matrix.rows(), matrix.cols(), matrix.getColsValue().size(),
+                                (typename EigenMatrixMap::StorageIndex*)matrix.rowBegin.data(),
+                                (typename EigenMatrixMap::StorageIndex*)matrix.colsIndex.data(),
+                                matrix.colsValue.data());
+        }, sofapython3::doc::forceField::assembleKMatrix);
+
+        PythonFactory::registerType<ForceField<TDOFType>>([](sofa::core::objectmodel::Base* object)
+        {
+            return py::cast(dynamic_cast<ForceField<TDOFType>*>(object));
+        });
     }
 
 
 void moduleAddForceField(py::module &m) {
-    declare_forcefield<Vec3dTypes>(m, "Vec3d");
-    declare_forcefield<Vec2dTypes>(m, "Vec2d");
-    declare_forcefield<Vec1dTypes>(m, "Vec1d");
-    declare_forcefield<Rigid3dTypes>(m, "Rigid3d");
-    declare_forcefield<Rigid2dTypes>(m, "Rigid2d");
+    declare_forcefield<Vec3dTypes>(m);
+    declare_forcefield<Vec2dTypes>(m);
+    declare_forcefield<Vec1dTypes>(m);
+    declare_forcefield<Rigid3dTypes>(m);
+    declare_forcefield<Rigid2dTypes>(m);
 }
 
 }  // namespace sofapython3

bindings/Sofa/src/SofaPython3/Sofa/Core/Binding_ForceField_doc.h

@@ -25,4 +25,20 @@ namespace sofapython3::doc::forceField
 static auto forceFieldClass = R"(
                          An overridable class to create your own customized force field.
                          )";
+
+static constexpr const char* assembleKMatrix = R"(
+Assemble the stiffness matrix of a force field.
+
+Note that this function is not free. It assembles the stiffness matrix whether or not it has been assembled previously
+to add it into the global system matrix. Besides, the function does not prevent side effects of a second matrix assembly
+in a single time step.
+
+Typical usage example:
+FEM = root.addObject('HexahedronFEMForceField', name="FEM", youngModulus="4000", poissonRatio="0.3", method="large")
+...
+stiffness_matrix = self.force_field.assembleKMatrix()
+
+Returns:
+    A scipy.sparse.csr_matrix object representing the stiffness matrix of the force field
+)";
 }

examples/access_stiffness_matrix.py

@@ -0,0 +1,49 @@
+# Required import for python
+import Sofa
+import SofaRuntime
+from Sofa import SofaLinearSolver
+from scipy import sparse
+from scipy import linalg
+
+# Function called when the scene graph is being created
+def createScene(root):
+
+    root.addObject('VisualStyle', displayFlags="showBehaviorModels showForceFields")
+
+    root.addObject('DefaultAnimationLoop')
+    root.addObject('DefaultVisualManagerLoop')
+
+    root.addObject('EulerImplicitSolver', rayleighStiffness="0.1", rayleighMass="0.1")
+    root.addObject('SparseLDLSolver', applyPermutation="false", template="CompressedRowSparseMatrixd")
+
+    root.addObject('MechanicalObject', name="DoFs")
+    root.addObject('UniformMass', name="mass", totalMass="320")
+    root.addObject('RegularGridTopology', name="grid", nx="4", ny="4", nz="20", xmin="-9", xmax="-6", ymin="0", ymax="3", zmin="0", zmax="19")
+    root.addObject('BoxROI', name="box", box="-10 -1 -0.0001  -5 4 0.0001")
+    root.addObject('FixedConstraint', indices="@box.indices")
+    FEM = root.addObject('HexahedronFEMForceField', name="FEM", youngModulus="4000", poissonRatio="0.3", method="large")
+
+    root.addObject(MatrixAccessController('MatrixAccessor', name='matrixAccessor', force_field=FEM))
+
+    return root
+
+
+class MatrixAccessController(Sofa.Core.Controller):
+
+
+    def __init__(self, *args, **kwargs):
+        Sofa.Core.Controller.__init__(self, *args, **kwargs)
+        self.force_field = kwargs.get("force_field")
+
+    def onAnimateEndEvent(self, event):
+        stiffness_matrix = self.force_field.assembleKMatrix()
+
+        print('====================================')
+        print('Stiffness matrix')
+        print('====================================')
+        print('dtype: ' + str(stiffness_matrix.dtype))
+        print('shape: ' + str(stiffness_matrix.shape))
+        print('ndim: ' + str(stiffness_matrix.ndim))
+        print('nnz: ' + str(stiffness_matrix.nnz))
+        print('norm: ' + str(sparse.linalg.norm(stiffness_matrix)))
+