Example of a ForceField implemented with JAX (#557)

* Add the example

* Igne jax example on main CI

* Fix ignoring the jax examples

* Fix addKToMatrix() with different options

* Apply suggestions from code review

---------

Co-authored-by: Paul Baksic <paul.baksic@outlook.fr>
Co-authored-by: Hugo <hugo.talbot@sofa-framework.org>

Author: 3 people

examples/.scene-tests

@@ -10,3 +10,6 @@ iterations "access_stiffness_matrix.py" "3"
 
 # Ignore additional examples
 ignore "additional-examples/.*"
+
+# Ignore jax examples
+ignore "jax/*"

examples/jax/forcefield.py

@@ -0,0 +1,180 @@
+"""
+Toy example of a force field leveraging autodiff with JAX.
+
+JAX can be installed via e.g. `pip install -U jax[cuda12]`
+"""
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import Sofa
+
+
+# Some configuration for JAX: device and precision
+# jax.config.update('jax_default_device', jax.devices('cpu')[0])
+jax.config.update("jax_default_device", jax.devices("gpu")[0])  # default "gpu"
+jax.config.update("jax_enable_x64", True)  # default False (ie use float32)
+
+
+@jax.jit  # JIT (just-in-time compilation) for better performance
+def get_force(position, length, stiffness):
+    """
+    Spring between the origin and the given position.
+
+    position: array of shape (n_particles, n_dimensions)
+    length: scalar or array of shape (n_particles, 1)
+    stiffness: scalar or array of shape (n_particles, 1)
+    """
+    distance = jnp.sqrt(jnp.sum(position**2, axis=1, keepdims=True))
+    direction = position / distance
+    return - stiffness * (distance - length) * direction
+
+
+@jax.jit  # JIT (just-in-time compilation) for better performance
+def get_dforce(position, length, stiffness, vector):
+    """
+    Compute the jacobian-vector product (jvp) using autodiff
+    """
+    def get_force_from_position(x):
+        return get_force(x, length, stiffness)
+    # Differentiate get_force() as a function of the position
+    return jax.jvp(get_force_from_position, (position,), (vector,))[1]
+
+
+@jax.jit  # JIT (just-in-time compilation) for better performance
+def get_kmatrix(position, length, stiffness):
+    """
+    Compute the jacobian using autodiff
+
+    Warning: The jacobian computed this way is a dense matrix.
+             Check `sparsejac` if you are interested in sparse jacobian with JAX.
+    """
+    def get_force_from_position(x):
+        return get_force(x, length, stiffness)
+    # Differentiate get_force() as a function of the position
+    return jax.jacrev(get_force_from_position)(position)
+
+
+class JaxForceField(Sofa.Core.ForceFieldVec3d):
+
+    def __init__(self, length, stiffness, *args, **kwargs):
+        Sofa.Core.ForceFieldVec3d.__init__(self, *args, **kwargs)
+        self.length = length
+        self.stiffness = stiffness
+        self.dense_to_sparse = None
+
+    def addForce(self, mechanical_parameters, out_force, position, velocity):
+        with out_force.writeableArray() as wa:
+            wa[:] += get_force(position.value, self.length, self.stiffness)
+
+    def addDForce(self, mechanical_parameters, df, dx):
+        with df.writeableArray() as wa:
+            wa[:] += get_dforce(self.mstate.position.value, self.length, self.stiffness, dx.value) * mechanical_parameters['kFactor']
+
+    # Option 1: Return the jacobian as a dense array (must have shape (n, n, 1) to be interpreted as such).
+    #           Note: Very slow for big sparse matrices.
+    # def addKToMatrix(self, mechanical_parameters, n_particles, n_dimensions):
+    #     jacobian = get_kmatrix(self.mstate.position.value, self.length, self.stiffness)
+    #     return np.array(jacobian).reshape((n_particles*n_dimensions, n_particles*n_dimensions, 1))
+
+    # Option 2: Return the non-zero coefficients of the jacobian as an array with rows (i, j, value).
+    #           Note: The extraction of the non-zero coefficients is faster with JAX on GPU.
+    # def addKToMatrix(self, mechanical_parameters, n_particles, n_dimensions):
+    #     jacobian = get_kmatrix(self.mstate.position.value, self.length, self.stiffness)
+    #     jacobian = jacobian.reshape((n_particles*n_dimensions, n_particles*n_dimensions))
+    #     i, j = jacobian.nonzero()
+    #     sparse_jacobian = jnp.stack([i, j, jacobian[i, j]], axis=1)
+    #     return np.array(sparse_jacobian)
+
+    # Option 2 optimization: We know the sparsity of the jacobian in advance (diagonal by 3x3 blocks).
+    def addKToMatrix(self, mechanical_parameters, n_particles, n_dimensions):
+        if self.dense_to_sparse is None:
+            # i = [0, 0, 0,  1, 1, 1,  2, 2, 2,  3, 3, 3,  ...]
+            # j = [0, 1, 2,  0, 1, 2,  0, 1, 2,  3, 4, 5,  ...]
+            i = jnp.repeat(jnp.arange(n_particles*n_dimensions), 3)
+            j = jnp.repeat(jnp.arange(n_particles*n_dimensions).reshape((-1, 3)), 3, axis=0).reshape(-1)
+            self.dense_to_sparse = lambda jac: jnp.stack([i, j, jac[i, j]], axis=1)
+            self.dense_to_sparse = jax.jit(self.dense_to_sparse)  # slightly faster with jit
+
+        jacobian = get_kmatrix(self.mstate.position.value, self.length, self.stiffness)
+        jacobian = jacobian.reshape((n_particles*n_dimensions, n_particles*n_dimensions))
+        sparse_jacobian = self.dense_to_sparse(jacobian)
+        sparse_jacobian = np.array(sparse_jacobian)
+        # Note: with the computations optimized, the conversion below can account for
+        #       ~90% of the time spent in this function.
+        return np.array(sparse_jacobian)
+
+
+def createScene(root, method="implicit-matrix-assembly", n_particles=1_000, use_sofa=False):
+    root.dt = 1e-3
+    root.gravity = (0, -9.8, 0)
+    root.box = (-5, -5, -5, 5, 5, 5)
+    root.addObject(
+        "RequiredPlugin",
+        pluginName=[
+            'Sofa.Component.Visual',
+            'Sofa.Component.ODESolver.Forward',
+            'Sofa.Component.ODESolver.Backward',
+            'Sofa.Component.LinearSolver.Iterative',
+            'Sofa.Component.LinearSolver.Direct',
+            'Sofa.Component.StateContainer',
+            'Sofa.Component.Mass',
+            'Sofa.Component.SolidMechanics.FEM.Elastic',
+            'Sofa.Component.SolidMechanics.Spring',
+        ]
+    )
+    root.addObject("DefaultAnimationLoop")
+    root.addObject("VisualStyle", displayFlags="showBehaviorModels showForceFields")
+
+    physics = root.addChild("Physics")
+
+    if method.lower() == "explicit":  # Requires the implementation of 'addForce'
+        physics.addObject("EulerExplicitSolver", name="eulerExplicit")
+    elif method.lower() == "implicit-matrix-free":  # Requires the implementation of 'addForce' and 'addDForce'
+        physics.addObject("EulerImplicitSolver", name="eulerImplicit")
+        physics.addObject("CGLinearSolver", template="GraphScattered", name="solver", iterations=50, tolerance=1e-5, threshold=1e-5)
+    elif method == "implicit-matrix-assembly":  # Requires the implementation of 'addForce', 'addDForce' and 'addKToMatrix'
+        physics.addObject("EulerImplicitSolver", name="eulerImplicit")
+        physics.addObject("SparseLDLSolver", name="solver", template="CompressedRowSparseMatrixd")
+
+    position = np.random.uniform(-1, 1, (n_particles, 3))
+    velocity = np.zeros_like(position)
+    length = np.random.uniform(0.8, 1.2, size=(n_particles, 1))
+    stiffness = 100.0
+
+    particles = physics.addChild("Particles")
+    particles.addObject("MechanicalObject", name="state", template="Vec3d", position=position, velocity=velocity, showObject=True)
+    particles.addObject("UniformMass", name="mass", totalMass=n_particles)
+
+    if not use_sofa:  # Use the force field implemented with JAX
+        particles.addObject(JaxForceField(length=length, stiffness=stiffness))
+    else:  # Use a SOFA equivalent for comparison
+        root.addObject("MechanicalObject", name="origin", template="Vec3d", position="0 0 0")
+        particles.addObject("SpringForceField", name="force", object1="@/origin", object2="@/Physics/Particles/state", indices1=np.zeros(n_particles, dtype=np.int32), indices2=np.arange(n_particles), length=length, stiffness=stiffness*np.ones(n_particles), damping=np.zeros(n_particles))
+
+
+def main():
+    import argparse
+    import SofaRuntime
+    import SofaImGui
+    import Sofa.Gui
+
+    parser = argparse.ArgumentParser(description="Example of a scene using a ForceField implemented with JAX")
+    parser.add_argument("--method", default="implicit-matrix-assembly", help="must be 'explicit', 'implicit-matrix-free' or 'implicit-matrix-assembly'")
+    parser.add_argument("--particles", type=int, default=1000, help="number of particles (default 1000)")
+    parser.add_argument("--use-sofa", action="store_true", help="use a force field from SOFA instead of the one implemented with JAX")
+    args = parser.parse_args()
+
+    root=Sofa.Core.Node("root")
+    createScene(root, method=args.method, n_particles=args.particles, use_sofa=args.use_sofa)
+    Sofa.Simulation.initRoot(root)
+
+    Sofa.Gui.GUIManager.Init("myscene", "imgui")
+    Sofa.Gui.GUIManager.createGUI(root, __file__)
+    Sofa.Gui.GUIManager.SetDimension(1600, 900)
+    Sofa.Gui.GUIManager.MainLoop(root)
+    Sofa.Gui.GUIManager.closeGUI()
+
+
+if __name__ == "__main__":
+    main()