Convert qLD to CSR format.

PiperOrigin-RevId: 723955038
Change-Id: I30c3dc7f59739e89ae5fff8841432bc74717ec1b
This commit is contained in:
Yuval Tassa
2025-02-06 08:58:45 -08:00
committed by Copybara-Service
parent cb1696eb34
commit c27d3758c2
22 changed files with 151 additions and 153 deletions
+1 -1
View File
@@ -368,7 +368,7 @@ def make_data(
'efc_aref': (nefc, float),
'efc_force': (nefc, float),
'_qM_sparse': (m.nM, float),
'_qLD_sparse': (m.nM, float),
'_qLD_sparse': (m.nC, float),
'_qLDiagInv_sparse': (m.nv, float),
}
+4 -1
View File
@@ -91,7 +91,10 @@ class SmoothTest(absltest.TestCase):
_assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
# factor_m
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qLD')
qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name
for i in range(m.nC):
qLDLegacy[d.mapM2C[i]] = d.qLD[i]
_assert_eq(qLDLegacy, dx.qLD, 'qLD')
_assert_attr_eq(d, dx, 'qLDiagInv')
_assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
_assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')
+1 -1
View File
@@ -1348,7 +1348,7 @@ class Data(PyTreeNode):
efc_aref: reference pseudo-acceleration (nefc,)
efc_force: constraint force in constraint space (nefc,)
_qM_sparse: qM in sparse representation (nM,)
_qLD_sparse: qLD in sparse representation (nM,)
_qLD_sparse: qLD in sparse representation (nC,)
_qLDiagInv_sparse: qLDiagInv in sparse representation (nv,)
""" # fmt: skip
# constant sizes: