From 643c1248b895457da9ce9b3d492cd41b763c630f Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Mon, 23 Oct 2023 10:00:33 -0700
Subject: [PATCH] Skip flex edge Jacobian if not needed.

When equality constraints and built-in passive forces are not present, edge Jacobians are unused. This gives a 3x speedup in `mj_kinematics` in `floppy_flex` (from 0.2ms/step to 0.06ms/step).

Changed the default options of `floppy_flex` to match those of `floppy` for a fair comparison.

PiperOrigin-RevId: 575850622
Change-Id: I4c618536fd17ef930153bbd1f1da7d9e0e44abbe
---
 model/plugin/elasticity/floppy_flex.xml |  4 ++--
 src/engine/engine_core_smooth.c         | 10 ++++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/model/plugin/elasticity/floppy_flex.xml b/model/plugin/elasticity/floppy_flex.xml
index a8645f2f..b6a262c3 100644
--- a/model/plugin/elasticity/floppy_flex.xml
+++ b/model/plugin/elasticity/floppy_flex.xml
@@ -22,7 +22,7 @@
 
   <compiler autolimits="true"/>
 
-  <option solver="CG" tolerance="1e-6" timestep=".001" integrator="implicitfast"/>
+  <option solver="Newton" tolerance="1e-6" timestep=".001" integrator="Euler"/>
 
   <size memory="100M"/>
 
@@ -34,11 +34,11 @@
     <flexcomp type="grid" count="24 4 4" spacing=".1 .1 .1" pos=".1 0 1.5"
               radius=".0" rgba="0 .7 .7 1" name="softbody" dim="3" mass="7">
       <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001" selfcollide="none"/>
-      <edge damping="1"/>
       <plugin plugin="mujoco.elasticity.solid">
         <config key="poisson" value="0.2"/>
         <!--Units are in Pa (SI)-->
         <config key="young" value="5e4"/>
+        <config key="damping" value="0.002"/>
       </plugin>
     </flexcomp>
 
diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
index af8b10e6..a8194d82 100644
--- a/src/engine/engine_core_smooth.c
+++ b/src/engine/engine_core_smooth.c
@@ -541,6 +541,11 @@ void mj_flex(const mjModel* m, mjData* d) {
       continue;
     }
 
+    // skip Jacobian if no built-in passive force is needed
+    int skipjacobian = !m->flex_edgeequality[f] &&
+                       !m->flex_edgedamping[f] &&
+                       !m->flex_edgestiffness[f];
+
     // process edges of this flex
     int vbase = m->flex_vertadr[f];
     int ebase = m->flex_edgeadr[f];
@@ -557,6 +562,11 @@ void mj_flex(const mjModel* m, mjData* d) {
       mju_sub3(vec, pos2, pos1);
       d->flexedge_length[ebase+e] = mju_normalize3(vec);
 
+      // skip Jacobian if not needed
+      if (skipjacobian) {
+        continue;
+      }
+
       // sparse edge Jacobian
       if (issparse) {
         // set rowadr