From 40d927c9fa24992e399e10d66fefed899ad3ff3e Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Fri, 27 Oct 2023 07:17:14 -0700
Subject: [PATCH] Move 2d and 3d plugin force computation to `elasticity.h`.

Change physical properties of `floppy_flex.xml` to resemble `floppy.xml`.

Make `Membrane` backward compatible with `Composite`.

PiperOrigin-RevId: 577190384
Change-Id: I5bc56ec62faa134156e839bec45d685c22eba5c4
---
 model/plugin/elasticity/floppy_flex.xml |  4 +-
 plugin/elasticity/elasticity.cc         |  6 +-
 plugin/elasticity/elasticity.h          | 55 +++++++++++++++++
 plugin/elasticity/membrane.cc           | 81 +++++++++++++------------
 plugin/elasticity/membrane.h            |  9 ++-
 plugin/elasticity/solid.cc              | 76 +++++++----------------
 plugin/elasticity/solid.h               |  1 +
 7 files changed, 131 insertions(+), 101 deletions(-)

diff --git a/model/plugin/elasticity/floppy_flex.xml b/model/plugin/elasticity/floppy_flex.xml
index b6a262c3..8609e3fa 100644
--- a/model/plugin/elasticity/floppy_flex.xml
+++ b/model/plugin/elasticity/floppy_flex.xml
@@ -22,7 +22,7 @@
 
   <compiler autolimits="true"/>
 
-  <option solver="Newton" tolerance="1e-6" timestep=".001" integrator="Euler"/>
+  <option solver="CG" tolerance="1e-6" timestep=".001"/>
 
   <size memory="100M"/>
 
@@ -32,7 +32,7 @@
 
   <worldbody>
     <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">
+              radius=".0" rgba="0 .7 .7 1" name="softbody" dim="3" mass="25">
       <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001" selfcollide="none"/>
       <plugin plugin="mujoco.elasticity.solid">
         <config key="poisson" value="0.2"/>
diff --git a/plugin/elasticity/elasticity.cc b/plugin/elasticity/elasticity.cc
index 088f8b6b..da95793e 100644
--- a/plugin/elasticity/elasticity.cc
+++ b/plugin/elasticity/elasticity.cc
@@ -64,8 +64,10 @@ int CreateStencils(std::vector<T>& elements,
         elements[t].edges[e] = it->second;
       }
 
-      if (!edgeidx.empty()) {
-        assert(elements[t].edges[e] == edgeidx[T::kNumEdges*t+e]);
+      if (!edgeidx.empty()) {  // SHOULD NOT OCCUR
+        if (elements[t].edges[e] != edgeidx[T::kNumEdges*t+e]) {
+          mju_error("edge ordering is incoherent between flex and plugin");
+        }
       }
     }
   }
diff --git a/plugin/elasticity/elasticity.h b/plugin/elasticity/elasticity.h
index 7bc35720..52670c75 100644
--- a/plugin/elasticity/elasticity.h
+++ b/plugin/elasticity/elasticity.h
@@ -48,6 +48,13 @@ inline void UpdateSquaredLengths(std::vector<mjtNum>& len,
   }
 }
 
+inline void UpdateSquaredLengthsFlex(std::vector<mjtNum>& len,
+                                     const mjtNum* flexedge_length) {
+  for (int e = 0; e < len.size(); e++) {
+    len[e] = flexedge_length[e]*flexedge_length[e];
+  }
+}
+
 struct Stencil2D {
   static constexpr int kNumEdges = 3;
   static constexpr int kNumVerts = 3;
@@ -78,6 +85,54 @@ void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3],
   }
 }
 
+template <typename T>
+inline void ComputeForce(mjtNum* qfrc_passive,
+                         const std::vector<T>& elements,
+                         const std::vector<mjtNum>& metric,
+                         const std::vector<mjtNum>& elongationglob,
+                         const mjtNum* xpos) {
+  for (int t = 0; t < elements.size(); t++)  {
+    const int* v = elements[t].vertices;
+
+    // compute length gradient with respect to dofs
+    mjtNum gradient[T::kNumEdges][2][3];
+    GradSquaredLengths<T>(gradient, xpos, v);
+
+    // extract elongation of edges belonging to this element
+    mjtNum elongation[T::kNumEdges];
+    for (int e = 0; e < T::kNumEdges; e++) {
+      int idx = elements[t].edges[e];
+      elongation[e] = elongationglob[idx];
+    }
+
+    // we now multiply the elongations by the precomputed metric tensor,
+    // notice that if metric=diag(1/reference) then this would yield a
+    // mass-spring model
+
+    // compute local force
+    mjtNum force[T::kNumVerts*3] = {0};
+    int offset = T::kNumEdges*T::kNumEdges;
+    for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
+      for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
+        for (int i = 0; i < 2; i++) {
+          for (int x = 0; x < 3; x++) {
+            force[3 * T::edge[ed2][i] + x] +=
+                elongation[ed1] * gradient[ed2][i][x] *
+                metric[offset * t + T::kNumEdges * ed1 + ed2];
+          }
+        }
+      }
+    }
+
+    // insert into global force
+    for (int i = 0; i < T::kNumVerts; i++) {
+      for (int x = 0; x < 3; x++) {
+        qfrc_passive[3*v[i]+x] -= force[3*i+x];
+      }
+    }
+  }
+}
+
 // compute metric tensor of edge lengths inner product
 template <typename T>
 void inline MetricTensor(std::vector<mjtNum>& metric, int idx, mjtNum mu,
diff --git a/plugin/elasticity/membrane.cc b/plugin/elasticity/membrane.cc
index 38470139..d0a93454 100644
--- a/plugin/elasticity/membrane.cc
+++ b/plugin/elasticity/membrane.cc
@@ -84,10 +84,12 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
     mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
     mjtNum thick =
         strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
+    mjtNum damp =
+            strtod(mj_getPluginConfig(m, instance, "damping"), nullptr);
     std::vector<int> face, edge;
     String2Vector(mj_getPluginConfig(m, instance, "face"), face);
     String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
-    return Membrane(m, d, instance, nu, E, thick, face, edge);
+    return Membrane(m, d, instance, nu, E, thick, damp, face, edge);
   } else {
     mju_warning("Invalid parameter specification in shell plugin");
     return std::nullopt;
@@ -96,9 +98,10 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
 
 // plugin constructor
 Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
-                   mjtNum E, mjtNum thick, const std::vector<int>& simplex,
+                   mjtNum E, mjtNum thick, mjtNum damp,
+                   const std::vector<int>& simplex,
                    const std::vector<int>& edgeidx)
-    : thickness(thick) {
+    : f0(-1), damping(damp), thickness(thick) {
   // count plugin bodies
   nv = ne = 0;
   for (int i = 1; i < m->nbody; i++) {
@@ -152,50 +155,48 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
     // compute metric tensor
     MetricTensor<Stencil2D>(metric, t, mu, la, basis);
   }
+
+  // allocate array
+  ne = edges.size();
+  reference.assign(ne, 0);
+  deformed.assign(ne, 0);
+  previous.assign(ne, 0);
+  elongation.assign(ne, 0);
+
+  // compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
+  UpdateSquaredLengths(reference, edges, m->body_pos+3*i0);
+
+  // save previous lengths
+  previous = reference;
 }
 
 void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
-  for (int t = 0; t < nt; t++)  {
-    int* v = elements[t].vertices;
+  mjtNum kD = damping / m->opt.timestep;
 
-    // compute length gradient with respect to dofs
-    mjtNum gradient[kNumEdges][2][3];
-    GradSquaredLengths<Stencil2D>(gradient, d->xpos+3*i0, v);
+  // update edge lengths
+  if (f0 < 0) {
+    UpdateSquaredLengths(deformed, edges, d->xpos+3*i0);
+  } else {
+    UpdateSquaredLengthsFlex(deformed,
+                             d->flexedge_length + m->flex_edgeadr[f0]);
+  }
 
-    // compute elongation
-    mjtNum elongation[kNumEdges];
-    for (int e = 0; e < kNumEdges; e++) {
-      int idx = elements[t].edges[e] + m->flex_edgeadr[f0];
-      mjtNum deformed = d->flexedge_length[idx]*d->flexedge_length[idx];
-      mjtNum reference = m->flexedge_length0[idx]*m->flexedge_length0[idx];
-      elongation[e] = deformed - reference;
-    }
+  // we add generalized Rayleigh damping as decribed in Section 5.2 of
+  // Kharevych et al., "Geometric, Variational Integrators for Computer
+  // Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
 
-    // we now multiply the elongations by the precomputed metric tensor,
-    // notice that if metric=diag(1/reference) then this would yield a
-    // mass-spring model
+  for (int idx = 0; idx < ne; idx++) {
+    elongation[idx] = deformed[idx] - reference[idx] +
+                    ( deformed[idx] -  previous[idx] ) * kD;
+  }
 
-    // compute local force
-    mjtNum force[kNumVerts*3] = {0};
-    int offset = kNumEdges*kNumEdges;
-    for (int ed1 = 0; ed1 < kNumEdges; ed1++) {
-      for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
-        for (int i = 0; i < 2; i++) {
-          for (int x = 0; x < 3; x++) {
-            force[3 * Stencil2D::edge[ed2][i] + x] +=
-                elongation[ed1] * gradient[ed2][i][x] *
-                metric[offset * t + kNumEdges * ed1 + ed2];
-          }
-        }
-      }
-    }
+  // compute gradient of elastic energy and insert into passive force
+  ComputeForce<Stencil2D>(d->qfrc_passive + m->body_dofadr[i0], elements,
+                          metric, elongation, d->xpos + 3 * i0);
 
-    // insert into global force
-    for (int i = 0; i < kNumVerts; i++) {
-      for (int x = 0; x < 3; x++) {
-        d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x];
-      }
-    }
+  // update stored lengths
+  if (kD > 0) {
+    previous = deformed;
   }
 }
 
@@ -208,7 +209,7 @@ void Membrane::RegisterPlugin() {
   plugin.name = "mujoco.elasticity.membrane";
   plugin.capabilityflags |= mjPLUGIN_PASSIVE;
 
-  const char* attributes[] = {"face", "edge", "young", "poisson", "thickness"};
+  const char* attributes[] = {"face", "edge", "young", "poisson", "thickness", "damping"};
   plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
   plugin.attributes = attributes;
   plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
diff --git a/plugin/elasticity/membrane.h b/plugin/elasticity/membrane.h
index 10fa2b60..4b39d426 100644
--- a/plugin/elasticity/membrane.h
+++ b/plugin/elasticity/membrane.h
@@ -52,13 +52,18 @@ class Membrane {
   std::vector<std::pair<int, int> > edges;  // edge to vertex map     (ne x 2)
 
   // precomputed quantities
-  std::vector<mjtNum> metric;               // geom-induced metric     (nt x 9)
+  std::vector<mjtNum> metric;               // geom-induced metric    (nt x 9)
+  std::vector<mjtNum> reference;            // reference lengths      (ne x 1)
+  std::vector<mjtNum> deformed;             // deformed lengths       (ne x 1)
+  std::vector<mjtNum> previous;             // previous-step lengths  (ne x 1)
+  std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
 
+  mjtNum damping;
   mjtNum thickness;
 
  private:
   Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-           mjtNum thick, const std::vector<int>& simplex,
+           mjtNum thick, mjtNum damp, const std::vector<int>& simplex,
            const std::vector<int>& edgeidx);
 };
 
diff --git a/plugin/elasticity/solid.cc b/plugin/elasticity/solid.cc
index d2c5c853..cca4ff78 100644
--- a/plugin/elasticity/solid.cc
+++ b/plugin/elasticity/solid.cc
@@ -165,75 +165,41 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   reference.assign(ne, 0);
   deformed.assign(ne, 0);
   previous.assign(ne, 0);
+  elongation.assign(ne, 0);
 
-  // compute edge lengths at equilibrium
+  // compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
   UpdateSquaredLengths(reference, edges, m->body_pos+3*i0);
+
+  // save previous lengths
   previous = reference;
 }
 
 void Solid::Compute(const mjModel* m, mjData* d, int instance) {
-  // update edges if no flex
+  mjtNum kD = damping / m->opt.timestep;
+
+  // update edge lengths
   if (f0 < 0) {
     UpdateSquaredLengths(deformed, edges, d->xpos+3*i0);
+  } else {
+    UpdateSquaredLengthsFlex(deformed,
+                             d->flexedge_length + m->flex_edgeadr[f0]);
   }
 
-  // loop over all elements
-  for (int t = 0; t < nt; t++)  {
-    int* v = elements[t].vertices;
+  // we add generalized Rayleigh damping as decribed in Section 5.2 of
+  // Kharevych et al., "Geometric, Variational Integrators for Computer
+  // Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
 
-    // compute length gradient with respect to dofs
-    mjtNum gradient[kNumEdges][2][3];
-    GradSquaredLengths<Stencil3D>(gradient, d->xpos+3*i0, v);
-
-    // we add generalized Rayleigh damping as decribed in Section 5.2 of
-    // Kharevych et al., "Geometric, Variational Integrators for Computer
-    // Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
-
-    // compute elongation
-    mjtNum elongation[kNumEdges];
-    for (int e = 0; e < kNumEdges; e++) {
-      if (f0 < 0) {
-        int idx = elements[t].edges[e];
-        mjtNum kD = damping / m->opt.timestep;
-        elongation[e] = deformed[idx] - reference[idx] +
-                      ( deformed[idx] -  previous[idx] ) * kD;
-      } else {
-        int idx = elements[t].edges[e] + m->flex_edgeadr[f0];
-        mjtNum deformed = d->flexedge_length[idx]*d->flexedge_length[idx];
-        mjtNum reference = m->flexedge_length0[idx]*m->flexedge_length0[idx];
-        elongation[e] = deformed - reference;
-      }
-    }
-
-    // we now multiply the elongations by the precomputed metric tensor,
-    // notice that if metric=diag(1/reference) then this would yield a
-    // mass-spring model
-
-    // compute local force
-    mjtNum force[kNumVerts*3] = {0};
-    int offset = kNumEdges*kNumEdges;
-    for (int ed1 = 0; ed1 < kNumEdges; ed1++) {
-      for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
-        for (int i = 0; i < 2; i++) {
-          for (int x = 0; x < 3; x++) {
-            force[3 * Stencil3D::edge[ed2][i] + x] +=
-                elongation[ed1] * gradient[ed2][i][x] *
-                metric[offset * t + kNumEdges * ed1 + ed2];
-          }
-        }
-      }
-    }
-
-    // insert into global force
-    for (int i = 0; i < kNumVerts; i++) {
-      for (int x = 0; x < 3; x++) {
-        d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x];
-      }
-    }
+  for (int idx = 0; idx < ne; idx++) {
+    elongation[idx] = deformed[idx] - reference[idx] +
+                    ( deformed[idx] -  previous[idx] ) * kD;
   }
 
+  // compute gradient of elastic energy and insert into passive force
+  ComputeForce<Stencil3D>(d->qfrc_passive + m->body_dofadr[i0], elements,
+                          metric, elongation, d->xpos + 3 * i0);
+
   // update stored lengths
-  if (f0 < 0) {
+  if (kD > 0) {
     previous = deformed;
   }
 }
diff --git a/plugin/elasticity/solid.h b/plugin/elasticity/solid.h
index a2505d44..ae1f2fe5 100644
--- a/plugin/elasticity/solid.h
+++ b/plugin/elasticity/solid.h
@@ -54,6 +54,7 @@ class Solid {
   std::vector<mjtNum> reference;            // reference lengths      (ne x 1)
   std::vector<mjtNum> deformed;             // deformed lengths       (ne x 1)
   std::vector<mjtNum> previous;             // previous-step lengths  (ne x 1)
+  std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
 
   mjtNum damping;