From 8f1ea05bef131e71bb374678ca0a777322dea2ae Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 29 Feb 2024 07:23:46 -0800
Subject: [PATCH] Add disc mesh generation to user_flexcomp.

PiperOrigin-RevId: 611462985
Change-Id: I0fc195d8733172f5f4347b04b5de0edfdd12613d
---
 model/plugin/elasticity/pancake_flex.xml | 50 ++++++++++++++++++
 src/user/user_flexcomp.cc                | 67 ++++++++++++++++++++----
 src/user/user_flexcomp.h                 |  3 ++
 src/xml/xml_native_reader.cc             |  2 +
 4 files changed, 112 insertions(+), 10 deletions(-)
 create mode 100644 model/plugin/elasticity/pancake_flex.xml

diff --git a/model/plugin/elasticity/pancake_flex.xml b/model/plugin/elasticity/pancake_flex.xml
new file mode 100644
index 00000000..38e839f9
--- /dev/null
+++ b/model/plugin/elasticity/pancake_flex.xml
@@ -0,0 +1,50 @@
+<!-- Copyright 2024 DeepMind Technologies Limited
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<mujoco model="Pancake">
+  <include file="scene.xml"/>
+
+  <extension>
+    <plugin plugin="mujoco.elasticity.shell"/>
+  </extension>
+
+  <option timestep="0.001" solver="CG" tolerance="1e-6" integrator="implicitfast"/>
+
+  <size memory="10M"/>
+
+  <visual>
+    <map stiffness="100"/>
+  </visual>
+
+  <default>
+    <default class="wall">
+      <geom type="plane" size=".5 .5 .05"/>
+    </default>
+  </default>
+
+  <worldbody>
+    <flexcomp type="disc" count="9 9 1" spacing=".1 .1 .1" pos="0 0 .5"
+              radius=".01" mass=".5" name="plate" dim="2">
+      <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
+      <edge equality="true" damping="10"/>
+      <plugin plugin="mujoco.elasticity.shell">
+        <config key="poisson" value="0"/>
+        <config key="thickness" value="8e-3"/>
+        <!--Units are in Pa (SI)-->
+        <config key="young" value="3e5"/>
+      </plugin>
+    </flexcomp>
+  </worldbody>
+</mujoco>
diff --git a/src/user/user_flexcomp.cc b/src/user/user_flexcomp.cc
index eefbc391..94fb3b25 100644
--- a/src/user/user_flexcomp.cc
+++ b/src/user/user_flexcomp.cc
@@ -84,6 +84,7 @@ mjCFlexcomp::mjCFlexcomp(void) {
 // make flexcomp object
 bool mjCFlexcomp::Make(mjCModel* model, mjmBody* body, char* error, int error_sz) {
   mjmFlex* dflex = def.spec.flex;
+  int dim = dflex->dim;
   bool radial = (type==mjFCOMPTYPE_BOX ||
                  type==mjFCOMPTYPE_CYLINDER ||
                  type==mjFCOMPTYPE_ELLIPSOID);
@@ -96,9 +97,14 @@ bool mjCFlexcomp::Make(mjCModel* model, mjmBody* body, char* error, int error_sz
     return comperr(error, "Parent body must have name", error_sz);
   }
 
+  // check dim
+  if (dim<1 || dim>3) {
+    return comperr(error, "Invalid dim, must be between 1 and 3", error_sz);
+  }
+
   // check counts
   for (int i=0; i<3; i++) {
-    if (count[i]<1 || (radial && count[i]<2)) {
+    if (count[i]<1 || ((radial && count[i]<2) && dim==3)) {
       return comperr(error, "Count too small", error_sz);
     }
   }
@@ -142,6 +148,11 @@ bool mjCFlexcomp::Make(mjCModel* model, mjmBody* body, char* error, int error_sz
     res = MakeBox(error, error_sz);
     break;
 
+  case mjFCOMPTYPE_SQUARE:
+  case mjFCOMPTYPE_DISC:
+    res = MakeSquare(error, error_sz);
+    break;
+
   case mjFCOMPTYPE_MESH:
     res = MakeMesh(model, error, error_sz);
     break;
@@ -161,12 +172,6 @@ bool mjCFlexcomp::Make(mjCModel* model, mjmBody* body, char* error, int error_sz
     return false;
   }
 
-  // get dim and check
-  int dim = dflex->dim;
-  if (dim<1 || dim>3) {
-    return comperr(error, "Invalid dim, must be between 1 and 3", error_sz);
-  }
-
   // force flatskin shading for box, cylinder and 3D grid
   if (type==mjFCOMPTYPE_BOX || type==mjFCOMPTYPE_CYLINDER ||
       (type==mjFCOMPTYPE_GRID && dim==3)) {
@@ -540,12 +545,15 @@ bool mjCFlexcomp::MakeGrid(char* error, int error_sz) {
 
   // 2D
   else if (dim==2) {
-    int quad2tri[2][3] = {{0, 1, 2}, {0, 2, 3}};
     for (int ix=0; ix<count[0]; ix++) {
       for (int iy=0; iy<count[1]; iy++) {
+        int quad2tri[2][3] = {{0, 1, 2}, {0, 2, 3}};
+
         // add point
-        point.push_back(spacing[0]*(ix - 0.5*(count[0]-1)));
-        point.push_back(spacing[1]*(iy - 0.5*(count[1]-1)));
+        mjtNum pos[2] = {spacing[0]*(ix- 0.5*(count[0]-1)),
+                         spacing[1]*(iy- 0.5*(count[1]-1))};
+        point.push_back(pos[0]);
+        point.push_back(pos[1]);
         point.push_back(0);
 
         // add texture coordinates, if not specified explicitly
@@ -554,6 +562,14 @@ bool mjCFlexcomp::MakeGrid(char* error, int error_sz) {
           texcoord.push_back(iy/(mjtNum)mjMAX(count[1]-1, 1));
         }
 
+        // flip triangles if radial projection is requested
+        if (((pos[0] < -mjEPS && pos[1] > -mjEPS) ||
+             (pos[0] > -mjEPS && pos[1] < -mjEPS)) &&
+            type == mjFCOMPTYPE_DISC) {
+          quad2tri[0][2] = 3;
+          quad2tri[1][0] = 1;
+        }
+
         // add elements
         if (ix<count[0]-1 && iy<count[1]-1) {
           int vert[4] = {
@@ -695,6 +711,37 @@ void mjCFlexcomp::BoxProject(double* pos, int ix, int iy, int iz) {
 
 
 
+// make 2d square or disc
+bool mjCFlexcomp::MakeSquare(char* error, int error_sz) {
+  // set 2D
+  def.spec.flex->dim = 2;
+
+  // create square
+  if (!MakeGrid(error, error_sz)) {
+    return false;
+  }
+
+  // do projection
+  if (type==mjFCOMPTYPE_DISC) {
+    double size[2] = {
+      0.5*spacing[0]*(count[0]-1),
+      0.5*spacing[1]*(count[1]-1),
+    };
+
+    for (int i=0; i<point.size()/3; i++) {
+      mjtNum* pos = point.data() + i*3;
+      double L0 = mjMAX(mju_abs(pos[0]), mju_abs(pos[1]));
+      mjuu_normvec(pos, 2);
+      pos[0] *= size[0]*L0;
+      pos[1] *= size[1]*L0;
+    }
+  }
+
+  return true;
+}
+
+
+
 // make 3d box, ellipsoid or cylinder
 bool mjCFlexcomp::MakeBox(char* error, int error_sz) {
   double pos[3];
diff --git a/src/user/user_flexcomp.h b/src/user/user_flexcomp.h
index d966af68..896a55a5 100644
--- a/src/user/user_flexcomp.h
+++ b/src/user/user_flexcomp.h
@@ -29,6 +29,8 @@ typedef enum _mjtFcompType {
   mjFCOMPTYPE_BOX,
   mjFCOMPTYPE_CYLINDER,
   mjFCOMPTYPE_ELLIPSOID,
+  mjFCOMPTYPE_SQUARE,
+  mjFCOMPTYPE_DISC,
   mjFCOMPTYPE_MESH,
   mjFCOMPTYPE_GMSH,
   mjFCOMPTYPE_DIRECT,
@@ -44,6 +46,7 @@ class mjCFlexcomp {
 
   bool MakeGrid(char* error, int error_sz);
   bool MakeBox(char* error, int error_sz);
+  bool MakeSquare(char* error, int error_sz);
   bool MakeMesh(mjCModel* model, char* error, int error_sz);
   bool MakeGMSH(mjCModel* model, char* error, int error_sz);
   void LoadGMSH(mjCModel* model, mjResource* resource);
diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc
index 26e5c18f..774df796 100644
--- a/src/xml/xml_native_reader.cc
+++ b/src/xml/xml_native_reader.cc
@@ -750,6 +750,8 @@ const mjMap fcomp_map[mjNFCOMPTYPES] = {
   {"box",         mjFCOMPTYPE_BOX},
   {"cylinder",    mjFCOMPTYPE_CYLINDER},
   {"ellipsoid",   mjFCOMPTYPE_ELLIPSOID},
+  {"square",      mjFCOMPTYPE_SQUARE},
+  {"disc",        mjFCOMPTYPE_DISC},
   {"mesh",        mjFCOMPTYPE_MESH},
   {"gmsh",        mjFCOMPTYPE_GMSH},
   {"direct",      mjFCOMPTYPE_DIRECT}