Add mjData.nA the number of non-zeros in the constraint inverse inertia matrix.

PiperOrigin-RevId: 714076988
Change-Id: I4c323a6e81dfb23f7972b662719bbcb688bacf4a

doc/includes/references.h

@@ -170,6 +170,7 @@ struct mjData_ {
   int     nl;                // number of limit constraints
   int     nefc;              // number of constraints
   int     nJ;                // number of non-zeros in constraint Jacobian
+  int     nA;                // number of non-zeros in constraint inverse inertia matrix
   int     nisland;           // number of detected constraint islands
 
   // global properties
@@ -391,8 +392,8 @@ struct mjData_ {
   // computed by mj_projectConstraint (PGS solver)
   int*    efc_AR_rownnz;     // number of non-zeros in AR                        (nefc x 1)
   int*    efc_AR_rowadr;     // row start address in colind array                (nefc x 1)
-  int*    efc_AR_colind;     // column indices in sparse AR                      (nefc x nefc)
+  int*    efc_AR_colind;     // column indices in sparse AR                      (nA x 1)
-  mjtNum* efc_AR;            // J*inv(M)*J' + R                                  (nefc x nefc)
+  mjtNum* efc_AR;            // J*inv(M)*J' + R                                  (nA x 1)
 
   //-------------------- arena-allocated: POSITION, VELOCITY dependent
 

include/mujoco/mjdata.h

@@ -198,6 +198,7 @@ struct mjData_ {
   int     nl;                // number of limit constraints
   int     nefc;              // number of constraints
   int     nJ;                // number of non-zeros in constraint Jacobian
+  int     nA;                // number of non-zeros in constraint inverse inertia matrix
   int     nisland;           // number of detected constraint islands
 
   // global properties
@@ -419,8 +420,8 @@ struct mjData_ {
   // computed by mj_projectConstraint (PGS solver)
   int*    efc_AR_rownnz;     // number of non-zeros in AR                        (nefc x 1)
   int*    efc_AR_rowadr;     // row start address in colind array                (nefc x 1)
-  int*    efc_AR_colind;     // column indices in sparse AR                      (nefc x nefc)
+  int*    efc_AR_colind;     // column indices in sparse AR                      (nA x 1)
-  mjtNum* efc_AR;            // J*inv(M)*J' + R                                  (nefc x nefc)
+  mjtNum* efc_AR;            // J*inv(M)*J' + R                                  (nA x 1)
 
   //-------------------- arena-allocated: POSITION, VELOCITY dependent
 

include/mujoco/mjxmacro.h

@@ -715,8 +715,8 @@
 #define MJDATA_ARENA_POINTERS_DUAL                           \
     X( int,      efc_AR_rownnz,     MJ_D(nefc), 1          ) \
     X( int,      efc_AR_rowadr,     MJ_D(nefc), 1          ) \
-    X( int,      efc_AR_colind,     MJ_D(nefc), MJ_D(nefc) ) \
+    X( int,      efc_AR_colind,     MJ_D(nA),   1          ) \
-    X( mjtNum,   efc_AR,            MJ_D(nefc), MJ_D(nefc) )
+    X( mjtNum,   efc_AR,            MJ_D(nA),   1          )
 
 // array fields of mjData that are used for constraint islands
 #define MJDATA_ARENA_POINTERS_ISLAND                 \
@@ -757,6 +757,7 @@
     X( int,       nl                 ) \
     X( int,       nefc               ) \
     X( int,       nJ                 ) \
+    X( int,       nA                 ) \
     X( int,       nisland            ) \
     X( mjtNum,    time               ) \
     X( uintptr_t, threadpool         )

introspect/structs.py

@@ -4701,6 +4701,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of non-zeros in constraint Jacobian',
              ),
+             StructFieldDecl(
+                 name='nA',
+                 type=ValueType(name='int'),
+                 doc='number of non-zeros in constraint inverse inertia matrix',
+             ),
              StructFieldDecl(
                  name='nisland',
                  type=ValueType(name='int'),
@@ -5797,7 +5802,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                      inner_type=ValueType(name='int'),
                  ),
                  doc='column indices in sparse AR',
-                 array_extent=('nefc', 'nefc'),
+                 array_extent=('nA',),
              ),
              StructFieldDecl(
                  name='efc_AR',
@@ -5805,7 +5810,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc="J*inv(M)*J' + R",
-                 array_extent=('nefc', 'nefc'),
+                 array_extent=('nA',),
              ),
              StructFieldDecl(
                  name='efc_vel',

python/mujoco/structs.cc

@@ -748,6 +748,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
   X(ne);
   X(nf);
   X(nJ);
+  X(nA);
   X(nefc);
   X(nisland);
   X(time);
@@ -825,6 +826,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
   X(ne);
   X(nf);
   X(nJ);
+  X(nA);
   X(nefc);
   X(nisland);
   X(time);

src/engine/engine_core_constraint.c

@@ -1943,7 +1943,7 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
 // driver: call all functions above
 void mj_makeConstraint(const mjModel* m, mjData* d) {
   // clear sizes
-  d->ne = d->nf = d->nl = d->nefc = d->nJ = 0;
+  d->ne = d->nf = d->nl = d->nefc = d->nJ = d->nA = 0;
 
   // disabled or Jacobian not allocated: return
   if (mjDISABLED(mjDSBL_CONSTRAINT)) {
@@ -1960,6 +1960,7 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
     d->nJ = nefc_allocated * m->nv;
   }
   d->nefc = nefc_allocated;
+  d->nA = d->nefc * d->nefc;
 
   // allocate efc arrays on arena
   if (!arenaAllocEfc(m, d)) {

src/engine/engine_io.c

@@ -1893,6 +1893,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
   d->nl = 0;
   d->nefc = 0;
   d->nJ = 0;
+  d->nA = 0;
   d->nisland = 0;
 
   // clear global properties

test/engine/engine_io_test.cc

@@ -157,6 +157,7 @@ TEST_F(EngineIoTest, ResetVariableSizes) {
   EXPECT_EQ(data->nf, 0);
   EXPECT_EQ(data->nefc, 0);
   EXPECT_EQ(data->nJ, 0);
+  EXPECT_EQ(data->nA, 0);
   EXPECT_EQ(data->ncon, 0);
 
   mj_deleteData(data);

unity/Runtime/Bindings/MjBindings.cs

@@ -4852,6 +4852,7 @@ public unsafe struct mjData_ {
   public int nl;
   public int nefc;
   public int nJ;
+  public int nA;
   public int nisland;
   public double time;
   public fixed double energy[2];