Change flex constraints to eigenmodes of the stiffness matrix.

This provides a reduction from 26 to 18 constraints for trilinear and from 162 to 75 for quadratic. The assembly of the constraints becomes trivial. In total the speedup for a trilinear 3x3x3 grid is about 3x.

PiperOrigin-RevId: 902502398
Change-Id: I764772c7adef78da5a644f64701f842d36e4b543

src/user/user_flexcomp.cc

@@ -691,6 +691,7 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
       mjs_setString(pe->name1, name.c_str());
     } else if (equality == 3) {
       // create one strain constraint per cell, storing cell index in eq_data
+      flex->has_strain_eq = true;
       int cell_cx = flex->spec.cellcount[0];
       int cell_cy = flex->spec.cellcount[1];
       int cell_cz = flex->spec.cellcount[2];

src/user/user_mesh.cc

@@ -3818,6 +3818,48 @@ void inline ComputeLinearStiffness(std::vector<double>& K,
   }
 }
 
+
+// Eigendecompose cell stiffness matrix and store scaled eigenvectors.
+// K_cell is n×n stored (negative convention: K_stored = -K_physical).
+// Output layout in `out`:
+//   [0]: neig (as double)
+//   [1 .. neig*n]: sqrt(λ_phys_i) * v_i, row-major
+// Returns number of retained eigenmodes.
+static int EigendecomposeStiffness(const double* K_cell_data,
+                                   double* out, int ndof) {
+  // copy K_cell for in-place decomposition
+  std::vector<double> mat(K_cell_data, K_cell_data + ndof * ndof);
+  std::vector<double> eigval(ndof);
+  std::vector<double> eigvec(ndof * ndof);
+
+  mjuu_eigendecompose(mat.data(), eigval.data(), eigvec.data(), ndof);
+
+  // K_stored = -K_physical, so physical eigenvalue = -eigval[i]
+  // retain modes where physical eigenvalue > threshold
+  double max_eigval = 0;
+  for (int i = 0; i < ndof; i++) {
+    max_eigval = std::max(max_eigval, std::abs(eigval[i]));
+  }
+  double threshold = max_eigval * 1e-8;
+
+  int neig = 0;
+  for (int i = 0; i < ndof; i++) {
+    double lambda_phys = -eigval[i];  // negate to get physical eigenvalue
+    if (lambda_phys > threshold) {
+      // store sqrt(λ) * eigenvector (column i of eigvec matrix)
+      double scale = std::sqrt(lambda_phys);
+      for (int j = 0; j < ndof; j++) {
+        out[1 + neig * ndof + j] = scale * eigvec[j * ndof + i];
+      }
+      neig++;
+    }
+  }
+
+  out[0] = static_cast<double>(neig);
+  return neig;
+}
+
+
 //------------------ class mjCFlex implementation --------------------------------------------------
 
 // constructor
@@ -4344,7 +4386,11 @@ void mjCFlex::Compile(const mjVFS* vfs) {
     stiffness_cached = LoadCachedStiffness();
   }
 
-  if (!stiffness_cached && young > 0 && interpolated) {
+  if (!stiffness_cached && interpolated && (young > 0 || has_strain_eq)) {
+    // use young=1 for strain constraints (eigenvectors are geometry-only)
+    double K_young = has_strain_eq ? 1e1 : young;
+    double K_poisson = has_strain_eq ? 0.3 : poisson;
+
     int npc = pow(spec.order + 1, 3);  // nodes per cell
     int ndof_cell = 3 * npc;
     int cx = spec.cellcount[0], cy = spec.cellcount[1], cz = spec.cellcount[2];
@@ -4379,11 +4425,17 @@ void mjCFlex::Compile(const mjVFS* vfs) {
 
           // compute per-cell stiffness
           std::vector<double> K_cell(ndof_cell * ndof_cell, 0);
-          ComputeLinearStiffness(K_cell, cell_pos.data(), young, poisson, spec.order);
+          ComputeLinearStiffness(K_cell, cell_pos.data(), K_young, K_poisson, spec.order);
+          double* out = stiffness.data() + cell_idx * ndof_cell * ndof_cell;
 
-          // copy into global stiffness array
-          mjuu_copyvec(stiffness.data() + cell_idx * ndof_cell * ndof_cell,
-                       K_cell.data(), ndof_cell * ndof_cell);
+          if (has_strain_eq) {
+            // eigendecompose: store [neig, sqrt(λ)*v_1, sqrt(λ)*v_2, ...]
+            std::fill(out, out + ndof_cell * ndof_cell, 0.0);
+            EigendecomposeStiffness(K_cell.data(), out, ndof_cell);
+          } else {
+            // store raw K for passive forces
+            std::copy(K_cell.begin(), K_cell.end(), out);
+          }
         }
       }
     }

src/user/user_model.cc

@@ -3657,6 +3657,8 @@ void mjCModel::CopyObjects(mjModel* m) {
         int b1 = pfl->vertbodyid[pfl->edge[k].first];
         int b2 = pfl->vertbodyid[pfl->edge[k].second];
         m->flexedge_rigid[edge_adr+k] = (bodies_[b1]->weldid == bodies_[b2]->weldid);
+      } else {
+        m->flexedge_rigid[edge_adr+k] = 0;
       }
     }
 

src/user/user_objects.h

@@ -983,6 +983,7 @@ class mjCFlex_ : public mjCBase {
   std::vector<int> edgeidx_;              // element edge ids
   std::vector<double> stiffness;          // elasticity stiffness matrix
   std::vector<double> bending;            // bending stiffness matrix
+  bool has_strain_eq = false;             // true if strain constraints reference this flex
 
   // variable-size data
   std::vector<std::string> vertbody_;     // vertex body names

src/user/user_util.cc

@@ -754,6 +754,81 @@ int mjuu_eig3(double eigval[3], double eigvec[9], double quat[4], const double m
   return iter;
 }
 
+
+// Jacobi eigenvalue decomposition of symmetric n×n matrix.
+// On output, eigenvalues are in eigval and eigenvectors are columns of eigvec.
+// Both arrays must be pre-allocated: eigval[n], eigvec[n*n].
+// The input matrix mat is destroyed.
+int mjuu_eigendecompose(double* mat, double* eigval, double* eigvec, int n) {
+  // initialize eigvec to identity
+  std::fill(eigvec, eigvec + n*n, 0.0);
+  for (int i = 0; i < n; i++) {
+    eigvec[i*n + i] = 1.0;
+  }
+
+  const int max_sweeps = 200;
+  const double tol = 1e-12;
+
+  int sweep;
+  for (sweep = 0; sweep < max_sweeps; sweep++) {
+    // check convergence: sum of squared off-diagonal elements
+    double off_diag = 0;
+    for (int i = 0; i < n; i++) {
+      for (int j = i+1; j < n; j++) {
+        off_diag += mat[i*n + j] * mat[i*n + j];
+      }
+    }
+    if (off_diag < tol * tol) break;
+
+    // sweep over all off-diagonal pairs
+    for (int p = 0; p < n; p++) {
+      for (int q = p+1; q < n; q++) {
+        double apq = mat[p*n + q];
+        if (std::abs(apq) < tol * 1e-3) continue;
+
+        // compute rotation angle
+        double app = mat[p*n + p];
+        double aqq = mat[q*n + q];
+        double tau = (aqq - app) / (2.0 * apq);
+        double t = (tau >= 0 ? 1.0 : -1.0) /
+                   (std::abs(tau) + std::sqrt(1.0 + tau*tau));
+        double c = 1.0 / std::sqrt(1.0 + t*t);
+        double s = t * c;
+
+        // update matrix (Jacobi rotation)
+        mat[p*n + p] -= t * apq;
+        mat[q*n + q] += t * apq;
+        mat[p*n + q] = 0;
+        mat[q*n + p] = 0;
+
+        for (int r = 0; r < n; r++) {
+          if (r == p || r == q) continue;
+          double mrp = mat[r*n + p];
+          double mrq = mat[r*n + q];
+          mat[r*n + p] = mat[p*n + r] = c*mrp - s*mrq;
+          mat[r*n + q] = mat[q*n + r] = s*mrp + c*mrq;
+        }
+
+        // accumulate eigenvectors
+        for (int r = 0; r < n; r++) {
+          double vrp = eigvec[r*n + p];
+          double vrq = eigvec[r*n + q];
+          eigvec[r*n + p] = c*vrp - s*vrq;
+          eigvec[r*n + q] = s*vrp + c*vrq;
+        }
+      }
+    }
+  }
+
+  // extract eigenvalues from diagonal
+  for (int i = 0; i < n; i++) {
+    eigval[i] = mat[i*n + i];
+  }
+
+  return sweep;
+}
+
+
 // transform vector by pose
 void mjuu_trnVecPose(double res[3], const double pos[3], const double quat[4],
                      const double vec[3]) {
@@ -1189,10 +1264,10 @@ template<typename T> std::string VectorToString(const std::vector<T>& v) {
   return s;
 }
 
-template std::string VectorToString(const std::vector<int>& v);
-template std::string VectorToString(const std::vector<float>& v);
-template std::string VectorToString(const std::vector<double>& v);
-template std::string VectorToString(const std::vector<std::string>& v);
+template MJAPI std::string VectorToString(const std::vector<int>& v);
+template MJAPI std::string VectorToString(const std::vector<float>& v);
+template MJAPI std::string VectorToString(const std::vector<double>& v);
+template MJAPI std::string VectorToString(const std::vector<std::string>& v);
 
 namespace {
 
@@ -1258,7 +1333,7 @@ template <typename T> std::vector<T> StringToVector(char* cs) {
   return v;
 }
 
-template<> std::vector<std::string> StringToVector(const std::string& s) {
+template<> MJAPI std::vector<std::string> StringToVector(const std::string& s) {
   std::vector<std::string> v;
   std::stringstream ss(s);
   std::string word;
@@ -1268,17 +1343,17 @@ template<> std::vector<std::string> StringToVector(const std::string& s) {
   return v;
 }
 
-template std::vector<int>    StringToVector(char* cs);
-template std::vector<float>  StringToVector(char* cs);
-template std::vector<double> StringToVector(char* cs);
+template MJAPI std::vector<int>    StringToVector(char* cs);
+template MJAPI std::vector<float>  StringToVector(char* cs);
+template MJAPI std::vector<double> StringToVector(char* cs);
 
 
 template <typename T> std::vector<T> StringToVector(const std::string& s) {
   return StringToVector<T>(const_cast<char*>(s.c_str()));
 }
-template std::vector<int>    StringToVector(const std::string& s);
-template std::vector<float>  StringToVector(const std::string& s);
-template std::vector<double> StringToVector(const std::string& s);
-template std::vector<unsigned char> StringToVector(const std::string& s);
+template MJAPI std::vector<int>    StringToVector(const std::string& s);
+template MJAPI std::vector<float>  StringToVector(const std::string& s);
+template MJAPI std::vector<double> StringToVector(const std::string& s);
+template MJAPI std::vector<unsigned char> StringToVector(const std::string& s);
 
 }  // namespace mujoco::user

src/user/user_util.h

@@ -26,6 +26,8 @@
 #include <utility>
 #include <vector>
 
+#include <mujoco/mjexport.h>
+
 const double mjEPS = 1E-14;     // minimum value in various calculations
 const double mjMINMASS = 1E-6;  // minimum mass allowed
 
@@ -157,6 +159,12 @@ double mjuu_updateFrame(double quat[4], double normal[3], const double edge[3],
 // eigenvalue decomposition of symmetric 3x3 matrix
 int mjuu_eig3(double eigval[3], double eigvec[9], double quat[4], const double mat[9]);
 
+// Jacobi eigenvalue decomposition of symmetric n×n matrix
+// eigval[n]: output eigenvalues, eigvec[n*n]: output eigenvectors (columns)
+// mat[n*n]: input matrix (destroyed on output)
+// returns number of sweeps used
+MJAPI int mjuu_eigendecompose(double* mat, double* eigval, double* eigvec, int n);
+
 // transform vector by pose
 void mjuu_trnVecPose(double res[3], const double pos[3], const double quat[4], const double vec[3]);
 
@@ -166,7 +174,7 @@ const char* mjuu_fullInertia(double quat[4], double inertia[3], const double ful
 namespace mujoco::user {
 
 // utility class for handling file paths
-class FilePath {
+class MJAPI FilePath {
  public:
   FilePath() = default;
   explicit FilePath(const std::string& str) : path_(PathReduce(str)) {}
@@ -251,11 +259,11 @@ struct Cleanup {
 std::vector<uint8_t> FileToMemory(const char* filename);
 
 // convert vector to string separating elements by whitespace
-template<typename T> std::string VectorToString(const std::vector<T>& v);
+template<typename T> MJAPI std::string VectorToString(const std::vector<T>& v);
 
 // convert string to vector
-template<typename T> std::vector<T> StringToVector(char *cs);
-template<typename T> std::vector<T> StringToVector(const std::string& s);
+template<typename T> MJAPI std::vector<T> StringToVector(char *cs);
+template<typename T> MJAPI std::vector<T> StringToVector(const std::string& s);
 
 }  // namespace mujoco::user