From 6c7ed667812bee182642891c12aca9b55b512372 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 17 Apr 2026 04:10:54 -0700 Subject: [PATCH] Implement multi-cell finite element method for interpolated flexes. This change introduces a `flex_cellcount` field to `mjModel` to specify the number of cells in each dimension for interpolated flexes. The stiffness computation, passive force calculation, and Jacobian derivatives are updated to operate on a per-cell basis, significantly improving performance by localizing computations to the nodes within each cell. PiperOrigin-RevId: 901216393 Change-Id: Ic23132e609de11e71bb7fef8d1f139daad2ec264 --- doc/XMLreference.rst | 21 +- doc/XMLschema.rst | 9 + doc/changelog.rst | 9 + doc/includes/references.h | 3 + doc/modeling.rst | 10 + include/mujoco/mjmodel.h | 1 + include/mujoco/mjspec.h | 2 + include/mujoco/mjxmacro.h | 1 + model/flex/bunny_multicell.xml | 42 ++ python/mujoco/introspect/structs.py | 21 + src/engine/engine_core_constraint.c | 307 +++++---- src/engine/engine_core_smooth.c | 27 +- src/engine/engine_core_util.c | 28 + src/engine/engine_core_util.h | 3 + src/engine/engine_derivative.c | 207 +++--- src/engine/engine_passive.c | 175 ++--- src/engine/engine_passive.h | 4 +- src/engine/engine_util_misc.c | 109 ++- src/engine/engine_util_misc.h | 14 +- src/engine/engine_vis_interact.c | 26 +- src/engine/engine_vis_visualize.c | 33 +- src/user/user_flexcomp.cc | 99 ++- src/user/user_flexcomp.h | 1 + src/user/user_init.c | 3 + src/user/user_mesh.cc | 69 +- src/user/user_model.cc | 38 +- src/user/user_objects.h | 4 +- src/xml/xml_native_reader.cc | 17 +- src/xml/xml_native_writer.cc | 7 + test/engine/engine_core_constraint_test.cc | 129 +--- test/engine/engine_core_util_test.cc | 740 +++++++++++++++++++++ test/engine/engine_forward_test.cc | 4 +- test/engine/engine_support_test.cc | 515 -------------- test/engine/engine_util_misc_test.cc | 122 +++- unity/Runtime/Bindings/MjBindings.cs | 1 + wasm/codegen/generated/bindings.cc | 15 + 36 files changed, 1777 insertions(+), 1039 deletions(-) create mode 100644 model/flex/bunny_multicell.xml create mode 100644 test/engine/engine_core_util_test.cc diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index a8ebff10..9a8eba6c 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -3700,8 +3700,14 @@ saving the XML: .. _body-flexcomp-count: :at:`count`: :at-val:`int(3), "10 10 10"` - The number of automatically generated points in each dimension. This and the next attribute only apply to types grid, - box, cylinder, ellipsoid. + Specifies the number of automatically generated points in each dimension for types **grid**, **box**, **cylinder**, + and **ellipsoid**. + +.. _body-flexcomp-cellcount: + +:at:`cellcount`: :at-val:`int(3), "1 1 1"` + Specifies the number of cells in each dimension for the background interpolation grid when using **trilinear** or + **quadratic** dofs. .. _body-flexcomp-spacing: @@ -4242,6 +4248,17 @@ cases, the user will specify a :el:`flexcomp` which will then automatically cons An array of MuJoCo body names (separated by white space) to which each node belongs. The number of body names should equal the number of nodes (nnode). See the flexcomp :ref:`dof` attribute for more details. +.. _deformable-flex-cellcount: + +:at:`cellcount`: :at-val:`int(3), optional` + When using **trilinear** or **quadratic** dofs, this specifies the number of cells in each dimension for the + background interpolation grid. + +.. _deformable-flex-dof: + +:at:`dof`: :at-val:`[trilinear, quadratic], optional` + Interpolation order for the flex. + .. _flex-edge: :el-prefix:`flex/` |-| **edge** |?| diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst index 4dbc6c29..60d828be 100755 --- a/doc/XMLschema.rst +++ b/doc/XMLschema.rst @@ -1417,6 +1417,9 @@ .. grid-item:: :ref:`count` + .. grid-item:: + :ref:`cellcount` + .. grid-item:: :ref:`spacing` @@ -1657,6 +1660,12 @@ .. grid-item:: :ref:`node` + .. grid-item:: + :ref:`cellcount` + + .. grid-item:: + :ref:`dof` + .. dropdown:: :ref:`contact` :octicon:`dot` diff --git a/doc/changelog.rst b/doc/changelog.rst index 880cd3a8..b073a3f2 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,6 +2,15 @@ Changelog ========= +Upcoming version (not yet released) +----------------------------------- + +General +^^^^^^^ + +- Added :ref:`multi-cell support` for trilinear and quadratic flexes. Note that the implicit + integrator uses a dense solver for the flex degrees of freedom, which can be slow for multi-cell flexes. + Version 3.7.0 (April 14, 2026) ------------------------------ diff --git a/doc/includes/references.h b/doc/includes/references.h index 7f23a52d..614a7516 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -1315,6 +1315,7 @@ struct mjModel_ { int* flex_matid; // material id for rendering (nflex x 1) int* flex_group; // group for visibility (nflex x 1) int* flex_interp; // interpolation (0: vertex, 1: nodes) (nflex x 1) + int* flex_cellnum; // finite cell num per dimension (nflex x 3) int* flex_nodeadr; // first node address (nflex x 1) int* flex_nodenum; // number of nodes (nflex x 1) int* flex_vertadr; // first vertex address (nflex x 1) @@ -2256,6 +2257,8 @@ typedef struct mjsFlex_ { // flex specification double damping; // Rayleigh's damping double thickness; // thickness (2D only) int elastic2d; // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both + int cellcount[3]; // grid cell count for finite cell method + int order; // interpolation order (1: trilinear, 2: quadratic) // mesh properties mjStringVec* nodebody; // node body names diff --git a/doc/modeling.rst b/doc/modeling.rst index f91f698e..017e68fc 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -1442,6 +1442,16 @@ instead to specify shear and volumetric stiffnesses separately using the `Poisso `__ of the material. For more details, see the `Saint Venant-Kirchhoff `__ hyperelastic model. +**Parametrization types**. + +While the default behavior of :el:`flexcomp` produces a "full" flex where every node corresponds to a MuJoCo body, it +also supports specialized :ref:`parametrizations` for volumetric objects: **trilinear** and +**quadratic**. Instead of directly simulating all nodes, these options define a background grid of cells. The positions +of the interior vertices are computed by interpolating the positions of the cell corners. Trilinear flexes use 8-node +hexahedral cells with linear interpolation along each axis, while quadratic flexes use 27-node cells with quadratic +interpolation, allowing for curved deformation modes. These grid-based parametrizations require fewer degrees of freedom +than full flexes and can result in significantly faster simulation times, especially for large volumetric soft bodies. + **Creation and visualization**. .. code-block:: xml diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 8a45f67b..9b39cbb6 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -978,6 +978,7 @@ struct mjModel_ { int* flex_matid; // material id for rendering (nflex x 1) int* flex_group; // group for visibility (nflex x 1) int* flex_interp; // interpolation (0: vertex, 1: nodes) (nflex x 1) + int* flex_cellnum; // finite cell num per dimension (nflex x 3) int* flex_nodeadr; // first node address (nflex x 1) int* flex_nodenum; // number of nodes (nflex x 1) int* flex_vertadr; // first vertex address (nflex x 1) diff --git a/include/mujoco/mjspec.h b/include/mujoco/mjspec.h index f886da53..73d7c74f 100644 --- a/include/mujoco/mjspec.h +++ b/include/mujoco/mjspec.h @@ -454,6 +454,8 @@ typedef struct mjsFlex_ { // flex specification double damping; // Rayleigh's damping double thickness; // thickness (2D only) int elastic2d; // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both + int cellcount[3]; // grid cell count for finite cell method + int order; // interpolation order (1: trilinear, 2: quadratic) // mesh properties mjStringVec* nodebody; // node body names diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 970ba622..87a95a31 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -454,6 +454,7 @@ X ( int, flex_matid, nflex, 1 ) \ X ( int, flex_group, nflex, 1 ) \ X ( int, flex_interp, nflex, 1 ) \ + X ( int, flex_cellnum, nflex, 3 ) \ X ( int, flex_nodeadr, nflex, 1 ) \ X ( int, flex_nodenum, nflex, 1 ) \ X ( int, flex_vertadr, nflex, 1 ) \ diff --git a/model/flex/bunny_multicell.xml b/model/flex/bunny_multicell.xml new file mode 100644 index 00000000..a77f5376 --- /dev/null +++ b/model/flex/bunny_multicell.xml @@ -0,0 +1,42 @@ + + + + + + diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index 4dded269..14a44f32 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -2668,6 +2668,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='interpolation (0: vertex, 1: nodes)', array_extent=('nflex',), ), + StructFieldDecl( + name='flex_cellnum', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='finite cell num per dimension', + array_extent=('nflex', 3), + ), StructFieldDecl( name='flex_nodeadr', type=PointerType( @@ -8237,6 +8245,19 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='2D passive forces; 0: none, 1: bending, 2: stretching, 3: both', # pylint: disable=line-too-long ), + StructFieldDecl( + name='cellcount', + type=ArrayType( + inner_type=ValueType(name='int'), + extents=(3,), + ), + doc='grid cell count for finite cell method', + ), + StructFieldDecl( + name='order', + type=ValueType(name='int'), + doc='interpolation order (1: trilinear, 2: quadratic)', + ), StructFieldDecl( name='nodebody', type=PointerType( diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 170d1f3c..41ca6a39 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -432,21 +432,31 @@ static int mj_vertBodyWeight(const mjModel* m, const mjData* d, int f, int* v, return 0; } + // compute parametric coordinates of the vertex in [0, 1]^3 mjtNum coord[3] = {0, 0, 0}; for (int i = 0; i < nw; i++) { - mju_addToScl3(coord, m->flex_vert0 + 3*v[i], vweight[i]); + mju_addToScl3(coord, m->flex_vert0 + 3*v[i], vweight[i]); } + + int order = m->flex_interp[f]; + int npc = (order+1)*(order+1)*(order+1); // number of nodes per cell + + // cell lookup: get local coords and node indices + mjtNum local[3]; + int nodeindices[27]; // max npc for quadratic: 3^3 = 27 + mju_cellLookup(coord, m->flex_cellnum+3*f, order, local, nodeindices); + + // evaluate basis functions for this cell's local nodes int nstart = m->flex_nodeadr[f]; - int nend = m->flex_nodeadr[f] + m->flex_nodenum[f]; int nb = 0; - for (int i = nstart; i < nend; i++) { - mjtNum w = mju_evalBasis(coord, i-nstart, m->flex_interp[f]); + for (int j = 0; j < npc; j++) { + mjtNum w = mju_evalBasis(local, j, order); if (w < 1e-5) { continue; } if (bweight) bweight[nb] = w; - body[nb++] = m->flex_nodebodyid[i]; + body[nb++] = m->flex_nodebodyid[nstart + nodeindices[j]]; } return nb; @@ -871,6 +881,11 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { break; } + int npc = (order+1)*(order+1)*(order+1); + int cx = m->flex_cellnum[3*f+0]; + int cy = m->flex_cellnum[3*f+1]; + int cz = m->flex_cellnum[3*f+2]; + // allocate stack for node positions and Jacobians mj_markStack(d); mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum); @@ -910,152 +925,164 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { } } - // loop over Gauss points - // get reference positions from m->flex_node0 (Cartesian positions at qpos0) + // reference positions for all nodes int nstart = m->flex_nodeadr[f]; mjtNum* refpos = mjSTACKALLOC(d, 3*nodenum, mjtNum); for (int n = 0; n < nodenum; n++) { mju_copy3(refpos + 3*n, m->flex_node0 + 3*(n + nstart)); } - // B-bar: precompute center-point values for volumetric constraint (trilinear only) - if (order == 1) { - mjtNum center[3] = {0.5, 0.5, 0.5}; - mjtNum Fcur_c[9], Fref_c[9], Fref_inv_center[9], F_center[9]; - - // compute deformation gradient at center - mju_defGradient(Fcur_c, center, xpos, order); - mju_defGradient(Fref_c, center, refpos, order); - mat3_inverse(Fref_c, Fref_inv_center); - mju_mulMatMat3(F_center, Fcur_c, Fref_inv_center); - - // compute C and E at center - mjtNum C_c[9], E_c[9]; - mju_mulMatTMat3(C_c, F_center, F_center); - mju_scl(E_c, C_c, 0.5, 9); - E_c[0] -= 0.5; - E_c[4] -= 0.5; - E_c[8] -= 0.5; - - // J = det(F) at center - mjtNum I1_center = E_c[0] + E_c[4] + E_c[8]; - mjtNum J_center = mat3_det(F_center); - - // compute shape function gradients at center (8 nodes for trilinear) - mjtNum grad_center[8][3]; - shape_gradients(order, center, grad_center); - - // add I1 and J-1 constraints at center (reduced integration for volumetric) - mjtNum* dSdx = mjSTACKALLOC(d, 3*nodenum, mjtNum); - for (int inv = 0; inv < 2; inv++) { - if (inv == 0) { - // I1 = tr(E), dI1/dE = I - cpos[0] = I1_center; - } else { - // J - 1 = det(F) - 1, dJ/dF = cofactor(F) - cpos[0] = J_center - 1.0; - } - - volumetric_dSdx(inv, nodenum, grad_center, F_center, Fref_inv_center, dSdx); - strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac); - - if (issparse) { - mjtNum* sparse_jac = mjSTACKALLOC(d, combined_nnz, mjtNum); - for (int k = 0; k < combined_nnz; k++) { - sparse_jac[k] = strain_jac[combined_chain[k]]; - } - mj_addConstraint(m, d, sparse_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, - combined_nnz, combined_chain); - } else { - mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL); - } - } - } - - // add I1 and J-1 constraints at center (reduced integration for volumetric) + // per-cell arrays + mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* refpos_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* dSdx_local = mjSTACKALLOC(d, 3*npc, mjtNum); mjtNum* dSdx = mjSTACKALLOC(d, 3*nodenum, mjtNum); - for (int g = 0; g < ngauss; g++) { - mjtNum* p = gauss[g]; + int gindices[125]; // max npc = 125 for quadratic - // F = Fcur * Fref_inv - mjtNum Fcur[9], Fref[9], Fref_inv[9], F[9]; - mju_defGradient(Fcur, p, xpos, order); - mju_defGradient(Fref, p, refpos, order); - mat3_inverse(Fref, Fref_inv); - mju_mulMatMat3(F, Fcur, Fref_inv); + // loop over cells + for (int ci = 0; ci < cx; ci++) { + for (int cj = 0; cj < cy; cj++) { + for (int ck = 0; ck < cz; ck++) { + // gather cell-local node positions + mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, refpos, xpos_c, NULL, + refpos_c, gindices, NULL); - // compute Green-Lagrange strain E = 0.5*(C - I) - mjtNum C[9], E[9]; - mju_mulMatTMat3(C, F, F); - for (int j = 0; j < 9; j++) { - E[j] = 0.5 * C[j]; - } - E[0] -= 0.5; - E[4] -= 0.5; - E[8] -= 0.5; + // B-bar: center-point volumetric constraints (trilinear) + if (order == 1) { + mjtNum center[3] = {0.5, 0.5, 0.5}; + mjtNum Fcur_c[9], Fref_c[9], Fref_inv_c[9], F_c[9]; - // compute 3 invariants of E - mjtNum I1 = E[0] + E[4] + E[8]; - mjtNum trE2 = E[0]*E[0] + E[1]*E[3] + E[2]*E[6] + - E[3]*E[1] + E[4]*E[4] + E[5]*E[7] + - E[6]*E[2] + E[7]*E[5] + E[8]*E[8]; - mjtNum I2 = 0.5 * (I1*I1 - trE2); - mjtNum I3 = mat3_det(E); + mju_defGradient(Fcur_c, center, xpos_c, order); + mju_defGradient(Fref_c, center, refpos_c, order); + mat3_inverse(Fref_c, Fref_inv_c); + mju_mulMatMat3(F_c, Fcur_c, Fref_inv_c); - // compute shape function gradients at this Gauss point - mjtNum grad[27][3]; - shape_gradients(order, p, grad); + mjtNum C_c[9], E_c[9]; + mju_mulMatTMat3(C_c, F_c, F_c); + mju_scl(E_c, C_c, 0.5, 9); + E_c[0] -= 0.5; E_c[4] -= 0.5; E_c[8] -= 0.5; - // trilinear: 3 constraints per Gauss point (I1, I2, I3 skipped - only shear) - // quadratic: 6 constraints per Gauss point - for (int s = 0; s < 6; s++) { - // skip I1, I2, I3 for trilinear (I1, J-1 at center; I2 is small for small strain) - if (order == 1 && (s == 0 || s == 1 || s == 2)) { - continue; - } + mjtNum I1_c = E_c[0] + E_c[4] + E_c[8]; + mjtNum J_c = mat3_det(F_c); - mjtNum dSdE[9]; - mju_zero(dSdE, 9); + mjtNum grad_c[8][3]; + shape_gradients(order, center, grad_c); - if (s == 0) { - // I1 = tr(E), dI1/dE = I (only for quadratic) - cpos[0] = I1; - dSdE[0] = dSdE[4] = dSdE[8] = 1.0; - } else if (s == 1) { - // I2 = 0.5*(tr(E)^2 - tr(E^2)), dI2/dE = tr(E)*I - E - cpos[0] = I2; - dSdE[0] = I1 - E[0]; - dSdE[4] = I1 - E[4]; - dSdE[8] = I1 - E[8]; - dSdE[1] = -E[1]; dSdE[3] = -E[3]; - dSdE[2] = -E[2]; dSdE[6] = -E[6]; - dSdE[5] = -E[5]; dSdE[7] = -E[7]; - } else if (s == 2) { - // I3 = det(E), dI3/dE = cofactor(E) - cpos[0] = I3; - mat3_cofactor(E, dSdE); - } else { - // off-diagonal entries: s=3->E12, s=4->E13, s=5->E23 - int offdiag_idx[3] = {1, 2, 5}; - int ij = offdiag_idx[s - 3]; - cpos[0] = E[ij]; - dSdE[ij] = 1.0; - } + for (int inv = 0; inv < 2; inv++) { + cpos[0] = (inv == 0) ? I1_c : J_c - 1.0; - // compute dS/dx for all nodes - invariant_dSdx(nodenum, grad, F, Fref_inv, dSdE, dSdx); - strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac); + // compute local dSdx + volumetric_dSdx(inv, npc, grad_c, F_c, Fref_inv_c, dSdx_local); - // add constraint - if (issparse) { - mjtNum* sparse_jac = mjSTACKALLOC(d, combined_nnz, mjtNum); - for (int k = 0; k < combined_nnz; k++) { - sparse_jac[k] = strain_jac[combined_chain[k]]; + // scatter to global dSdx + mju_zero(dSdx, 3*nodenum); + for (int n = 0; n < npc; n++) { + mju_addTo3(dSdx + 3*gindices[n], dSdx_local + 3*n); + } + + strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac); + + if (issparse) { + mj_markStack(d); + mjtNum* sj = mjSTACKALLOC(d, combined_nnz, mjtNum); + for (int k = 0; k < combined_nnz; k++) { + sj[k] = strain_jac[combined_chain[k]]; + } + mj_addConstraint(m, d, sj, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, + combined_nnz, combined_chain); + mj_freeStack(d); + } else { + mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL); + } + } + } + + // Gauss integration per cell + for (int g = 0; g < ngauss; g++) { + mjtNum* p = gauss[g]; + + // F = Fcur * Fref_inv + mjtNum Fcur[9], Fref[9], Fref_inv[9], F[9]; + mju_defGradient(Fcur, p, xpos_c, order); + mju_defGradient(Fref, p, refpos_c, order); + mat3_inverse(Fref, Fref_inv); + mju_mulMatMat3(F, Fcur, Fref_inv); + + // Green-Lagrange strain E = 0.5*(C - I) + mjtNum C[9], E[9]; + mju_mulMatTMat3(C, F, F); + for (int j = 0; j < 9; j++) { + E[j] = 0.5 * C[j]; + } + E[0] -= 0.5; E[4] -= 0.5; E[8] -= 0.5; + + // 3 invariants of E + mjtNum I1 = E[0] + E[4] + E[8]; + mjtNum trE2 = E[0]*E[0] + E[1]*E[3] + E[2]*E[6] + + E[3]*E[1] + E[4]*E[4] + E[5]*E[7] + + E[6]*E[2] + E[7]*E[5] + E[8]*E[8]; + mjtNum I2 = 0.5 * (I1*I1 - trE2); + mjtNum I3 = mat3_det(E); + + // shape function gradients at Gauss point + mjtNum grad[27][3]; + shape_gradients(order, p, grad); + + for (int s = 0; s < 6; s++) { + // skip I1,I2,I3 for trilinear (B-bar handles vol) + if (order == 1 && (s == 0 || s == 1 || s == 2)) { + continue; + } + + mjtNum dSdE[9]; + mju_zero(dSdE, 9); + + if (s == 0) { + cpos[0] = I1; + dSdE[0] = dSdE[4] = dSdE[8] = 1.0; + } else if (s == 1) { + cpos[0] = I2; + dSdE[0] = I1-E[0]; dSdE[4] = I1-E[4]; + dSdE[8] = I1-E[8]; + dSdE[1] = -E[1]; dSdE[3] = -E[3]; + dSdE[2] = -E[2]; dSdE[6] = -E[6]; + dSdE[5] = -E[5]; dSdE[7] = -E[7]; + } else if (s == 2) { + cpos[0] = I3; + mat3_cofactor(E, dSdE); + } else { + int offdiag_idx[3] = {1, 2, 5}; + int ij = offdiag_idx[s - 3]; + cpos[0] = E[ij]; + dSdE[ij] = 1.0; + } + + // compute local dS/dx for cell nodes + invariant_dSdx(npc, grad, F, Fref_inv, dSdE, + dSdx_local); + + // scatter to global dSdx + mju_zero(dSdx, 3*nodenum); + for (int n = 0; n < npc; n++) { + mju_addTo3(dSdx + 3*gindices[n], dSdx_local + 3*n); + } + + strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac); + + if (issparse) { + mj_markStack(d); + mjtNum* sj = mjSTACKALLOC(d, combined_nnz, mjtNum); + for (int k = 0; k < combined_nnz; k++) { + sj[k] = strain_jac[combined_chain[k]]; + } + mj_addConstraint(m, d, sj, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, + combined_nnz, combined_chain); + mj_freeStack(d); + } else { + mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL); + } + } } - mj_addConstraint(m, d, sparse_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, - combined_nnz, combined_chain); - } else { - mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL); } } } @@ -1899,10 +1926,13 @@ void mj_diagApprox(const mjModel* m, mjData* d) { int nstart = m->flex_nodeadr[flex_id]; int order = m->flex_interp[flex_id]; - // compute constraint count: trilinear (2 + 3*8 = 26), quadratic (6*27 = 162) + // compute constraint count per cell, then multiply by ncells int nquad = order + 1; int ngauss = nquad * nquad * nquad; - int nconstraint = (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss); + int ncells = m->flex_cellnum[3*flex_id+0] + * m->flex_cellnum[3*flex_id+1] + * m->flex_cellnum[3*flex_id+2]; + int nconstraint = ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss)); mjtNum avg_invweight = 0; for (int n = 0; n < nodenum; n++) { @@ -2510,7 +2540,10 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { } int nquad = order + 1; // 2 for order=1, 3 for order=2 int ngauss = nquad * nquad * nquad; // 8 or 27 - size = (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss); // 26 or 162 + int ncells = m->flex_cellnum[3*id[0]+0] + * m->flex_cellnum[3*id[0]+1] + * m->flex_cellnum[3*id[0]+2]; + size = ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss)); if (nnz) { // Count unique DOFs across all node bodies (matching instantiation) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 2ae4d5b4..8f3412f1 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -580,9 +580,11 @@ void mj_flex(const mjModel* m, mjData* d) { } } - // trilinear interpolation + // trilinear/quadratic interpolation else { - mjtNum nodexpos[3*mjMAXFLEXNODES]; + int nodenum = nend - nstart; + mj_markStack(d); + mjtNum* nodexpos = mjSTACKALLOC(d, 3*nodenum, mjtNum); if (m->flex_centered[f]) { for (int i=nstart; i < nend; i++) { mji_copy3(nodexpos + 3*(i-nstart), d->xpos + 3*m->flex_nodebodyid[i]); @@ -596,14 +598,26 @@ void mj_flex(const mjModel* m, mjData* d) { } int order = m->flex_interp[f]; - if (nend - nstart != (order + 1) * (order + 1) * (order + 1)) { + int cx = m->flex_cellnum[3*f+0]; + int cy = m->flex_cellnum[3*f+1]; + int cz = m->flex_cellnum[3*f+2]; + int nx_g = cx * order + 1; + int ny_g = cy * order + 1; + int nz_g = cz * order + 1; + if (nend - nstart != nx_g * ny_g * nz_g) { mjERROR("flex_interp_order mismatch"); } for (int i=vstart; i < vend; i++) { mju_zero3(d->flexvert_xpos+3*i); - mju_interpolate3D(d->flexvert_xpos+3*i, m->flex_vert0 + 3*i, nodexpos, order); + + // cell lookup: get local coords and node indices + mjtNum local[3]; + int nodeindices[27]; // max npc for quadratic: 3^3 = 27 + mju_cellLookup(m->flex_vert0 + 3*i, m->flex_cellnum+3*f, order, local, nodeindices); + mju_interpolate3D(d->flexvert_xpos+3*i, local, nodexpos, order, nodeindices); } + mj_freeStack(d); } } @@ -2617,7 +2631,10 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { if (order && nodenum) { int nquad = order + 1; int ngauss = nquad * nquad * nquad; - i += (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss); + int ncells = m->flex_cellnum[3*k+0] + * m->flex_cellnum[3*k+1] + * m->flex_cellnum[3*k+2]; + i += ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss)); } break; } diff --git a/src/engine/engine_core_util.c b/src/engine/engine_core_util.c index c49f3ee5..bf2b6994 100644 --- a/src/engine/engine_core_util.c +++ b/src/engine/engine_core_util.c @@ -987,6 +987,34 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], //-------------------------- miscellaneous utilities ----------------------------------------------- +// gather global node positions and velocities +void mju_flexGatherState(const mjModel* m, mjData* d, int f, mjtNum* xpos, mjtNum* vel) { + int nodenum = m->flex_nodenum[f]; + int nstart = m->flex_nodeadr[f]; + int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f]; + + // compute positions + if (m->flex_centered[f]) { + for (int i=0; i < nodenum; i++) { + mju_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]); + if (vel) { + mju_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]); + } + } + } else { + mjtNum screw[6]; + for (int i=0; i < nodenum; i++) { + mju_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart)); + mju_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]); + if (vel) { + mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0); + mju_copy3(vel + 3*i, screw + 3); + } + } + } +} + + // extract 6D force:torque for one contact, in contact frame void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) { mjContact* con; diff --git a/src/engine/engine_core_util.h b/src/engine/engine_core_util.h index 39633ff0..949aa257 100644 --- a/src/engine/engine_core_util.h +++ b/src/engine/engine_core_util.h @@ -129,6 +129,9 @@ MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], //-------------------------- miscellaneous --------------------------------------------------------- +// gather global node positions and velocities +MJAPI void mju_flexGatherState(const mjModel* m, mjData* d, int f, mjtNum* xpos, mjtNum* vel); + // extract 6D force:torque for one contact, in contact frame MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]); diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 5f1058d0..1eb4df9d 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -915,127 +915,146 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op, continue; } + int order = m->flex_interp[f]; + int npc = (order+1)*(order+1)*(order+1); + int cx = m->flex_cellnum[3*f+0]; + int cy = m->flex_cellnum[3*f+1]; + int cz = m->flex_cellnum[3*f+2]; + int nodenum = m->flex_nodenum[f]; int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f]; // standard stack allocation mj_markStack(d); mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum); - mjtNum* K_rot = mjSTACKALLOC(d, 9*nodenum*nodenum, mjtNum); - // sparse Jacobian allocations - int dim = 3 * nodenum; - int* rownnz = mjSTACKALLOC(d, dim, int); - int* rowadr = mjSTACKALLOC(d, dim, int); - mjtNum* J_val = mjSTACKALLOC(d, dim*nv, mjtNum); - int* J_colind = mjSTACKALLOC(d, dim*nv, int); + // per-cell arrays + int dim_c = 3 * npc; + mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* K_rot_cell = mjSTACKALLOC(d, dim_c*dim_c, mjtNum); + + // sparse Jacobian for one cell + int* J_rownnz = mjSTACKALLOC(d, dim_c, int); + int* J_rowadr = mjSTACKALLOC(d, dim_c, int); + mjtNum* J_val = mjSTACKALLOC(d, dim_c*nv, mjtNum); + int* J_colind = mjSTACKALLOC(d, dim_c*nv, int); // temp allocations for chain int* chain_colind = mjSTACKALLOC(d, nv, int); mjtNum* blk_jac = mjSTACKALLOC(d, 3*nv, mjtNum); - // compute positions, rotation and Jacobian - mjtNum quat[4] = {1, 0, 0, 0}; - mj_flexInterpState(m, d, f, xpos, NULL, quat); + // gather raw node positions (unrotated) + mju_flexGatherState(m, d, f, xpos, NULL); - // compute generalized stiffness in global frame: K_rot = R * K * R^T - mjtNum R[9]; - mju_quat2Mat(R, quat); // R = R_global2local - mjtNum RT[9]; - mju_transpose(RT, R, 3, 3); // RT = R_local2global + // loop over cells + int cell_idx = 0; + for (int ci = 0; ci < cx; ci++) { + for (int cj = 0; cj < cy; cj++) { + for (int ck = 0; ck < cz; ck++) { + // gather cell-local node positions + int gindices[125]; // max npc = 125 for quadratic + mjtNum quat[4]; + mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, NULL, + xpos_c, NULL, NULL, gindices, quat); - // blockwise rotation: K_rot(i,j) = scale * RT * K_local(i,j) * R - // note: k stores -K, so K_rot = scale * (-K_phys) - for (int i=0; i < nodenum; i++) { - for (int j=0; j < nodenum; j++) { - mjtNum blk[9], tmp[9]; + // R = R_global2local, RT = R_local2global + mjtNum R[9], RT[9]; + mju_quat2Mat(R, quat); + mju_transpose(RT, R, 3, 3); - // get K_local(i,j) - int adr = (3*i)*(3*nodenum) + 3*j; - for (int r=0; r < 3; r++) { - for (int c=0; c < 3; c++) { - blk[3*r+c] = k[adr + r*(3*nodenum) + c]; + // get cell stiffness + mjtNum* k_cell = k + cell_idx * 3*npc * 3*npc; + + // compute K_rot_cell = RT * K_cell * R (block-wise) + mju_zero(K_rot_cell, dim_c*dim_c); + for (int a = 0; a < npc; a++) { + for (int b = 0; b < npc; b++) { + mjtNum blk[9], tmp[9]; + + // get K_cell(a,b) 3x3 block + int adr_cell = (3*a)*(3*npc) + 3*b; + for (int r = 0; r < 3; r++) { + for (int c = 0; c < 3; c++) { + blk[3*r+c] = k_cell[adr_cell + r*(3*npc) + c]; + } + } + + // tmp = K * R + mju_mulMatMat3(tmp, blk, R); + // blk = RT * tmp = RT * K * R + mju_mulMatMat3(blk, RT, tmp); + + // store in K_rot_cell at (a, b) + int adr_out = (3*a)*dim_c + 3*b; + for (int r = 0; r < 3; r++) { + for (int c = 0; c < 3; c++) { + K_rot_cell[adr_out + r*dim_c + c] = scale * blk[3*r+c]; + } + } + } } - } - // tmp = K * R - mju_mulMatMat3(tmp, blk, R); + // construct sparse Jacobian for this cell's nodes + int current_adr = 0; + for (int n = 0; n < npc; n++) { + int bid = bodyid[gindices[n]]; + int chain_nnz = mj_bodyChain(m, bid, chain_colind); + mj_jacSparse(m, d, blk_jac, NULL, xpos+3*gindices[n], bid, + chain_nnz, chain_colind, /*flg_skipcommon=*/0); - // blk = RT * tmp = RT * K * R - mju_mulMatMat3(blk, RT, tmp); + for (int r = 0; r < 3; r++) { + int row_idx = 3*n + r; + J_rownnz[row_idx] = chain_nnz; + J_rowadr[row_idx] = current_adr; - // store scaled into K_rot - for (int r=0; r < 3; r++) { - for (int c=0; c < 3; c++) { - K_rot[adr + r*(3*nodenum) + c] = scale * blk[3*r+c]; - } - } - } - } - - // construct sparse Jacobian J_val - int current_adr = 0; - for (int i=0; i < nodenum; i++) { - // get chain for this node - int chain_nnz = mj_bodyChain(m, bodyid[i], chain_colind); - - // compute sparse Jacobian for this node (3 rows) - mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind, - /*flg_skipcommon=*/0); - - // copy to sparse structure - for (int r=0; r<3; r++) { - int row_idx = 3*i + r; - rownnz[row_idx] = chain_nnz; - rowadr[row_idx] = current_adr; - - for (int idx=0; idx= 0) { - J_reduced[i*ndof + local_idx] = J_val[adr + idx]; } + } } - } - // H -= J_reduced^T * K_rot * J_reduced - // K_rot * J_reduced (dim x ndof) - mjtNum* KJ = mjSTACKALLOC(d, dim*ndof, mjtNum); - mju_mulMatMat(KJ, K_rot, J_reduced, dim, dim, ndof); + // apply operation with cell's K_rot and J + if (op == mjFLEXOP_VEC) { + addJTBJ_mulSparse(m, d, res, vec, J_rownnz, J_rowadr, J_colind, + J_val, K_rot_cell, dim_c); + } else if (op == mjFLEXOP_ADDH) { + mj_markStack(d); + // H -= J_cell^T * K_rot_cell * J_cell + mjtNum* J_reduced = mjSTACKALLOC(d, dim_c*ndof, mjtNum); + mju_zero(J_reduced, dim_c*ndof); - // H[i, j] -= sum_k J_reduced[k, i] * KJ[k, j] - for (int i=0; i= 0) { + J_reduced[i*ndof + local_idx] = J_val[adr + idx]; + } + } + } + + // KJ = K_rot_cell * J_reduced (dim_c x ndof) + mjtNum* KJ = mjSTACKALLOC(d, dim_c*ndof, mjtNum); + mju_mulMatMat(KJ, K_rot_cell, J_reduced, dim_c, dim_c, ndof); + + // H[i,j] -= J_reduced[k,i] * KJ[k,j] + for (int i = 0; i < ndof; i++) { + for (int j = 0; j < ndof; j++) { + mjtNum val = 0; + for (int dim_idx = 0; dim_idx < dim_c; dim_idx++) { + val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j]; + } + res[i*ndof + j] -= val; + } + } + mj_freeStack(d); } - // res is H - res[i*ndof + j] -= val; + + cell_idx++; } } } diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index d67f8683..e6d3b6fa 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -58,61 +58,7 @@ static void inline GradSquaredLengths(mjtNum gradient[6][2][3], } } -// compute interpolated flex state: xpos, vel, quat -// f: flex index -// xpos: (output) 3*nodenum -// vel: (output) 3*nodenum, can be NULL -// quat: (output) 4, rotation from global to local -void mj_flexInterpState(const mjModel* m, mjData* d, int f, - mjtNum* xpos, mjtNum* vel, mjtNum* quat) { - int nodenum = m->flex_nodenum[f]; - int nstart = m->flex_nodeadr[f]; - int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f]; - mjtNum com[3] = {0}; - // compute positions - if (m->flex_centered[f]) { - for (int i=0; i < nodenum; i++) { - mji_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]); - if (vel) { - mji_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]); - } - } - } else { - mjtNum screw[6]; - for (int i=0; i < nodenum; i++) { - mji_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart)); - mji_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]); - if (vel) { - mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0); - mji_copy3(vel + 3*i, screw + 3); - } - } - } - - // compute center of mass - for (int i = 0; i < nodenum; i++) { - mji_addToScl3(com, xpos+3*i, 1.0/nodenum); - } - - // compute the Jacobian at the center of mass - mjtNum mat[9] = {0}; - mjtNum p[3] = {.5, .5, .5}; - mju_defGradient(mat, p, xpos, m->flex_interp[f]); - - // find rotation - mju_mat2Rot(quat, mat); - mju_negQuat(quat, quat); - - // rotate vertices to quat and add reference center of mass - for (int i = 0; i < nodenum; i++) { - mju_rotVecQuat(xpos+3*i, xpos+3*i, quat); - mji_addTo3(xpos+3*i, p); - if (vel) { - mju_rotVecQuat(vel+3*i, vel+3*i, quat); - } - } -} // spring and damper forces static void mj_springdamper(const mjModel* m, mjData* d) { @@ -284,42 +230,111 @@ static void mj_springdamper(const mjModel* m, mjData* d) { } if (m->flex_interp[f]) { + int order = m->flex_interp[f]; + int npc = (order+1)*(order+1)*(order+1); // nodes per cell + int cx = m->flex_cellnum[3*f+0]; + int cy = m->flex_cellnum[3*f+1]; + int cz = m->flex_cellnum[3*f+2]; + int ny_g = cy * order + 1; + int nz_g = cz * order + 1; + mj_markStack(d); - mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum); - mjtNum* displ = mjSTACKALLOC(d, 3*nodenum, mjtNum); - mjtNum* vel = mjSTACKALLOC(d, 3*nodenum, mjtNum); - mjtNum* frc = mjSTACKALLOC(d, 3*nodenum, mjtNum); - mjtNum* dmp = mjSTACKALLOC(d, 3*nodenum, mjtNum); + + // allocate global arrays + mjtNum* xpos_g = mjSTACKALLOC(d, 3*nodenum, mjtNum); + mjtNum* vel_g = mjSTACKALLOC(d, 3*nodenum, mjtNum); + mjtNum* frc_g = mjSTACKALLOC(d, 3*nodenum, mjtNum); + mjtNum* dmp_g = mjSTACKALLOC(d, 3*nodenum, mjtNum); mjtNum* xpos0 = m->flex_node0 + 3*m->flex_nodeadr[f]; int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f]; - mjtNum quat[4] = {1, 0, 0, 0}; - mj_flexInterpState(m, d, f, xpos, vel, quat); + // gather global node positions and velocities (unrotated) + mju_flexGatherState(m, d, f, xpos_g, vel_g); - // compute displacement - for (int i = 0; i < nodenum; i++) { - mji_addScl3(displ+3*i, xpos+3*i, xpos0+3*i, -1); + // zero global force accumulators + mju_zero(frc_g, 3*nodenum); + mju_zero(dmp_g, 3*nodenum); + + // per-cell arrays + mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* vel_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* xpos0_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* displ_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* frc_c = mjSTACKALLOC(d, 3*npc, mjtNum); + mjtNum* dmp_c = mjSTACKALLOC(d, 3*npc, mjtNum); + + // loop over cells + int cell_idx = 0; + for (int ci = 0; ci < cx; ci++) { + for (int cj = 0; cj < cy; cj++) { + for (int ck = 0; ck < cz; ck++) { + // gather cell-local node data + mjtNum quat[4]; + mjtNum p[3] = {.5, .5, .5}; + mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos_g, vel_g, xpos0, + xpos_c, vel_c, xpos0_c, NULL, quat); + + // rotate to corotational frame + for (int n = 0; n < npc; n++) { + mju_rotVecQuat(xpos_c+3*n, xpos_c+3*n, quat); + mji_addTo3(xpos_c+3*n, p); + mju_rotVecQuat(vel_c+3*n, vel_c+3*n, quat); + } + + // compute displacement + for (int n = 0; n < npc; n++) { + mji_addScl3(displ_c+3*n, xpos_c+3*n, xpos0_c+3*n, -1); + } + + // get cell stiffness matrix + mjtNum* k_cell = k + cell_idx * 3*npc * 3*npc; + + // compute force in corotational frame + if (enbl_spring) { + mju_mulMatVec(frc_c, k_cell, displ_c, 3*npc, 3*npc); + } + if (enbl_damper) { + mju_mulMatVec(dmp_c, k_cell, vel_c, 3*npc, 3*npc); + } + + // rotate back to global frame and scatter + mju_negQuat(quat, quat); + int local = 0; + for (int li = 0; li <= order; li++) { + for (int lj = 0; lj <= order; lj++) { + for (int lk = 0; lk <= order; lk++) { + int gi = ci*order + li; + int gj = cj*order + lj; + int gk = ck*order + lk; + int gidx = gi*ny_g*nz_g + gj*nz_g + gk; + mjtNum qfrc[3], qdmp[3]; + mji_rotVecQuat(qfrc, frc_c+3*local, quat); + mji_rotVecQuat(qdmp, dmp_c+3*local, quat); + if (enbl_spring) { + mji_addTo3(frc_g + 3*gidx, qfrc); + } + if (enbl_damper) { + mji_addTo3(dmp_g + 3*gidx, qdmp); + } + local++; + } + } + } + + cell_idx++; + } + } } - // compute force in the stretch frame - if (enbl_spring) mju_mulMatVec(frc, k, displ, 3*nodenum, 3*nodenum); - - // compute damping force in stretch frame - if (enbl_damper) mju_mulMatVec(dmp, k, vel, 3*nodenum, 3*nodenum); - - // rotate forces to global frame and add to qfrc - mju_negQuat(quat, quat); + // apply accumulated forces to bodies for (int i = 0; i < nodenum; i++) { - mjtNum qfrc[3], qdmp[3]; - mji_rotVecQuat(qfrc, frc+3*i, quat); - mji_rotVecQuat(qdmp, dmp+3*i, quat); - mju_scl3(qdmp, qdmp, m->flex_damping[f]); + mju_scl3(dmp_g+3*i, dmp_g+3*i, m->flex_damping[f]); if (m->flex_centered[f]) { - if (enbl_spring) mji_addTo3(d->qfrc_spring+m->body_dofadr[bodyid[i]], qfrc); - if (enbl_damper) mji_addTo3(d->qfrc_damper+m->body_dofadr[bodyid[i]], qdmp); + if (enbl_spring) mji_addTo3(d->qfrc_spring + m->body_dofadr[bodyid[i]], frc_g+3*i); + if (enbl_damper) mji_addTo3(d->qfrc_damper + m->body_dofadr[bodyid[i]], dmp_g+3*i); } else { - if (enbl_spring) mj_applyFT(m, d, qfrc, 0, xpos+3*i, bodyid[i], d->qfrc_spring); - if (enbl_damper) mj_applyFT(m, d, qdmp, 0, xpos+3*i, bodyid[i], d->qfrc_damper); + if (enbl_spring) mj_applyFT(m, d, frc_g+3*i, 0, xpos_g+3*i, bodyid[i], d->qfrc_spring); + if (enbl_damper) mj_applyFT(m, d, dmp_g+3*i, 0, xpos_g+3*i, bodyid[i], d->qfrc_damper); } } diff --git a/src/engine/engine_passive.h b/src/engine/engine_passive.h index ba8b3233..096cfefb 100644 --- a/src/engine/engine_passive.h +++ b/src/engine/engine_passive.h @@ -28,9 +28,7 @@ extern "C" { // all passive forces MJAPI void mj_passive(const mjModel* m, mjData* d); -// compute interpolated flex state: xpos, vel, quat -MJAPI void mj_flexInterpState(const mjModel* m, mjData* d, int f, - mjtNum* xpos, mjtNum* vel, mjtNum* quat); + //------------------------- fluid models ----------------------------------------------------------- diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index abaf7055..6ec01f72 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -612,11 +612,116 @@ mjtNum mju_evalBasis(const mjtNum x[3], int i, int order) { } } +// map global parametric coord to cell-local coord and build node indices +// coord: [0,1]^3 parametric coordinates +// cellnum: cell counts (cx, cy, cz) +// order: interpolation order (1=trilinear, 2=triquadratic) +// local: output local parametric coordinates within cell [0,1]^3 +// nodeindices: output array of global node indices for the cell (size (order+1)^3, may be NULL) +// returns: number of nodes per cell (order+1)^3 +int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNum local[3], + int* nodeindices) { + int cx = cellnum[0], cy = cellnum[1], cz = cellnum[2]; + + // find containing cell + int ci = (int)mju_floor(coord[0] * cx); + int cj = (int)mju_floor(coord[1] * cy); + int ck = (int)mju_floor(coord[2] * cz); + ci = mjMIN(ci, cx - 1); ci = mjMAX(ci, 0); + cj = mjMIN(cj, cy - 1); cj = mjMAX(cj, 0); + ck = mjMIN(ck, cz - 1); ck = mjMAX(ck, 0); + + // local parametric coordinates within cell + local[0] = mju_clip(coord[0] * cx - ci, 0, 1); + local[1] = mju_clip(coord[1] * cy - cj, 0, 1); + local[2] = mju_clip(coord[2] * cz - ck, 0, 1); + + // build node indices for this cell + if (nodeindices) { + int ny_g = cy * order + 1; + int nz_g = cz * order + 1; + int ni = 0; + for (int li = 0; li <= order; li++) { + for (int lj = 0; lj <= order; lj++) { + for (int lk = 0; lk <= order; lk++) { + int gi = ci*order + li; + int gj = cj*order + lj; + int gk = ck*order + lk; + nodeindices[ni++] = gi*ny_g*nz_g + gj*nz_g + gk; + } + } + } + } + + int npc = (order + 1) * (order + 1) * (order + 1); + return npc; +} + + // interpolate a function at x with given interpolation coefficients and order n -void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order) { +void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order, + const int* nodeindices) { int npoint = (order + 1) * (order + 1) * (order + 1); for (int j=0; j < npoint; j++) { - mju_addToScl3(res, coeff+3*j, mju_evalBasis(x, j, order)); + int idx = nodeindices ? nodeindices[j] : j; + mju_addToScl3(res, coeff+3*idx, mju_evalBasis(x, j, order)); + } +} + + +static void flexInterpRotation(int order, const mjtNum* xpos_c, + const mjtNum local[3], mjtNum* quat) { + mjtNum mat[9] = {0}; + + if (order > 0) { + mju_defGradient(mat, local, xpos_c, order); + } else { + // order 0: fallback to identity matrix + mat[0] = 1; + mat[4] = 1; + mat[8] = 1; + } + + // find rotation + quat[0] = 1; + quat[1] = 0; + quat[2] = 0; + quat[3] = 0; + mju_mat2Rot(quat, mat); + mju_negQuat(quat, quat); +} + + +// gather cell-local quantities and optionally compute rotation +void mju_flexGatherCellState(int order, int cy, int cz, int ci, int cj, int ck, + const mjtNum* xpos_g, const mjtNum* vel_g, const mjtNum* xpos0_g, + mjtNum* xpos_c, mjtNum* vel_c, mjtNum* xpos0_c, + int* nodeindices, mjtNum* quat) { + int ny_g = cy * order + 1; + int nz_g = cz * order + 1; + + int local = 0; + for (int li = 0; li <= order; li++) { + for (int lj = 0; lj <= order; lj++) { + for (int lk = 0; lk <= order; lk++) { + int gi = ci*order + li; + int gj = cj*order + lj; + int gk = ck*order + lk; + int gidx = gi*ny_g*nz_g + gj*nz_g + gk; + + if (xpos_c && xpos_g) mju_copy3(xpos_c + 3*local, xpos_g + 3*gidx); + if (vel_c && vel_g) mju_copy3(vel_c + 3*local, vel_g + 3*gidx); + if (xpos0_c && xpos0_g) mju_copy3(xpos0_c + 3*local, xpos0_g + 3*gidx); + if (nodeindices) nodeindices[local] = gidx; + + local++; + } + } + } + + if (quat && xpos_c) { + mjtNum p[3] = {.5, .5, .5}; + flexInterpRotation(order, xpos_c, p, quat); } } diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h index cac5bed5..b5574670 100644 --- a/src/engine/engine_util_misc.h +++ b/src/engine/engine_util_misc.h @@ -89,8 +89,20 @@ MJAPI void mju_defGradient(mjtNum res[9], const mjtNum p[3], const mjtNum* dof, // evaluate the basis function at x for the i-th node MJAPI mjtNum mju_evalBasis(const mjtNum x[3], int i, int order); +// map global parametric coord to cell-local coord and build node indices +MJAPI int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNum local[3], + int* nodeindices); + // interpolate a function at x with given interpolation coefficients and order n -MJAPI void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order); +MJAPI void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order, + const int* nodeindices); + +// gather cell-local quantities and optionally compute rotation +MJAPI void mju_flexGatherCellState(int order, int cy, int cz, int ci, int cj, int ck, + const mjtNum* xpos_g, const mjtNum* vel_g, + const mjtNum* xpos0_g, mjtNum* xpos_c, mjtNum* vel_c, + mjtNum* xpos0_c, int* nodeindices, mjtNum* quat); + // ----------------------------- Base64 ------------------------------------------------------------ diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c index 1c7b602c..d1ce7ffe 100644 --- a/src/engine/engine_vis_interact.c +++ b/src/engine/engine_vis_interact.c @@ -863,24 +863,30 @@ int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt, flexdist = newdist; if (m->flex_interp[i]) { mjtNum* coord = m->flex_vert0 + 3*(m->flex_vertadr[i] + vertid); + int order = m->flex_interp[i]; + int npc = (order+1)*(order+1)*(order+1); + + // cell lookup: get local coords and node indices + mjtNum loc[3]; + int nodeindices[27]; // max npc for quadratic: 3^3 = 27 + mju_cellLookup(coord, m->flex_cellnum+3*i, order, loc, nodeindices); + + // find node with largest weight in this cell int nodeid = -1; int nstart = m->flex_nodeadr[i]; - int nend = nstart + m->flex_nodenum[i]; mjtNum w = 0; - for (int j = nstart; j < nend; j++) { - if (mju_evalBasis(coord, j-nstart, m->flex_interp[i]) > w) { - w = mju_evalBasis(coord, j-nstart, m->flex_interp[i]); - nodeid = j; + for (int j = 0; j < npc; j++) { + mjtNum ww = mju_evalBasis(loc, j, order); + if (ww > w) { + w = ww; + nodeid = nodeindices[j]; } } - if (nodeid < 0) { - mjERROR("flex %d: node closest to vertex %d not found", i, vertid); - } - flexbodyid = m->flex_nodebodyid[m->flex_nodeadr[i] + nodeid]; + flexbodyid = m->flex_nodebodyid[nstart + nodeid]; if (m->flex_centered[i]) { mju_copy3(flexpnt, d->xpos + 3*flexbodyid); } else { - mju_mulMatVec3(flexpnt, d->xmat + 9*flexbodyid, m->flex_node + 3*nodeid); + mju_mulMatVec3(flexpnt, d->xmat + 9*flexbodyid, m->flex_node + 3*(nstart + nodeid)); mju_addTo3(flexpnt, d->xpos + 3*flexbodyid); } } else { diff --git a/src/engine/engine_vis_visualize.c b/src/engine/engine_vis_visualize.c index b1b74d21..c0072dbb 100644 --- a/src/engine/engine_vis_visualize.c +++ b/src/engine/engine_vis_visualize.c @@ -1434,10 +1434,9 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, } // control points box - mjtNum xpos[mjMAXFLEXNODES]; + mjtNum* xpos = mjSTACKALLOC(d, 3*m->flex_nodenum[f], mjtNum); int nstart = m->flex_nodeadr[f]; int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f]; - int nnode = m->flex_interp[f]+1; if (m->flex_centered[f]) { for (int i=0; i < m->flex_nodenum[f]; i++) { mju_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]); @@ -1448,15 +1447,23 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, mju_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]); } } - for (int i=0; i < nnode; i++) { - for (int j=0; j < nnode; j++) { - for (int k=0; k < nnode; k++) { - int nn = nnode*nnode; - int offset = 3*(nn*(i+0) + nnode*(j+0) + k); - int offset1 = 3*(nn*(i+1) + nnode*(j+0) + k); - int offset2 = 3*(nn*(i+0) + nnode*(j+1) + k); - int offset3 = 3*(nn*(i+0) + nnode*(j+0) + (k+1)); - if (i < nnode-1) { + + int cx = m->flex_cellnum[3*f+0]; + int cy = m->flex_cellnum[3*f+1]; + int cz = m->flex_cellnum[3*f+2]; + int order = m->flex_interp[f]; + int NX = cx * order + 1; + int NY = cy * order + 1; + int NZ = cz * order + 1; + + for (int i=0; i < NX; i++) { + for (int j=0; j < NY; j++) { + for (int k=0; k < NZ; k++) { + int offset = 3*(i*NY*NZ + j*NZ + k); + int offset1 = 3*((i+1)*NY*NZ + j*NZ + k); + int offset2 = 3*(i*NY*NZ + (j+1)*NZ + k); + int offset3 = 3*(i*NY*NZ + j*NZ + (k+1)); + if (i < NX-1) { mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN); if (!thisgeom) { return; @@ -1465,7 +1472,7 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, mjv_connector(thisgeom, mjGEOM_LINE, 3, xpos+offset, xpos+offset1); releaseGeom(&thisgeom, scn); } - if (j < nnode-1) { + if (j < NY-1) { mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN); if (!thisgeom) { return; @@ -1474,7 +1481,7 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, mjv_connector(thisgeom, mjGEOM_LINE, 3, xpos+offset, xpos+offset2); releaseGeom(&thisgeom, scn); } - if (k < nnode-1) { + if (k < NZ-1) { mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN); if (!thisgeom) { return; diff --git a/src/user/user_flexcomp.cc b/src/user/user_flexcomp.cc index 56fcb615..4035a176 100644 --- a/src/user/user_flexcomp.cc +++ b/src/user/user_flexcomp.cc @@ -77,6 +77,7 @@ bool IsValidElementOrNodeHeader22(const std::string& line) { mjCFlexcomp::mjCFlexcomp(void) { type = mjFCOMPTYPE_GRID; count[0] = count[1] = count[2] = 10; + cellcount[0] = cellcount[1] = cellcount[2] = -1; mjuu_setvec(spacing, 0.02, 0.02, 0.02); mjuu_setvec(scale, 1, 1, 1); mass = 1; @@ -269,10 +270,19 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf // construct pinned array int nnode = 0; - if (doftype == mjFCOMPDOF_TRILINEAR) { - nnode = 8; - } else if (doftype == mjFCOMPDOF_QUADRATIC) { - nnode = 27; + if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) { + int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2; + // multi-cell count for mesh/direct/gmsh, else single cell + int cx = 1, cy = 1, cz = 1; + if (type == mjFCOMPTYPE_MESH || type == mjFCOMPTYPE_DIRECT || + type == mjFCOMPTYPE_GMSH) { + if (cellcount[0] >= 0) { + cx = cellcount[0]; + cy = cellcount[1]; + cz = cellcount[2]; + } + } + nnode = (cx*order+1) * (cy*order+1) * (cz*order+1); } pinned = vector(std::max(npnt, nnode), rigid); @@ -562,36 +572,81 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf } } - // create nodal mesh for trilinear interpolation + // create nodal mesh for trilinear/quadratic interpolation if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) { - int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2; - flex->SetOrder(order); - std::vector node(3*(order+1)*(order+1)*(order+1), 0); + flex->spec.order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2; + + if (cellcount[0] >= 0) { + flex->spec.cellcount[0] = cellcount[0]; + flex->spec.cellcount[1] = cellcount[1]; + flex->spec.cellcount[2] = cellcount[2]; + } + + // total number of nodes with shared boundaries + int nx = flex->spec.cellcount[0] * flex->spec.order + 1; + int ny = flex->spec.cellcount[1] * flex->spec.order + 1; + int nz = flex->spec.cellcount[2] * flex->spec.order + 1; + int nnode = nx * ny * nz; + + std::vector node(3 * nnode, 0); int idx = 0; - double step = 1.0 / (double)order; + + // Simpson's rule weights for quadratic mass distribution double massP2[3] = {1. / 6., 2. / 3., 1. / 6.}; - for (int i=0; i <= order; i++) { - for (int j=0; j <= order; j++) { - for (int k=0; k <= order; k++) { + + // compute per-node mass for trilinear: + // mass / nnode (uniform), or use Simpson for quadratic + double node_mass_uniform = mass / nnode; + + for (int gi = 0; gi < nx; gi++) { + for (int gj = 0; gj < ny; gj++) { + for (int gk = 0; gk < nz; gk++) { + // parametric position in [0, 1]^3 + double s = (double)gi / (flex->spec.cellcount[0] * flex->spec.order); + double t = (double)gj / (flex->spec.cellcount[1] * flex->spec.order); + double u = (double)gk / (flex->spec.cellcount[2] * flex->spec.order); + + // physical position + double px = minmax[0] + s * (minmax[3] - minmax[0]); + double py = minmax[1] + t * (minmax[4] - minmax[1]); + double pz = minmax[2] + u * (minmax[5] - minmax[2]); + if (pinned[idx]) { - node[3*idx+0] = minmax[0] + i * step * (minmax[3] - minmax[0]); - node[3*idx+1] = minmax[1] + j * step * (minmax[4] - minmax[1]); - node[3*idx+2] = minmax[2] + k * step * (minmax[5] - minmax[2]); - mjs_appendString(pf->nodebody, mjs_getName(body->element)->c_str()); + node[3*idx+0] = px; + node[3*idx+1] = py; + node[3*idx+2] = pz; + mjs_appendString(pf->nodebody, + mjs_getName(body->element)->c_str()); idx++; continue; } mjsBody* pb = mjs_addBody(body, 0); - pb->pos[0] = minmax[0] + i * step * (minmax[3] - minmax[0]); - pb->pos[1] = minmax[1] + j * step * (minmax[4] - minmax[1]); - pb->pos[2] = minmax[2] + k * step * (minmax[5] - minmax[2]); + pb->pos[0] = px; + pb->pos[1] = py; + pb->pos[2] = pz; mjuu_zerovec(pb->ipos, 3); + + // mass distribution if (doftype == mjFCOMPDOF_TRILINEAR) { - pb->mass = mass / 8; + pb->mass = node_mass_uniform; } else { - pb->mass = mass * massP2[i] * massP2[j] * massP2[k]; + // local index within the cell for mass computation + int li = gi % flex->spec.order; + int lj = gj % flex->spec.order; + int lk = gk % flex->spec.order; + // boundary nodes: average mass contribution + int ncells_i = (gi > 0 && gi < nx-1 && li == 0) ? 2 : 1; + int ncells_j = (gj > 0 && gj < ny-1 && lj == 0) ? 2 : 1; + int ncells_k = (gk > 0 && gk < nz-1 && lk == 0) ? 2 : 1; + // use Simpson weights scaled by cell count + double wi = massP2[li == 0 ? 0 : li]; + double wj = massP2[lj == 0 ? 0 : lj]; + double wk = massP2[lk == 0 ? 0 : lk]; + pb->mass = mass * wi * wj * wk * ncells_i * ncells_j * ncells_k + / (flex->spec.cellcount[0] * flex->spec.cellcount[1] * flex->spec.cellcount[2]); } + pb->inertia[0] = pb->mass*(2.0*inertiabox*inertiabox)/3.0; pb->inertia[1] = pb->mass*(2.0*inertiabox*inertiabox)/3.0; pb->inertia[2] = pb->mass*(2.0*inertiabox*inertiabox)/3.0; @@ -607,7 +662,7 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf // construct node name, add to nodebody char txt[100]; - mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), i, j, k); + mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), gi, gj, gk); mjs_setName(pb->element, txt); mjs_appendString(pf->nodebody, mjs_getName(pb->element)->c_str()); diff --git a/src/user/user_flexcomp.h b/src/user/user_flexcomp.h index 09e13c41..8624be7b 100644 --- a/src/user/user_flexcomp.h +++ b/src/user/user_flexcomp.h @@ -78,6 +78,7 @@ class mjCFlexcomp { std::string name; // flex name mjtFcompType type; // flexcomp type int count[3]; // grid count in each dimension + int cellcount[3]; // number of cells for interpolation double spacing[3]; // spacing between grid elements double scale[3]; // scaling for mesh and direct double origin[3]; // origin for generating a 3D mesh from a convex 2D mesh diff --git a/src/user/user_init.c b/src/user/user_init.c index 82db341b..1e522776 100644 --- a/src/user/user_init.c +++ b/src/user/user_init.c @@ -224,6 +224,9 @@ void mjs_defaultFlex(mjsFlex* flex) { // set other defaults flex->dim = 2; flex->radius = 0.005; + flex->cellcount[0] = 1; + flex->cellcount[1] = 1; + flex->cellcount[2] = 1; flex->internal = 0; flex->selfcollide = mjFLEXSELF_AUTO; flex->activelayers = 1; diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index dc2e5257..51a940a4 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -3961,7 +3961,10 @@ std::string mjCFlex::ComputeStiffnessCacheKey() const { combine(std::hash{}(young)); combine(std::hash{}(poisson)); - combine(std::hash{}(order_)); + combine(std::hash{}(spec.order)); + combine(std::hash{}(spec.cellcount[0])); + combine(std::hash{}(spec.cellcount[1])); + combine(std::hash{}(spec.cellcount[2])); // compute bounding box from vertex positions if (!vert_.empty()) { @@ -4086,10 +4089,20 @@ void mjCFlex::Compile(const mjVFS* vfs) { // set nnode nnode = static_cast(nodebody_.size()); - if (nnode && !order_) { - order_ = std::pow(nnode, 1.0 / 3) - 1; - if (nnode != std::pow(order_ + 1, 3)) { - throw mjCError(this, "number of nodes must be %d^3 but it is %d", nullptr, order_, nnode); + if (nnode && !spec.order) { + throw mjCError(this, "Interpolation order must be explicitly specified (dof is missing)"); + } + + // check node compatibility with count and dof + if (spec.order > 0) { + int expected_nodes = (spec.cellcount[0] * spec.order + 1) * + (spec.cellcount[1] * spec.order + 1) * + (spec.cellcount[2] * spec.order + 1); + if (nnode != expected_nodes) { + std::string msg = "number of nodes (" + std::to_string(nnode) + + ") does not match cellcount and dof expected (" + + std::to_string(expected_nodes) + ")"; + throw mjCError(this, msg.c_str()); } } @@ -4329,12 +4342,48 @@ void mjCFlex::Compile(const mjVFS* vfs) { } if (!stiffness_cached && young > 0 && interpolated) { - int n = pow(order_ + 1, 3); - int ndof = 3 * n; - if (stiffness.size() < ndof * ndof) { - stiffness.resize(ndof * ndof, 0); + 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]; + int ncells = cx * cy * cz; + int ny_global = cy * spec.order + 1; + int nz_global = cz * spec.order + 1; + + // total stiffness = ncells * ndof_cell^2 + stiffness.resize(ncells * ndof_cell * ndof_cell, 0); + + // compute stiffness per cell + for (int ci = 0; ci < cx; ci++) { + for (int cj = 0; cj < cy; cj++) { + for (int ck = 0; ck < cz; ck++) { + int cell_idx = ci * cy * cz + cj * cz + ck; + + // gather cell's local node positions + std::vector cell_pos(3 * npc); + int local = 0; + for (int li = 0; li <= spec.order; li++) { + for (int lj = 0; lj <= spec.order; lj++) { + for (int lk = 0; lk <= spec.order; lk++) { + int gi = ci * spec.order + li; + int gj = cj * spec.order + lj; + int gk = ck * spec.order + lk; + int global = gi * ny_global * nz_global + gj * nz_global + gk; + mjuu_copyvec(cell_pos.data() + 3*local, nodexpos.data() + 3*global, 3); + local++; + } + } + } + + // compute per-cell stiffness + std::vector K_cell(ndof_cell * ndof_cell, 0); + ComputeLinearStiffness(K_cell, cell_pos.data(), young, poisson, spec.order); + + // copy into global stiffness array + mjuu_copyvec(stiffness.data() + cell_idx * ndof_cell * ndof_cell, + K_cell.data(), ndof_cell * ndof_cell); + } + } } - ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_); } // create bounding volume hierarchy diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 64048736..85245020 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -2190,8 +2190,13 @@ void mjCModel::SetSizes() { nflexshelldata += (int)flexes_[i]->shell.size(); nflexevpair += (int)flexes_[i]->evpair.size()/2; nflextexcoord += (flexes_[i]->HasTexcoord() ? flexes_[i]->get_texcoord().size()/2 : 0); - if (flexes_[i]->order_ != 0) { - extra_stiffness_size += (3 * flexes_[i]->nnode) * (3 * flexes_[i]->nnode); + if (flexes_[i]->spec.order != 0) { + int npc = (int)pow(flexes_[i]->spec.order + 1, 3); + int ndof_cell = 3 * npc; + int ncells = flexes_[i]->spec.cellcount[0] * + flexes_[i]->spec.cellcount[1] * + flexes_[i]->spec.cellcount[2]; + extra_stiffness_size += ncells * ndof_cell * ndof_cell; } if (flexes_[i]->interpolated || flexes_[i]->rigid) { continue; @@ -3467,18 +3472,30 @@ void mjCModel::CopyObjects(mjModel* m) { mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4); // elasticity - if (pfl->order_ == 0) { + if (pfl->spec.order == 0) { m->flex_stiffnessadr[i] = 21 * elem_adr; } else { m->flex_stiffnessadr[i] = current_extra_stiffness_adr; - current_extra_stiffness_adr += (3 * pfl->nnode) * (3 * pfl->nnode); + int npc = (int)pow(pfl->spec.order + 1, 3); + int ndof_cell = 3 * npc; + int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2]; + current_extra_stiffness_adr += ncells * ndof_cell * ndof_cell; } if (!pfl->stiffness.empty()) { - mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i], pfl->stiffness.data(), pfl->stiffness.size()); + mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i], + pfl->stiffness.data(), pfl->stiffness.size()); } else { - int size = (pfl->order_ == 0) ? 21 * pfl->nelem : (3 * pfl->nnode) * (3 * pfl->nnode); - mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], size); + int stiff_size; + if (pfl->spec.order == 0) { + stiff_size = 21 * pfl->nelem; + } else { + int npc = (int)pow(pfl->spec.order + 1, 3); + int ndof_cell = 3 * npc; + int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2]; + stiff_size = ncells * ndof_cell * ndof_cell; + } + mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], stiff_size); } if (!pfl->bending.empty()) { mjuu_copyvec(m->flex_bending + 17 * edge_adr, pfl->bending.data(), pfl->bending.size()); @@ -3613,7 +3630,12 @@ void mjCModel::CopyObjects(mjModel* m) { } // set interpolation type, only two types for now - m->flex_interp[i] = pfl->order_; + m->flex_interp[i] = pfl->spec.order; + + // set cell count for multi-cell finite cell method + m->flex_cellnum[3*i+0] = pfl->spec.cellcount[0]; + m->flex_cellnum[3*i+1] = pfl->spec.cellcount[1]; + m->flex_cellnum[3*i+2] = pfl->spec.cellcount[2]; // convert edge pairs to int array, set edge rigid for (int k=0; k < pfl->nedge; k++) { diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 232d0470..8ca60124 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -1042,7 +1042,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex { static constexpr int kNumEdges[3] = {1, 3, 6}; // number of edges per element indexed by dim - void SetOrder(int order) { order_ = order; } // set interpolation order + private: void Compile(const mjVFS* vfs); // compiler @@ -1052,7 +1052,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex { std::vector vert0_; // vertex positions in [0, 1]^d in the bounding box std::vector node0_; // node Cartesian positions - int order_ = 0; // interpolation order + // stiffness caching std::string ComputeStiffnessCacheKey() const; diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc index df10509f..d38bde20 100644 --- a/src/xml/xml_native_reader.cc +++ b/src/xml/xml_native_reader.cc @@ -315,7 +315,7 @@ std::vector MJCF[nMJCF] = { {">"}, {">"}, {"flexcomp", "*", "name", "type", "group", "dim", "dof", - "count", "spacing", "radius", "rigid", "mass", "inertiabox", + "count", "cellcount", "spacing", "radius", "rigid", "mass", "inertiabox", "scale", "file", "point", "element", "texcoord", "material", "rgba", "flatskin", "pos", "quat", "axisangle", "xyaxes", "zaxis", "euler", "origin"}, {"<"}, @@ -334,8 +334,8 @@ std::vector MJCF[nMJCF] = { {"deformable", "*"}, {"<"}, - {"flex", "*", "name", "group", "dim", "radius", "material", - "rgba", "flatskin", "body", "vertex", "element", "texcoord", "elemtexcoord", "node"}, + {"flex", "*", "name", "group", "dim", "radius", "material", "rgba", "flatskin", "body", + "vertex", "element", "texcoord", "elemtexcoord", "node", "cellcount", "dof"}, {"<"}, {"contact", "?", "contype", "conaffinity", "condim", "priority", "friction", "solmix", "solref", "solimp", "margin", "gap", @@ -1501,6 +1501,16 @@ void mjXReader::OneFlex(XMLElement* elem, mjsFlex* flex) { ReadAttrInt(elem, "dim", &flex->dim); ReadAttrInt(elem, "group", &flex->group); + flex->cellcount[0] = 1; + flex->cellcount[1] = 1; + flex->cellcount[2] = 1; + ReadAttr(elem, "cellcount", 3, flex->cellcount, text); + + flex->order = 0; + if (MapValue(elem, "dof", &n, fdof_map, mjNFCOMPDOFS)) { + flex->order = (n == mjFCOMPDOF_QUADRATIC) ? 2 : (n == mjFCOMPDOF_TRILINEAR ? 1 : 0); + } + // read data vectors if (ReadAttrTxt(elem, "body", text, true)) { mjs_setStringVec(flex->vertbody, text.c_str()); @@ -2794,6 +2804,7 @@ void mjXReader::OneFlexcomp(XMLElement* elem, mjsBody* body, const mjVFS* vfs) { fcomp.type = (mjtFcompType)n; } ReadAttr(elem, "count", 3, fcomp.count, text); + ReadAttr(elem, "cellcount", 3, fcomp.cellcount, text); ReadAttr(elem, "spacing", 3, fcomp.spacing, text); ReadAttr(elem, "scale", 3, fcomp.scale, text); ReadAttr(elem, "mass", 1, &fcomp.mass, text); diff --git a/src/xml/xml_native_writer.cc b/src/xml/xml_native_writer.cc index e9aa9ea7..007a76a3 100644 --- a/src/xml/xml_native_writer.cc +++ b/src/xml/xml_native_writer.cc @@ -141,6 +141,13 @@ void mjXWriter::OneFlex(XMLElement* elem, const mjCFlex* flex) { WriteAttrKey(elem, "flatskin", bool_map, 2, flex->flatskin, defflex.flatskin); WriteAttrInt(elem, "dim", flex->dim, defflex.dim); WriteAttrInt(elem, "group", flex->group, defflex.group); + WriteAttr(elem, "cellcount", 3, flex->spec.cellcount, defflex.spec.cellcount); + if (flex->spec.order != defflex.spec.order) { + string dof_str = "full"; + if (flex->spec.order == 1) dof_str = "trilinear"; + else if (flex->spec.order == 2) dof_str = "quadratic"; + WriteAttrTxt(elem, "dof", dof_str); + } // data vectors if (!flex->get_vertbody().empty()) { diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index 7b055936..d795a404 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -35,136 +35,9 @@ namespace { using ::testing::NotNull; using ::testing::Pointwise; + using CoreConstraintTest = MujocoTest; -// compute rotation residual following formula in mj_instantiateEquality -void RotationResidual(const mjModel *model, mjData *data, - const mjtNum qpos[7], const mjtNum dqpos[6], - mjtNum res[3]) { - // copy configuration, compute required quantities with mj_step1 - mju_copy(data->qpos, qpos, 7); - - // perturb configuration if given - if (dqpos) { - mj_integratePos(model, data->qpos, dqpos, 1); - } - - // update relevant quantities - mj_step1(model, data); - - // compute orientation residual - mjtNum quat1[4], quat2[4], quat3[4]; - mju_copy4(quat1, data->xquat+4*1); - mju_negQuat(quat2, data->xquat+4*2); - mju_mulQuat(quat3, quat2, quat1); - mju_copy3(res, quat3+1); -} - -// validate rotational Jacobian used in welds -TEST_F(CoreConstraintTest, WeldRotJacobian) { -#ifdef mjUSESINGLE - GTEST_SKIP() << "FD Jacobian with eps=1e-6 below float32 precision"; -#endif - constexpr char xml[] = R"( - - - )"; - char error[1024]; - mjModel* model = LoadModelFromString(xml, error, sizeof(error)); - ASSERT_THAT(model, testing::NotNull()) << error; - ASSERT_EQ(model->nq, 7); - ASSERT_EQ(model->nv, 6); - static const int nv = 6; // for increased readability - mjData* data = mj_makeData(model); - - // arbitrary initial values for the ball and hinge joints - mjtNum qpos0[7] = {.5, .5, .5, .5, .7, .8, .9}; - - // compute required quantities using mj_step1 - mj_step1(model, data); - - // get orientation error - mjtNum res[3]; - RotationResidual(model, data, qpos0, NULL, res); - - // compute Jacobian with finite-differencing - mjtNum jacFD[3*nv]; - mjtNum dqpos[nv] = {0}; - mjtNum dres[3]; - const mjtNum eps = 1e-6; - for (int i=0; i < nv; i++) { - // nudge i-th dof - dqpos[i] = eps; - - // get nudged residual - RotationResidual(model, data, qpos0, dqpos, dres); - - // remove nudge - dqpos[i] = 0.0; - - // compute Jacobian column - for (int j=0; j < 3; j++) { - jacFD[nv*j + i] = (dres[j] - res[j]) / eps; - } - } - - // reset mjData to qpos0 - mju_copy(data->qpos, qpos0, 7); - mj_step1(model, data); - - // intermediate quaternions quat1 and quat2 - mjtNum quat1[4], negQuat2[4]; - mju_copy4(quat1, data->xquat+4*1); - mju_negQuat(negQuat2, data->xquat+4*2); - - // get analytical Jacobian following formula in mj_instantiateEquality - mjtNum jacdif[3*nv], jac0[3*nv], jac1[3*nv]; - mjtNum point[3] = {0}; - - // rotational Jacobian difference - mj_jacDifPair(model, data, NULL, 2, 1, point, point, - NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model), - /*flg_skipcommon=*/0); - - // formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1 - mjtNum axis[3], quat3[4], quat4[4]; - for (int j=0; j < nv; j++) { - // axis = [jac1-jac2]_col(j) - axis[0] = jacdif[0*nv+j]; - axis[1] = jacdif[1*nv+j]; - axis[2] = jacdif[2*nv+j]; - - // apply formula - mju_mulQuatAxis(quat3, negQuat2, axis); - mju_mulQuat(quat4, quat3, quat1); - - // correct Jacobian - jacdif[0*nv+j] = 0.5*quat4[1]; - jacdif[1*nv+j] = 0.5*quat4[2]; - jacdif[2*nv+j] = 0.5*quat4[3]; - } - - // test that analytical and finite-differenced Jacobians match - EXPECT_THAT(AsVector(jacFD, 3*nv), - Pointwise(MjNear(eps, 1e-3), AsVector(jacdif, 3*nv))); - - mj_deleteData(data); - mj_deleteModel(model); -} - // test formulas for penetration at rest TEST_F(CoreConstraintTest, RestPenetration) { constexpr char xml[] = R"( diff --git a/test/engine/engine_core_util_test.cc b/test/engine/engine_core_util_test.cc new file mode 100644 index 00000000..8cd73dca --- /dev/null +++ b/test/engine/engine_core_util_test.cc @@ -0,0 +1,740 @@ +// Copyright 2026 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. + +// Tests for engine/engine_core_util.c. + +#include "src/engine/engine_core_util.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include "test/fixture.h" + +namespace mujoco { +namespace { + +using ::testing::NotNull; +using ::testing::Pointwise; + +using FlexGatherStateTest = MujocoTest; + +TEST_F(FlexGatherStateTest, mju_flexGatherState_Grid) { + static constexpr char xml[] = R"( + + + + + + + + + )"; + + char error[1024]; + mjModel* model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + mjData* data = mj_makeData(model); + mj_forward(model, data); + + ASSERT_EQ(model->nflex, 1); + int f = 0; + int nodenum = model->flex_nodenum[f]; + int nstart = model->flex_nodeadr[f]; + + // Simulate a rotated state (90 degrees around Z axis) + ASSERT_TRUE(model->flex_centered[f]); + for (int i = 0; i < nodenum; i++) { + int b = model->flex_nodebodyid[nstart + i]; + mjtNum x = data->xpos[3*b + 0]; + mjtNum y = data->xpos[3*b + 1]; + mjtNum z = data->xpos[3*b + 2]; + + // Rotate 90 degrees around Z: (x, y, z) -> (-y, x, z) + data->xpos[3*b + 0] = -y; + data->xpos[3*b + 1] = x; + data->xpos[3*b + 2] = z; + } + + std::vector xpos(3 * nodenum); + mju_flexGatherState(model, data, f, xpos.data(), NULL); + + // Verify that gathered xpos matches the rotated data->xpos + for (int i = 0; i < nodenum; i++) { + int b = model->flex_nodebodyid[nstart + i]; + EXPECT_NEAR(xpos[3*i + 0], data->xpos[3*b + 0], 1e-5); + EXPECT_NEAR(xpos[3*i + 1], data->xpos[3*b + 1], 1e-5); + EXPECT_NEAR(xpos[3*i + 2], data->xpos[3*b + 2], 1e-5); + } + + mj_deleteData(data); + mj_deleteModel(model); +} + + +using AngMomMatTest = MujocoTest; + +static constexpr char AngMomTestingModel[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +// compare subtree angular momentum computed in two ways +TEST_F(AngMomMatTest, CompareAngMom) { + char error[1024]; + mjModel* model = + LoadModelFromString(AngMomTestingModel, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + int nv = model->nv; + int bodyid = mj_name2id(model, mjOBJ_BODY, "link1"); + + mjData* data = mj_makeData(model); + + // reset to the keyframe with some angular velocities + mj_resetDataKeyframe(model, data, 0); + mj_forward(model, data); + + // get the reference value of angular momentum + mj_subtreeVel(model, data); + mjtNum angmom_ref[3]; + mju_copy3(angmom_ref, data->subtree_angmom+3*bodyid); + + // compute angular momentum using the angular momentum matrix + mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); + mj_angmomMat(model, data, angmom_mat, bodyid); + mjtNum angmom_test[3]; + mju_mulMatVec(angmom_test, angmom_mat, data->qvel, 3, nv); + + // compare the two angular momentum values + for (int i = 0; i < 3; i++) { + EXPECT_THAT(angmom_ref[i], MjNear(angmom_test[i], 1e-8, 1e-4)); + } + + mju_free(angmom_mat); + mj_deleteData(data); + mj_deleteModel(model); +} + +// compare subtree angular momentum matrix: analytical and findiff +TEST_F(AngMomMatTest, CompareAngMomMats) { + char error[1024]; + mjModel* model = + LoadModelFromString(AngMomTestingModel, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + int nv = model->nv; + int bodyid = mj_name2id(model, mjOBJ_BODY, "link1"); + mjData* data = mj_makeData(model); + mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); + mjtNum* angmom_mat_fd = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); + + // reset to the keyframe with some angular velocities + mj_resetDataKeyframe(model, data, 0); + mj_forward(model, data); + + // compute the angular momentum matrix using the analytical method + mj_angmomMat(model, data, angmom_mat, bodyid); + + // compute the angular momentum matrix using finite differences + static constexpr mjtNum eps = MjTol(1e-6, 1e-3); + for (int i = 0; i < nv; i++) { + // reset vel, forward nudge i-th dof, get angmom + mju_copy(data->qvel, model->key_qvel, model->nv); + data->qvel[i] += eps; + mj_forward(model, data); + mj_subtreeVel(model, data); + mjtNum agmf[3]; + mju_copy3(agmf, data->subtree_angmom+3*bodyid); + + // reset vel, backward nudge i-th dof, get angmom + mju_copy(data->qvel, model->key_qvel, model->nv); + data->qvel[i] -= eps; + mj_forward(model, data); + mj_subtreeVel(model, data); + mjtNum agmb[3]; + mju_copy3(agmb, data->subtree_angmom+3*bodyid); + + // finite-difference the angmom matrix + for (int j = 0; j < 3; j++) { + angmom_mat_fd[nv*j+i] = (agmf[j] - agmb[j]) / (2 * eps); + } + } + + // compare the two matrices + for (int i = 0; i < 3*nv; i++) { + EXPECT_THAT(angmom_mat_fd[i], MjNear(angmom_mat[i], 1e-8, 2e-4)); + } + + mju_free(angmom_mat_fd); + mju_free(angmom_mat); + mj_deleteData(data); + mj_deleteModel(model); +} + +using JacobianTest = MujocoTest; +static const mjtNum max_abs_err = std::numeric_limits::epsilon(); + +static constexpr char kJacobianTestingModel[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +// compare analytic and finite-differenced subtree-com Jacobian +TEST_F(JacobianTest, SubtreeJac) { + char error[1024]; + mjModel* model = + LoadModelFromString(kJacobianTestingModel, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + int nv = model->nv; + int bodyid = mj_name2id(model, mjOBJ_BODY, "main"); + mjData* data = mj_makeData(model); + mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); + mjtNum* qpos = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nq); + mjtNum* nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv); + + // all we need for Jacobians are kinematics and CoM-related quantities + mj_kinematics(model, data); + mj_comPos(model, data); + + // get subtree CoM Jacobian of free body + mj_jacSubtreeCom(model, data, jac_subtree, bodyid); + + // save current subtree-com and qpos, clear nudge + mjtNum subtree_com[3]; + mju_copy3(subtree_com, data->subtree_com+3*bodyid); + mju_copy(qpos, data->qpos, model->nq); + mju_zero(nudge, nv); + + // compare analytic Jacobian to finite-difference approximation + static const mjtNum eps = 1e-6; + for (int i=0; i < nv; i++) { + // reset qpos, nudge i-th dof, update data->qpos, reset nudge + mju_copy(data->qpos, qpos, model->nq); + nudge[i] = 1; + mj_integratePos(model, data->qpos, nudge, eps); + nudge[i] = 0; + + // kinematics and comPos to get nudged com + mj_kinematics(model, data); + mj_comPos(model, data); + + // compare finite-differenced and analytic Jacobian + for (int j=0; j < 3; j++) { + mjtNum findiff = (data->subtree_com[3*bodyid+j] - subtree_com[j]) / eps; + EXPECT_THAT(jac_subtree[nv*j+i], MjNear(findiff, eps, 1e-2)); + } + } + + mju_free(nudge); + mju_free(qpos); + mju_free(jac_subtree); + mj_deleteData(data); + mj_deleteModel(model); +} + +// confirm that applying linear forces via the subtree-com Jacobian only creates +// the expected linear accelerations (no accelerations of internal joints) +TEST_F(JacobianTest, SubtreeJacNoInternalAcc) { + char error[1024]; + mjModel* model = + LoadModelFromString(kJacobianTestingModel, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + int nv = model->nv; + int bodyid = mj_name2id(model, mjOBJ_BODY, "main"); + mjData* data = mj_makeData(model); + mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); + + // all we need for Jacobians are kinematics and CoM-related quantities + mj_kinematics(model, data); + mj_comPos(model, data); + + // get subtree CoM Jacobian of free body + mj_jacSubtreeCom(model, data, jac_subtree, bodyid); + + // uncomment for debugging + // mju_printMat(jac_subtree, 3, nv); + + // call fwdPosition since we'll need the factorised mass matrix in the test + mj_fwdPosition(model, data); + + // treating the subtree Jacobian as the projection of 3 axis-aligned unit + // forces into joint space, solve for the resulting accelerations in-place + mj_solveM(model, data, jac_subtree, jac_subtree, 3); + + // expect to find accelerations of magnitude 1/subtreemass in the first 3 + // coordinates of the free joint and 0s elsewhere, since applying forces to + // the CoM should accelerate the whole mechanism without any internal motion + int body_dofadr = model->body_dofadr[bodyid]; + mjtNum invtreemass = 1.0/model->body_subtreemass[bodyid]; + for (int r = 0; r < 3; r++) { + for (int c = 0; c < nv; c++) { + mjtNum expected = c - body_dofadr == r ? invtreemass : 0.0; + EXPECT_THAT(jac_subtree[nv*r+c], MjNear(expected, max_abs_err, 1e-4)); + } + } + + mju_free(jac_subtree); + mj_deleteData(data); + mj_deleteModel(model); +} + +static constexpr char kQuat[] = R"( + + + + + + + + + + + + +)"; + +static constexpr char kFreeBall[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +static constexpr char kQuatlessPendulum[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +static constexpr char kTelescope[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +static constexpr char kHinge[] = R"( + + + + + + + + + + + + + +)"; + +// compare mj_jacDot with finite-differenced mj_jac +TEST_F(JacobianTest, JacDot) { + for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) { + char error[1024]; + mjModel* model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + int nv = model->nv; + mjData* data = mj_makeData(model); + + // load keyframe if present, step for a bit + if (model->nkey) mj_resetDataKeyframe(model, data, 0); + while (data->time < 0.1) { + mj_step(model, data); + } + + // minimal call required for mj_jacDot outputs to be valid + mj_kinematics(model, data); + mj_comPos(model, data); + mj_comVel(model, data); + + // get bodyid + int bodyid = mj_name2id(model, mjOBJ_BODY, "query"); + EXPECT_GT(bodyid, 0); + + // get site position + int siteid = mj_name2id(model, mjOBJ_SITE, "query"); + EXPECT_GT(siteid, -1); + mjtNum point[3]; + mju_copy3(point, data->site_xpos+3*siteid); + + // jac, jac_dot + std::vector jacp(3*nv); + std::vector jacr(3*nv); + mj_jac(model, data, jacp.data(), jacr.data(), point, bodyid); + std::vector jacp_dot(3*nv); + std::vector jacr_dot(3*nv); + mj_jacDot(model, data, jacp_dot.data(), jacr_dot.data(), point, bodyid); + + // jac_h: jacobian after integrating qpos with a timestep of h + constexpr mjtNum h = MjTol(1e-7, 5e-4); + mj_integratePos(model, data->qpos, data->qvel, h); + mj_kinematics(model, data); + mj_comPos(model, data); + std::vector jacp_h(3*nv); + std::vector jacr_h(3*nv); + mju_copy3(point, data->site_xpos+3*siteid); // get updated site position + mj_jac(model, data, jacp_h.data(), jacr_h.data(), point, bodyid); + + // jac_dot_h finite-difference approximation + std::vector jacp_dot_h(3*nv); + mju_sub(jacp_dot_h.data(), jacp_h.data(), jacp.data(), 3*nv); + mju_scl(jacp_dot_h.data(), jacp_dot_h.data(), 1/h, 3*nv); + std::vector jacr_dot_h(3*nv); + mju_sub(jacr_dot_h.data(), jacr_h.data(), jacr.data(), 3*nv); + mju_scl(jacr_dot_h.data(), jacr_dot_h.data(), 1/h, 3*nv); + + // compare finite-differenced and analytic + mjtNum tol = 1e-5; + EXPECT_THAT(jacp_dot, Pointwise(MjNear(tol, 5e-2), jacp_dot_h)); + EXPECT_THAT(jacr_dot, Pointwise(MjNear(tol, 5e-2), jacr_dot_h)); + + mj_deleteData(data); + mj_deleteModel(model); + } +} + +// compare mj_jacDotSparse with dense mj_jacDot +TEST_F(JacobianTest, JacDotSparse) { + for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) { + char error[1024]; + mjModel* model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + int nv = model->nv; + mjData* data = mj_makeData(model); + + // load keyframe if present, step for a bit + if (model->nkey) mj_resetDataKeyframe(model, data, 0); + while (data->time < 0.1) { + mj_step(model, data); + } + + // minimal call required for mj_jacDot outputs to be valid + mj_kinematics(model, data); + mj_comPos(model, data); + mj_comVel(model, data); + + // get bodyid and site position + int bodyid = mj_name2id(model, mjOBJ_BODY, "query"); + EXPECT_GT(bodyid, 0); + int siteid = mj_name2id(model, mjOBJ_SITE, "query"); + EXPECT_GT(siteid, -1); + mjtNum point[3]; + mju_copy3(point, data->site_xpos+3*siteid); + + // dense jacDot + std::vector jacp_dense(3*nv); + std::vector jacr_dense(3*nv); + mj_jacDot(model, data, jacp_dense.data(), jacr_dense.data(), point, bodyid); + + // compute body chain using public mjModel fields + std::vector chain(nv); + int NV = 0; + int weldbody = model->body_weldid[bodyid]; + if (weldbody) { + int da = model->body_dofadr[weldbody] + model->body_dofnum[weldbody] - 1; + while (da >= 0) { + chain[NV++] = da; + da = model->dof_parentid[da]; + } + std::reverse(chain.begin(), chain.begin() + NV); + } + EXPECT_GT(NV, 0); + + // sparse jacDot + std::vector jacp_sparse(3*NV); + std::vector jacr_sparse(3*NV); + mj_jacDotSparse(model, data, jacp_sparse.data(), jacr_sparse.data(), + point, bodyid, NV, chain.data()); + + // expand sparse to dense and compare + std::vector jacp_expanded(3*nv, 0); + std::vector jacr_expanded(3*nv, 0); + for (int ci = 0; ci < NV; ci++) { + int di = chain[ci]; + for (int r = 0; r < 3; r++) { + jacp_expanded[di+r*nv] = jacp_sparse[ci+r*NV]; + jacr_expanded[di+r*nv] = jacr_sparse[ci+r*NV]; + } + } + + // expect bitwise equality + EXPECT_EQ(jacp_expanded, jacp_dense); + EXPECT_EQ(jacr_expanded, jacr_dense); + + mj_deleteData(data); + mj_deleteModel(model); + } +} + + +// validate rotational Jacobian used in welds +TEST_F(JacobianTest, WeldRotJacobian) { +#ifdef mjUSESINGLE + GTEST_SKIP() << "FD Jacobian with eps=1e-6 below float32 precision"; +#endif + constexpr char xml[] = R"( + + + )"; + char error[1024]; + mjModel* model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model, testing::NotNull()) << error; + ASSERT_EQ(model->nq, 7); + ASSERT_EQ(model->nv, 6); + static const int nv = 6; // for increased readability + mjData* data = mj_makeData(model); + + // arbitrary initial values for the ball and hinge joints + mjtNum qpos0[7] = {.5, .5, .5, .5, .7, .8, .9}; + + // compute required quantities using mj_step1 + mj_step1(model, data); + + // get orientation error + mjtNum res[3]; + // compute rotation residual following formula in mj_instantiateEquality + auto RotationResidual = [](const mjModel *model, mjData *data, + const mjtNum qpos[7], const mjtNum dqpos[6], + mjtNum res[3]) { + // copy configuration, compute required quantities with mj_step1 + mju_copy(data->qpos, qpos, 7); + + // perturb configuration if given + if (dqpos) { + mj_integratePos(model, data->qpos, dqpos, 1); + } + + // update relevant quantities + mj_step1(model, data); + + // compute orientation residual + mjtNum quat1[4], quat2[4], quat3[4]; + mju_copy4(quat1, data->xquat+4*1); + mju_negQuat(quat2, data->xquat+4*2); + mju_mulQuat(quat3, quat2, quat1); + mju_copy3(res, quat3+1); + }; + + RotationResidual(model, data, qpos0, NULL, res); + + // compute Jacobian with finite-differencing + mjtNum jacFD[3*nv]; + mjtNum dqpos[nv] = {0}; + mjtNum dres[3]; + const mjtNum eps = 1e-6; + for (int i=0; i < nv; i++) { + // nudge i-th dof + dqpos[i] = eps; + + // get nudged residual + RotationResidual(model, data, qpos0, dqpos, dres); + + // remove nudge + dqpos[i] = 0.0; + + // compute Jacobian column + for (int j=0; j < 3; j++) { + jacFD[nv*j + i] = (dres[j] - res[j]) / eps; + } + } + + // reset mjData to qpos0 + mju_copy(data->qpos, qpos0, 7); + mj_step1(model, data); + + // intermediate quaternions quat1 and quat2 + mjtNum quat1[4], negQuat2[4]; + mju_copy4(quat1, data->xquat+4*1); + mju_negQuat(negQuat2, data->xquat+4*2); + + // get analytical Jacobian following formula in mj_instantiateEquality + mjtNum jacdif[3*nv], jac0[3*nv], jac1[3*nv]; + mjtNum point[3] = {0}; + + // rotational Jacobian difference + mj_jacDifPair(model, data, NULL, 2, 1, point, point, + NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model), + /*flg_skipcommon=*/0); + + // formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1 + mjtNum axis[3], quat3[4], quat4[4]; + for (int j=0; j < nv; j++) { + // axis = [jac1-jac2]_col(j) + axis[0] = jacdif[0*nv+j]; + axis[1] = jacdif[1*nv+j]; + axis[2] = jacdif[2*nv+j]; + + // apply formula + mju_mulQuatAxis(quat3, negQuat2, axis); + mju_mulQuat(quat4, quat3, quat1); + + // correct Jacobian + jacdif[0*nv+j] = 0.5*quat4[1]; + jacdif[1*nv+j] = 0.5*quat4[2]; + jacdif[2*nv+j] = 0.5*quat4[3]; + } + + // test that analytical and finite-differenced Jacobians match + EXPECT_THAT(AsVector(jacFD, 3*nv), + Pointwise(MjNear(eps, 1e-3), AsVector(jacdif, 3*nv))); + + mj_deleteData(data); + mj_deleteModel(model); +} + +} // namespace +} // namespace mujoco + + diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc index 4b0caf9c..3e36934f 100644 --- a/test/engine/engine_forward_test.cc +++ b/test/engine/engine_forward_test.cc @@ -3107,7 +3107,7 @@ TEST_F(ForwardTest, FlexParentCoupling) { - @@ -3152,7 +3152,7 @@ TEST_F(ForwardTest, FlexParentCoupling) { if (diff > max_diff) max_diff = diff; } - EXPECT_LT(max_diff, MjTol(2e-5, 5e-3)) + EXPECT_LT(max_diff, MjTol(2e-5, 1.5e-2)) << "Implicit integrator should match Euler at small timestep"; mj_deleteData(data); diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index 629517ae..f52d374f 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -14,12 +14,9 @@ // Tests for engine/{engine_support.c and engine_core_util.c} -#include "src/engine/engine_core_util.h" #include "src/engine/engine_support.h" -#include #include -#include #include #include #include @@ -41,521 +38,9 @@ using ::testing::Ne; using ::testing::NotNull; using ::testing::Pointwise; -using AngMomMatTest = MujocoTest; -static constexpr char AngMomTestingModel[] = R"( - - - - - - - - - - - - - - - - - - - - - - - - -)"; -// compare subtree angular momentum computed in two ways -TEST_F(AngMomMatTest, CompareAngMom) { - char error[1024]; - mjModel* model = - LoadModelFromString(AngMomTestingModel, error, sizeof(error)); - ASSERT_THAT(model, NotNull()) << error; - int nv = model->nv; - int bodyid = mj_name2id(model, mjOBJ_BODY, "link1"); - mjData* data = mj_makeData(model); - - // reset to the keyframe with some angular velocities - mj_resetDataKeyframe(model, data, 0); - mj_forward(model, data); - - // get the reference value of angular momentum - mj_subtreeVel(model, data); - mjtNum angmom_ref[3]; - mju_copy3(angmom_ref, data->subtree_angmom+3*bodyid); - - // compute angular momentum using the angular momentum matrix - mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); - mj_angmomMat(model, data, angmom_mat, bodyid); - mjtNum angmom_test[3]; - mju_mulMatVec(angmom_test, angmom_mat, data->qvel, 3, nv); - - // compare the two angular momentum values - for (int i = 0; i < 3; i++) { - EXPECT_THAT(angmom_ref[i], MjNear(angmom_test[i], 1e-8, 1e-4)); - } - - mju_free(angmom_mat); - mj_deleteData(data); - mj_deleteModel(model); -} - -// compare subtree angular momentum matrix: analytical and findiff -TEST_F(AngMomMatTest, CompareAngMomMats) { - char error[1024]; - mjModel* model = - LoadModelFromString(AngMomTestingModel, error, sizeof(error)); - ASSERT_THAT(model, NotNull()) << error; - int nv = model->nv; - int bodyid = mj_name2id(model, mjOBJ_BODY, "link1"); - mjData* data = mj_makeData(model); - mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); - mjtNum* angmom_mat_fd = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); - - // reset to the keyframe with some angular velocities - mj_resetDataKeyframe(model, data, 0); - mj_forward(model, data); - - // compute the angular momentum matrix using the analytical method - mj_angmomMat(model, data, angmom_mat, bodyid); - - // compute the angular momentum matrix using finite differences - static constexpr mjtNum eps = MjTol(1e-6, 1e-3); - for (int i = 0; i < nv; i++) { - // reset vel, forward nudge i-th dof, get angmom - mju_copy(data->qvel, model->key_qvel, model->nv); - data->qvel[i] += eps; - mj_forward(model, data); - mj_subtreeVel(model, data); - mjtNum agmf[3]; - mju_copy3(agmf, data->subtree_angmom+3*bodyid); - - // reset vel, backward nudge i-th dof, get angmom - mju_copy(data->qvel, model->key_qvel, model->nv); - data->qvel[i] -= eps; - mj_forward(model, data); - mj_subtreeVel(model, data); - mjtNum agmb[3]; - mju_copy3(agmb, data->subtree_angmom+3*bodyid); - - // finite-difference the angmom matrix - for (int j = 0; j < 3; j++) { - angmom_mat_fd[nv*j+i] = (agmf[j] - agmb[j]) / (2 * eps); - } - } - - // compare the two matrices - for (int i = 0; i < 3*nv; i++) { - EXPECT_THAT(angmom_mat_fd[i], MjNear(angmom_mat[i], 1e-8, 2e-4)); - } - - mju_free(angmom_mat_fd); - mju_free(angmom_mat); - mj_deleteData(data); - mj_deleteModel(model); -} - -using JacobianTest = MujocoTest; -static const mjtNum max_abs_err = std::numeric_limits::epsilon(); - -static constexpr char kJacobianTestingModel[] = R"( - - - - - - - - - - - - - - - - - - - - - - - - - - - - -)"; - -// compare analytic and finite-differenced subtree-com Jacobian -TEST_F(JacobianTest, SubtreeJac) { - char error[1024]; - mjModel* model = - LoadModelFromString(kJacobianTestingModel, error, sizeof(error)); - ASSERT_THAT(model, NotNull()) << error; - int nv = model->nv; - int bodyid = mj_name2id(model, mjOBJ_BODY, "main"); - mjData* data = mj_makeData(model); - mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); - mjtNum* qpos = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nq); - mjtNum* nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv); - - // all we need for Jacobians are kinematics and CoM-related quantities - mj_kinematics(model, data); - mj_comPos(model, data); - - // get subtree CoM Jacobian of free body - mj_jacSubtreeCom(model, data, jac_subtree, bodyid); - - // save current subtree-com and qpos, clear nudge - mjtNum subtree_com[3]; - mju_copy3(subtree_com, data->subtree_com+3*bodyid); - mju_copy(qpos, data->qpos, model->nq); - mju_zero(nudge, nv); - - // compare analytic Jacobian to finite-difference approximation - static const mjtNum eps = 1e-6; - for (int i=0; i < nv; i++) { - // reset qpos, nudge i-th dof, update data->qpos, reset nudge - mju_copy(data->qpos, qpos, model->nq); - nudge[i] = 1; - mj_integratePos(model, data->qpos, nudge, eps); - nudge[i] = 0; - - // kinematics and comPos to get nudged com - mj_kinematics(model, data); - mj_comPos(model, data); - - // compare finite-differenced and analytic Jacobian - for (int j=0; j < 3; j++) { - mjtNum findiff = (data->subtree_com[3*bodyid+j] - subtree_com[j]) / eps; - EXPECT_THAT(jac_subtree[nv*j+i], MjNear(findiff, eps, 1e-2)); - } - } - - mju_free(nudge); - mju_free(qpos); - mju_free(jac_subtree); - mj_deleteData(data); - mj_deleteModel(model); -} - -// confirm that applying linear forces via the subtree-com Jacobian only creates -// the expected linear accelerations (no accelerations of internal joints) -TEST_F(JacobianTest, SubtreeJacNoInternalAcc) { - char error[1024]; - mjModel* model = - LoadModelFromString(kJacobianTestingModel, error, sizeof(error)); - ASSERT_THAT(model, NotNull()) << error; - int nv = model->nv; - int bodyid = mj_name2id(model, mjOBJ_BODY, "main"); - mjData* data = mj_makeData(model); - mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); - - // all we need for Jacobians are kinematics and CoM-related quantities - mj_kinematics(model, data); - mj_comPos(model, data); - - // get subtree CoM Jacobian of free body - mj_jacSubtreeCom(model, data, jac_subtree, bodyid); - - // uncomment for debugging - // mju_printMat(jac_subtree, 3, nv); - - // call fwdPosition since we'll need the factorised mass matrix in the test - mj_fwdPosition(model, data); - - // treating the subtree Jacobian as the projection of 3 axis-aligned unit - // forces into joint space, solve for the resulting accelerations in-place - mj_solveM(model, data, jac_subtree, jac_subtree, 3); - - // expect to find accelerations of magnitude 1/subtreemass in the first 3 - // coordinates of the free joint and 0s elsewhere, since applying forces to - // the CoM should accelerate the whole mechanism without any internal motion - int body_dofadr = model->body_dofadr[bodyid]; - mjtNum invtreemass = 1.0/model->body_subtreemass[bodyid]; - for (int r = 0; r < 3; r++) { - for (int c = 0; c < nv; c++) { - mjtNum expected = c - body_dofadr == r ? invtreemass : 0.0; - EXPECT_THAT(jac_subtree[nv*r+c], MjNear(expected, max_abs_err, 1e-4)); - } - } - - mju_free(jac_subtree); - mj_deleteData(data); - mj_deleteModel(model); -} - -static constexpr char kQuat[] = R"( - - - - - - - - - - - - -)"; - -static constexpr char kFreeBall[] = R"( - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -)"; - -static constexpr char kQuatlessPendulum[] = R"( - - - - - - - - - - - - - - - - - - - - - - - - -)"; - -static constexpr char kTelescope[] = R"( - - - - - - - - - - - - - - - - - - - - - - - - -)"; - -static constexpr char kHinge[] = R"( - - - - - - - - - - - - - -)"; - -// compare mj_jacDot with finite-differenced mj_jac -TEST_F(JacobianTest, JacDot) { - for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) { - char error[1024]; - mjModel* model = LoadModelFromString(xml, error, sizeof(error)); - ASSERT_THAT(model, NotNull()) << error; - int nv = model->nv; - mjData* data = mj_makeData(model); - - // load keyframe if present, step for a bit - if (model->nkey) mj_resetDataKeyframe(model, data, 0); - while (data->time < 0.1) { - mj_step(model, data); - } - - // minimal call required for mj_jacDot outputs to be valid - mj_kinematics(model, data); - mj_comPos(model, data); - mj_comVel(model, data); - - // get bodyid - int bodyid = mj_name2id(model, mjOBJ_BODY, "query"); - EXPECT_GT(bodyid, 0); - - // get site position - int siteid = mj_name2id(model, mjOBJ_SITE, "query"); - EXPECT_GT(siteid, -1); - mjtNum point[3]; - mju_copy3(point, data->site_xpos+3*siteid); - - // jac, jac_dot - vector jacp(3*nv); - vector jacr(3*nv); - mj_jac(model, data, jacp.data(), jacr.data(), point, bodyid); - vector jacp_dot(3*nv); - vector jacr_dot(3*nv); - mj_jacDot(model, data, jacp_dot.data(), jacr_dot.data(), point, bodyid); - - // jac_h: jacobian after integrating qpos with a timestep of h - constexpr mjtNum h = MjTol(1e-7, 5e-4); - mj_integratePos(model, data->qpos, data->qvel, h); - mj_kinematics(model, data); - mj_comPos(model, data); - vector jacp_h(3*nv); - vector jacr_h(3*nv); - mju_copy3(point, data->site_xpos+3*siteid); // get updated site position - mj_jac(model, data, jacp_h.data(), jacr_h.data(), point, bodyid); - - // jac_dot_h finite-difference approximation - vector jacp_dot_h(3*nv); - mju_sub(jacp_dot_h.data(), jacp_h.data(), jacp.data(), 3*nv); - mju_scl(jacp_dot_h.data(), jacp_dot_h.data(), 1/h, 3*nv); - vector jacr_dot_h(3*nv); - mju_sub(jacr_dot_h.data(), jacr_h.data(), jacr.data(), 3*nv); - mju_scl(jacr_dot_h.data(), jacr_dot_h.data(), 1/h, 3*nv); - - // compare finite-differenced and analytic - mjtNum tol = 1e-5; - EXPECT_THAT(jacp_dot, Pointwise(MjNear(tol, 5e-2), jacp_dot_h)); - EXPECT_THAT(jacr_dot, Pointwise(MjNear(tol, 5e-2), jacr_dot_h)); - - mj_deleteData(data); - mj_deleteModel(model); - } -} - -// compare mj_jacDotSparse with dense mj_jacDot -TEST_F(JacobianTest, JacDotSparse) { - for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) { - char error[1024]; - mjModel* model = LoadModelFromString(xml, error, sizeof(error)); - ASSERT_THAT(model, NotNull()) << error; - int nv = model->nv; - mjData* data = mj_makeData(model); - - // load keyframe if present, step for a bit - if (model->nkey) mj_resetDataKeyframe(model, data, 0); - while (data->time < 0.1) { - mj_step(model, data); - } - - // minimal call required for mj_jacDot outputs to be valid - mj_kinematics(model, data); - mj_comPos(model, data); - mj_comVel(model, data); - - // get bodyid and site position - int bodyid = mj_name2id(model, mjOBJ_BODY, "query"); - EXPECT_GT(bodyid, 0); - int siteid = mj_name2id(model, mjOBJ_SITE, "query"); - EXPECT_GT(siteid, -1); - mjtNum point[3]; - mju_copy3(point, data->site_xpos+3*siteid); - - // dense jacDot - vector jacp_dense(3*nv); - vector jacr_dense(3*nv); - mj_jacDot(model, data, jacp_dense.data(), jacr_dense.data(), point, bodyid); - - // compute body chain using public mjModel fields - vector chain(nv); - int NV = 0; - int weldbody = model->body_weldid[bodyid]; - if (weldbody) { - int da = model->body_dofadr[weldbody] + model->body_dofnum[weldbody] - 1; - while (da >= 0) { - chain[NV++] = da; - da = model->dof_parentid[da]; - } - std::reverse(chain.begin(), chain.begin() + NV); - } - EXPECT_GT(NV, 0); - - // sparse jacDot - vector jacp_sparse(3*NV); - vector jacr_sparse(3*NV); - mj_jacDotSparse(model, data, jacp_sparse.data(), jacr_sparse.data(), - point, bodyid, NV, chain.data()); - - // expand sparse to dense and compare - vector jacp_expanded(3*nv, 0); - vector jacr_expanded(3*nv, 0); - for (int ci = 0; ci < NV; ci++) { - int di = chain[ci]; - for (int r = 0; r < 3; r++) { - jacp_expanded[di+r*nv] = jacp_sparse[ci+r*NV]; - jacr_expanded[di+r*nv] = jacr_sparse[ci+r*NV]; - } - } - - // expect bitwise equality - EXPECT_EQ(jacp_expanded, jacp_dense); - EXPECT_EQ(jacr_expanded, jacr_dense); - - mj_deleteData(data); - mj_deleteModel(model); - } -} using Name2idTest = MujocoTest; diff --git a/test/engine/engine_util_misc_test.cc b/test/engine/engine_util_misc_test.cc index 2acd6125..d1eb8314 100644 --- a/test/engine/engine_util_misc_test.cc +++ b/test/engine/engine_util_misc_test.cc @@ -430,13 +430,90 @@ TEST_F(InterpolationTest, mju_interpolate3D) { expected[0] = quadratic_function_1(sample[0], sample[1], sample[2]); expected[1] = quadratic_function_2(sample[0], sample[1], sample[2]); expected[2] = quadratic_function_3(sample[0], sample[1], sample[2]); - mju_interpolate3D(res, sample, coeff, order); + mju_interpolate3D(res, sample, coeff, order, NULL); EXPECT_NEAR(res[0], expected[0], MjTol(1e-10, 1e-5)); EXPECT_NEAR(res[1], expected[1], MjTol(1e-10, 1e-5)); EXPECT_NEAR(res[2], expected[2], MjTol(1e-10, 1e-5)); } } +TEST_F(InterpolationTest, mju_cellLookup_SingleCell) { + // single cell (1x1x1): local coords should equal global coords + int cellnum[3] = {1, 1, 1}; + mjtNum coord[3] = {0.3, 0.7, 0.5}; + mjtNum local[3]; + int nodeindices[8]; + + int npc = mju_cellLookup(coord, cellnum, 1, local, nodeindices); + EXPECT_EQ(npc, 8); + EXPECT_NEAR(local[0], 0.3, MjTol(1e-12, 1e-6)); + EXPECT_NEAR(local[1], 0.7, MjTol(1e-12, 1e-6)); + EXPECT_NEAR(local[2], 0.5, MjTol(1e-12, 1e-6)); + + // for trilinear 1x1x1: nodes are 0..7 in lexicographic order + for (int i = 0; i < 8; i++) { + EXPECT_EQ(nodeindices[i], i); + } +} + +TEST_F(InterpolationTest, mju_cellLookup_MultiCell) { + // 2x3x4 grid, trilinear: 3x4x5 = 60 nodes + int cellnum[3] = {2, 3, 4}; + int order = 1; + int ny_g = 3*1 + 1; // 4 + int nz_g = 4*1 + 1; // 5 + + // point at (0.75, 0.5, 0.125) -> cell (1, 1, 0) + mjtNum coord[3] = {0.75, 0.5, 0.125}; + mjtNum local[3]; + int nodeindices[8]; + + int npc = mju_cellLookup(coord, cellnum, order, local, nodeindices); + EXPECT_EQ(npc, 8); + + // cell (1,1,0): local = (0.75*2 - 1, 0.5*3 - 1, 0.125*4 - 0) + EXPECT_NEAR(local[0], 0.5, 1e-12); + EXPECT_NEAR(local[1], 0.5, 1e-12); + EXPECT_NEAR(local[2], 0.5, 1e-12); + + // expected node indices for cell (1,1,0), trilinear: + // (gi, gj, gk) for li,lj,lk in {0,1} + // gi = 1+li, gj = 1+lj, gk = 0+lk + // gidx = gi*ny_g*nz_g + gj*nz_g + gk + int expected[8]; + int ni = 0; + for (int li = 0; li <= 1; li++) { + for (int lj = 0; lj <= 1; lj++) { + for (int lk = 0; lk <= 1; lk++) { + expected[ni++] = (1+li)*ny_g*nz_g + (1+lj)*nz_g + lk; + } + } + } + for (int i = 0; i < 8; i++) { + EXPECT_EQ(nodeindices[i], expected[i]); + } +} + +TEST_F(InterpolationTest, mju_cellLookup_Boundary) { + // point exactly at coord=1.0 should clamp to last cell + int cellnum[3] = {3, 3, 3}; + mjtNum coord[3] = {1.0, 1.0, 1.0}; + mjtNum local[3]; + + mju_cellLookup(coord, cellnum, 1, local, NULL); + // cell (2,2,2), local = (1*3 - 2, 1*3 - 2, 1*3 - 2) = (1, 1, 1) + EXPECT_NEAR(local[0], 1.0, 1e-12); + EXPECT_NEAR(local[1], 1.0, 1e-12); + EXPECT_NEAR(local[2], 1.0, 1e-12); + + // point at coord=0.0 should map to first cell + mjtNum coord0[3] = {0.0, 0.0, 0.0}; + mju_cellLookup(coord0, cellnum, 1, local, NULL); + EXPECT_NEAR(local[0], 0.0, 1e-12); + EXPECT_NEAR(local[1], 0.0, 1e-12); + EXPECT_NEAR(local[2], 0.0, 1e-12); +} + TEST_F(InterpolationTest, mju_defGradient) { int order = 1; mjtNum mat[9]; @@ -521,7 +598,48 @@ TEST_F(InterpolationTest, mju_defGradient) { EXPECT_THAT(mat, Pointwise(MjNear(1e-8, 1e-6), rot7)); } -// --------------------------------- Base64 ------------------------------------ +TEST_F(InterpolationTest, mju_flexInterpState_MultiCell) { + int order = 1; // trilinear + int cy = 2; + int cz = 2; + int nodenum = 27; // 3x3x3 + + std::vector xpos(3 * nodenum); + mjtNum quat[4]; + + // Populate xpos directly for a grid centered at origin, rotated 90 deg around + // Z Original grid points: {-0.1, 0.0, 0.1}^3 Rotated: (x, y, z) -> (-y, x, z) + int idx = 0; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + for (int k = 0; k < 3; k++) { + mjtNum x = (i - 1) * 0.1; + mjtNum y = (j - 1) * 0.1; + mjtNum z = (k - 1) * 0.1; + + // Apply rotation + xpos[3*idx + 0] = -y; + xpos[3*idx + 1] = x; + xpos[3*idx + 2] = z; + idx++; + } + } + } + + int npc = (order+1)*(order+1)*(order+1); + std::vector xpos_c(3 * npc); + + mju_flexGatherCellState(order, cy, cz, 0, 0, 0, xpos.data(), NULL, NULL, + xpos_c.data(), NULL, NULL, NULL, quat); + + // Expected quaternion for -90 deg around Z (global to local): + // [sqrt(0.5), 0, 0, -sqrt(0.5)] + mjtNum expected_val = mju_sqrt(0.5); + EXPECT_NEAR(quat[0], expected_val, 1e-5); + EXPECT_NEAR(quat[1], 0.0, 1e-5); + EXPECT_NEAR(quat[2], 0.0, 1e-5); + EXPECT_NEAR(quat[3], -expected_val, 1e-5); +} using Base64Test = MujocoTest; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 94d15988..d7d4c47d 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -1200,6 +1200,7 @@ public unsafe struct mjModel_ { public int* flex_matid; public int* flex_group; public int* flex_interp; + public int* flex_cellnum; public int* flex_nodeadr; public int* flex_nodenum; public int* flex_vertadr; diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 3fcacd7d..80fb63d5 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -2446,6 +2446,15 @@ struct MjsFlex { void set_elastic2d(int value) { ptr_->elastic2d = value; } + emscripten::val cellcount() const { + return emscripten::val(emscripten::typed_memory_view(3, ptr_->cellcount)); + } + int order() const { + return ptr_->order; + } + void set_order(int value) { + ptr_->order = value; + } mjStringVec &nodebody() const { return *(ptr_->nodebody); } @@ -4640,6 +4649,9 @@ struct MjModel { emscripten::val flex_interp() const { return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_interp)); } + emscripten::val flex_cellnum() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->nflex * 3, ptr_->flex_cellnum)); + } emscripten::val flex_nodeadr() const { return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_nodeadr)); } @@ -11806,6 +11818,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("flex_bending", &MjModel::flex_bending) .property("flex_bvhadr", &MjModel::flex_bvhadr) .property("flex_bvhnum", &MjModel::flex_bvhnum) + .property("flex_cellnum", &MjModel::flex_cellnum) .property("flex_centered", &MjModel::flex_centered) .property("flex_conaffinity", &MjModel::flex_conaffinity) .property("flex_condim", &MjModel::flex_condim) @@ -12569,6 +12582,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("info", &MjsExclude::info, &MjsExclude::set_info, reference()); emscripten::class_("MjsFlex") .property("activelayers", &MjsFlex::activelayers, &MjsFlex::set_activelayers, reference()) + .property("cellcount", &MjsFlex::cellcount) .property("conaffinity", &MjsFlex::conaffinity, &MjsFlex::set_conaffinity, reference()) .property("condim", &MjsFlex::condim, &MjsFlex::set_condim, reference()) .property("contype", &MjsFlex::contype, &MjsFlex::set_contype, reference()) @@ -12590,6 +12604,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("material", &MjsFlex::material, &MjsFlex::set_material, reference()) .property("node", &MjsFlex::node, reference()) .property("nodebody", &MjsFlex::nodebody, reference()) + .property("order", &MjsFlex::order, &MjsFlex::set_order, reference()) .property("passive", &MjsFlex::passive, &MjsFlex::set_passive, reference()) .property("poisson", &MjsFlex::poisson, &MjsFlex::set_poisson, reference()) .property("priority", &MjsFlex::priority, &MjsFlex::set_priority, reference())