Implement sparse Jacobian time derivative

PiperOrigin-RevId: 897609541
Change-Id: Ic86b6026cfc369c584d3edfe12a9a2fd14cc4357

src/engine/engine_core_util.c

@@ -660,6 +660,92 @@ void mj_jacDot(const mjModel* m, const mjData* d,
 }
 
 
+// compute 3/6-by-NV sparse Jacobian time derivative of global point attached to given body
+void mj_jacDotSparse(const mjModel* m, const mjData* d,
+                     mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
+                     int NV, const int* chain) {
+  mjtNum offset[3];
+  mjtNum pvel[6];
+
+  // clear jacobians, compute offset and pvel if required
+  if (jacp) {
+    mju_zero(jacp, 3*NV);
+    const mjtNum* com = d->subtree_com+3*m->body_rootid[body];
+    mju_sub3(offset, point, com);
+    mju_transformSpatial(pvel, d->cvel+6*body, 0, point, com, 0);
+  }
+  if (jacr) {
+    mju_zero(jacr, 3*NV);
+  }
+
+  // skip fixed bodies
+  body = m->body_weldid[body];
+
+  // no movable body found: nothing to do
+  if (!body) {
+    return;
+  }
+
+  // get last dof that affects this body
+  int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
+
+  // start at end of chain (chain is in increasing order)
+  int ci = NV-1;
+
+  // backward pass over dof ancestor chain
+  while (da >= 0) {
+    // find chain index for this dof
+    while (ci >= 0 && chain[ci] > da) {
+      ci--;
+    }
+
+    // dof not in chain: SHOULD NOT OCCUR
+    if (ci < 0 || chain[ci] != da) {
+      mjERROR("dof index %d not found in chain", da);
+    }
+
+    mjtNum cdof_dot[6];
+    mji_copy6(cdof_dot, d->cdof_dot+6*da);
+    mjtNum* cdof = d->cdof+6*da;
+
+    // check for quaternion
+    mjtJoint type = m->jnt_type[m->dof_jntid[da]];
+    int dofadr = m->jnt_dofadr[m->dof_jntid[da]];
+    int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && da >= dofadr + 3);
+
+    // compute cdof_dot for quaternion (use current body cvel)
+    if (is_quat) {
+      mji_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[da], cdof);
+    }
+
+    // construct rotation jacobian
+    if (jacr) {
+      jacr[ci+0*NV] += cdof_dot[0];
+      jacr[ci+1*NV] += cdof_dot[1];
+      jacr[ci+2*NV] += cdof_dot[2];
+    }
+
+    // construct translation jacobian (correct for rotation)
+    if (jacp) {
+      // first correction term, account for varying cdof
+      mjtNum tmp1[3];
+      mji_cross(tmp1, cdof_dot, offset);
+
+      // second correction term, account for point translational velocity
+      mjtNum tmp2[3];
+      mji_cross(tmp2, cdof, pvel + 3);
+
+      jacp[ci+0*NV] += cdof_dot[3] + tmp1[0] + tmp2[0];
+      jacp[ci+1*NV] += cdof_dot[4] + tmp1[1] + tmp2[1];
+      jacp[ci+2*NV] += cdof_dot[5] + tmp1[2] + tmp2[2];
+    }
+
+    // advance to parent dof
+    da = m->dof_parentid[da];
+  }
+}
+
+
 // compute subtree angular momentum matrix
 void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
   int nv = m->nv;