From 55c7f45dc5c30ff39ac38e7166835912db08f2d3 Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Wed, 24 May 2023 11:01:58 -0700
Subject: [PATCH] Fix sparse matmul logic in `addJTBJSparse`.

Also rename arguments and loop variables in `addJTBJSparse` for additional clarity.

PiperOrigin-RevId: 534908360
Change-Id: I143ef5aef7ff90091369d408740f7f44d9e19378
---
 doc/changelog.rst              |  3 +++
 src/engine/engine_derivative.c | 38 ++++++++++++++++++----------------
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index 32fb7a1a..8844f301 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -17,6 +17,9 @@ Bug fixes
   invalid. In such case, now MuJoCo only accepts a non-watertight geometry if :ref:`shellinertia<body-geom-shellinertia>`
   is equal to ``true``.
 
+- Fixed the sparse Jacobian multiplication logic that is used to compute derivatives for tendon damping and fluid force,
+  which affects the behaviour of the :ref:`implicit and implicitfast integrators<geIntegration>`.
+
 Plugins
 ^^^^^^^
 
diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c
index 48ee4ab1..7e5833c3 100644
--- a/src/engine/engine_derivative.c
+++ b/src/engine/engine_derivative.c
@@ -708,39 +708,41 @@ static void addJTBJ(const mjModel* m, mjData* d, const mjtNum* J, const mjtNum*
 
 
 // add J'*B*J to qDeriv, sparse version
-static void addJTBJSparse(const mjModel* m, mjData* d, const mjtNum* J,
-  const mjtNum* B, int n, int offset,
-  const int* rownnz, const int* rowadr, const int* colind) {
+static void addJTBJSparse(
+    const mjModel* m, mjData* d, const mjtNum* J,
+    const mjtNum* B, int n, int offset,
+    const int* J_rownnz, const int* J_rowadr, const int* J_colind) {
   int nv = m->nv;
 
   // allocate row
   mjMARKSTACK;
   mjtNum* row = mj_stackAlloc(d, nv);
-  mju_zero(row, nv);
 
-  // process non-zero elements of B
-  for (int i=0; i<n; i++) {
-    for (int j=0; j<n; j++) {
+  // compute qDeriv(k,p) += sum_{i,j} ( J(i,k)*B(i,j)*J(j,p) )
+  for (int i = 0; i < n; i++) {
+    for (int j = 0; j < n; j++) {
       if (!B[i*n+j]) {
         continue;
       }
 
-      // process non-zero elements of J(i,k)
-      for (int k=0; k<rownnz[offset+i]; k++) {
-        int ik = rowadr[offset+i] + k;
+      mju_zero(row, nv);
+
+      // loop over non-zero elements of J(i,:)
+      for (int k = 0; k < J_rownnz[offset+i]; k++) {
+        int ik = J_rowadr[offset+i] + k;
         mjtNum scl = J[ik]*B[i*n+j];
 
-        // process non-zero elements of J(j,p)
-        for (int p=0; p<rownnz[offset+j]; p++) {
-          int jp = rowadr[offset+j] + p;
-
-          // row[p] = J(i,k)*B(i,j)*J(j,p)
-          row[p] = scl * J[jp];
+        // loop over non-zero elements of J(j,:)
+        // (pJ is the sparse column index into J)
+        for (int pJ = 0; pJ < J_rownnz[offset+j]; pJ++) {
+          int adr = J_rowadr[offset+j] + pJ;
+          row[J_colind[adr]] = scl * J[adr];
         }
 
         // add row to qDeriv(k,:)
-        for (int s=0; s<d->D_rownnz[k]; s++) {
-          int adr = d->D_rowadr[k] + s;
+        // (pD is the sparse column index into qDeriv)
+        for (int pD = 0; pD < d->D_rownnz[k]; pD++) {
+          int adr = d->D_rowadr[k] + pD;
           d->qDeriv[adr] += row[d->D_colind[adr]];
         }
       }