Transition to buildStiffnessMatrix over addKToMatrix in ForceField

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

@@ -30,6 +30,7 @@
 #include <sofa/core/MechanicalParams.h>
 #include <sofa/core/behavior/DefaultMultiMatrixAccessor.h>
 #include <sofa/linearalgebra/CompressedRowSparseMatrix.h>
+#include <sofa/core/behavior/BaseLocalForceFieldMatrix.h>
 
 #include <pybind11/eigen.h>
 
@@ -136,7 +137,7 @@ namespace sofapython3
     void ForceField_Trampoline<TDOFType>::addKToMatrix(const MechanicalParams* mparams, const MultiMatrixAccessor* dfId)
     {
         PythonEnvironment::gil acquire;
-        
+
         MultiMatrixAccessor::MatrixRef mref = dfId->getMatrix(this->mstate);
         sofa::linearalgebra::BaseMatrix* mat = mref.matrix;
 
@@ -183,8 +184,74 @@ namespace sofapython3
 
     }
 
-
     template<class TDOFType>
+    py::object ForceField_Trampoline<TDOFType>::_buildStiffnessMatrix(int nIndices, int nDofs)
+    {
+        PYBIND11_OVERRIDE_PURE(py::object, ForceField<TDOFType>, buildStiffnessMatrix, nIndices, nDofs);
+    }
+
+    template <class TDOFType>
+    void ForceField_Trampoline<TDOFType>::buildStiffnessMatrix(
+        sofa::core::behavior::StiffnessMatrix *matrix)
+    {
+        PythonEnvironment::gil acquire;
+
+        // nNodes is the number of nodes (positions) of the object whose K matrix we're computing
+        const int nNodes = static_cast<int>(this->mstate->getSize());
+        // nDofs is the number of degrees of freedom per-element of the object whose K matrix we're computing
+        static constexpr int nDofs = Coord::total_size;
+
+        if (py::get_overload(this, "addKToMatrix")
+            && !py::get_overload(this, "buildStiffnessMatrix"))
+        {
+            msg_deprecated_once() << "The method `addKToMatrix` is deprecated. Please use `buildStiffnessMatrix` instead.";
+
+            // Runs the compatibility layer with `addKToMatrix`
+            ForceField<TDOFType>::buildStiffnessMatrix(matrix);
+            return;
+        }
+
+        py::object ret = _buildStiffnessMatrix(nNodes, nDofs);
+
+        if(!py::isinstance<py::array>(ret))
+        {
+            throw py::type_error("Can't read return value of `buildStiffnessMatrix`. A numpy array is expected");
+        }
+
+        auto dfdx =
+            matrix->getForceDerivativeIn(this->mstate)
+            .withRespectToPositionsIn(this->mstate);
+
+        // if ret is numpy array
+        auto r = py::cast<py::array>(ret);
+        if (r.ndim() == 3 && r.shape(2) == 1)
+        {
+            // read K as a plain 2D matrix
+            auto kMatrix = r.unchecked<double, 3>();
+            for (size_t x = 0 ; x < size_t(kMatrix.shape(0)) ; ++x)
+            {
+                for (size_t y = 0 ; y < size_t(kMatrix.shape(1)) ; ++y)
+                {
+                    dfdx(x, y) += kMatrix(x,y, 0);
+                }
+            }
+        }
+        else if (r.ndim() == 2 && r.shape(1) == 3)
+        {
+            // consider ret to be a list of tuples [(i,j,[val])]
+            auto kMatrix = r.unchecked<double, 2>();
+            for (auto x = 0 ; x < kMatrix.shape(0) ; ++x)
+            {
+                dfdx(int(size_t(kMatrix(x,0))), int(size_t(kMatrix(x,1)))) += kMatrix(x,2);
+            }
+        }
+        else
+        {
+            throw py::type_error("Can't read return value of `buildStiffnessMatrix`. The method should return either a plain 2D matrix or a vector of tuples (i, j, val)");
+        }
+    }
+
+    template <class TDOFType>
     void declare_forcefield(py::module &m) {
         const std::string pyclass_name = std::string("ForceField") + TDOFType::Name();
         py::class_<ForceField<TDOFType>, BaseComponent, ForceField_Trampoline<TDOFType>, py_shared_ptr<ForceField<TDOFType>>> f(m, pyclass_name.c_str(), py::dynamic_attr(), py::multiple_inheritance(), sofapython3::doc::forceField::forceFieldClass);

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

@@ -50,8 +50,13 @@ class ForceField_Trampoline : public sofa::core::behavior::ForceField<TDOFType>
     pybind11::object _addKToMatrix(const sofa::core::MechanicalParams* mparams, int nNodes, int nDofs);
     void addKToMatrix(const sofa::core::MechanicalParams* mparams, const sofa::core::behavior::MultiMatrixAccessor* dfId) override;
 
+    void buildStiffnessMatrix(sofa::core::behavior::StiffnessMatrix* matrix) override;
+
     SReal getPotentialEnergy(const sofa::core::MechanicalParams* /*mparams*/, const DataVecCoord& /*x*/) const override { return 0.0; }
 
+private:
+
+    pybind11::object _buildStiffnessMatrix(int nIndices, int nDofs);
 };
 
 void moduleAddForceField(pybind11::module &m);

examples/example-forcefield.py

@@ -14,19 +14,30 @@ def __init__(self, ks=1.0, kd=1.0, *args, **kwargs):
     def init(self):
         mstate = self.getContext().mechanical
         self.initpos = mstate.position.array().copy()
-        self.k = np.zeros((1,1))
-        self.f = []
-        self.d = 0.5
 
     def addForce(self, m, out_force, pos, vel):
         with out_force.writeableArray() as wa:
-            wa[:] += ( (self.initpos-pos.value) * self.ks.value  )
-                 
+            wa[:] += ( (self.initpos-pos.value) * self.ks.value )
+
     def addDForce(self, params, df, dx):
-        pass
-
-    #def addKToMatrix(self, params, number_of_nodes, number_of_dofs_per_node):
-    #    print(" Python::addKToMatrix: ", a, " ", b)
+        kFactor = params['kFactor']
+        with df.writeableArray() as dfwa:
+            dfwa[:] -= dx.value * self.ks.value * kFactor
+
+    def addKToMatrix(self, params, number_of_nodes, number_of_dofs_per_node):
+        K = []
+        kFactor = params['kFactor']
+        for n in range(0, number_of_nodes):
+            for d in range(0, number_of_dofs_per_node):
+                K.append([number_of_dofs_per_node * n + d, number_of_dofs_per_node * n + d, -kFactor * self.kd.value])
+        return np.asarray(K)
+
+    def buildStiffnessMatrix(self, number_of_nodes, number_of_dofs_per_node):
+        K = []
+        for n in range(0, number_of_nodes):
+            for d in range(0, number_of_dofs_per_node):
+                K.append([number_of_dofs_per_node * n + d, number_of_dofs_per_node * n + d, -self.ks.value])
+        return np.asarray(K)
 
 
 def createScene(root):
@@ -42,17 +53,24 @@ def createScene(root):
 
         root.addObject("LineAxis")
         root.addObject("DefaultAnimationLoop", name="loop")
-        root.addObject("EulerImplicitIntegrationScheme")
-        root.addObject("CGLinearSolver", tolerance=1e-12, threshold=1e-12, iterations=25)
 
-        o = root.addChild("Object")
-        c = o.addObject("MechanicalObject", name="mechanical", position=[0.0,0.0,0.0, 1.0,0.0,0.0])
-        c.showObject = True
-        c.showColor = [1.0,0.0,0.0,1.0]
-        c.drawMode = 1
+        def add_spring_simulation(node, start_position, color):
+            node.addObject("EulerImplicitIntegrationScheme", rayleighStiffness=0.1, rayleighMass=0.1)
+
+            o = node.addChild("Object")
+            c = o.addObject("MechanicalObject", name="mechanical", position=[start_position + 0.0,0.0,0.0, start_position + 1.0,0.0,0.0], showColor=color)
+            c.showObject = True
+            c.drawMode = 1
+
+            o.addObject("UniformMass", name="mass", totalMass=[0.1])
+            o.addObject( RestShapeForceField(name="spring", ks=2.0, kd=0.1))
 
-        o.addObject("UniformMass", name="mass", totalMass=[0.1])
-        o.addObject( RestShapeForceField(name="CPPObject", ks=2.0, kd=0.1))
+        with root.addChild("Direct") as direct:
+            add_spring_simulation(direct, 0.0, [0,1,0,1])
+            direct.addObject("SparseLDLSolver", template="CompressedRowSparseMatrixd")
+        with root.addChild("Iterative") as iterative:
+            add_spring_simulation(iterative, 0.2, [1,0,0,1])
+            iterative.addObject("CGLinearSolver", tolerance=1e-12, threshold=1e-12, iterations=25)
 
         return root