Allocate qM CSR matrix structure on mjData (currently empty).
PiperOrigin-RevId: 733247419 Change-Id: I814e574e3161875e89433a6e03f2c9e200d8b1d6
This commit is contained in:
committed by
Copybara-Service
parent
8f37ae1bf2
commit
c72c83aa3b
@@ -312,6 +312,10 @@ struct mjData_ {
|
|||||||
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
|
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
|
||||||
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
|
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
|
||||||
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
|
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
|
||||||
|
int* M_rownnz; // inertia: non-zeros in each row (nv x 1)
|
||||||
|
int* M_rowadr; // inertia: address of each row in M_colind (nv x 1)
|
||||||
|
int* M_colind; // inertia: column indices of non-zeros (nM x 1)
|
||||||
|
int* mapM2M; // index mapping from M (legacy) to M (CSR) (nM x 1)
|
||||||
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
|
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
|
||||||
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
|
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
|
||||||
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
|
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
|
||||||
|
|||||||
@@ -340,6 +340,10 @@ struct mjData_ {
|
|||||||
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
|
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
|
||||||
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
|
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
|
||||||
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
|
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
|
||||||
|
int* M_rownnz; // inertia: non-zeros in each row (nv x 1)
|
||||||
|
int* M_rowadr; // inertia: address of each row in M_colind (nv x 1)
|
||||||
|
int* M_colind; // inertia: column indices of non-zeros (nM x 1)
|
||||||
|
int* mapM2M; // index mapping from M (legacy) to M (CSR) (nM x 1)
|
||||||
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
|
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
|
||||||
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
|
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
|
||||||
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
|
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
|
||||||
|
|||||||
@@ -672,6 +672,10 @@
|
|||||||
X ( int, B_rownnz, nbody, 1 ) \
|
X ( int, B_rownnz, nbody, 1 ) \
|
||||||
X ( int, B_rowadr, nbody, 1 ) \
|
X ( int, B_rowadr, nbody, 1 ) \
|
||||||
X ( int, B_colind, nB, 1 ) \
|
X ( int, B_colind, nB, 1 ) \
|
||||||
|
X ( int, M_rownnz, nv, 1 ) \
|
||||||
|
X ( int, M_rowadr, nv, 1 ) \
|
||||||
|
X ( int, M_colind, nM, 1 ) \
|
||||||
|
X ( int, mapM2M, nM, 1 ) \
|
||||||
X ( int, C_rownnz, nv, 1 ) \
|
X ( int, C_rownnz, nv, 1 ) \
|
||||||
X ( int, C_rowadr, nv, 1 ) \
|
X ( int, C_rowadr, nv, 1 ) \
|
||||||
X ( int, C_colind, nC, 1 ) \
|
X ( int, C_colind, nC, 1 ) \
|
||||||
|
|||||||
@@ -339,6 +339,10 @@ def make_data(
|
|||||||
'B_rownnz': (m.nbody, jp.int32),
|
'B_rownnz': (m.nbody, jp.int32),
|
||||||
'B_rowadr': (m.nbody, jp.int32),
|
'B_rowadr': (m.nbody, jp.int32),
|
||||||
'B_colind': (m.nB, jp.int32),
|
'B_colind': (m.nB, jp.int32),
|
||||||
|
'M_rownnz': (m.nv, jp.int32),
|
||||||
|
'M_rowadr': (m.nv, jp.int32),
|
||||||
|
'M_colind': (m.nM, jp.int32),
|
||||||
|
'mapM2M': (m.nM, jp.int32),
|
||||||
'C_rownnz': (m.nv, jp.int32),
|
'C_rownnz': (m.nv, jp.int32),
|
||||||
'C_rowadr': (m.nv, jp.int32),
|
'C_rowadr': (m.nv, jp.int32),
|
||||||
'C_colind': (m.nC, jp.int32),
|
'C_colind': (m.nC, jp.int32),
|
||||||
|
|||||||
@@ -1316,6 +1316,10 @@ class Data(PyTreeNode):
|
|||||||
B_rownnz: body-dof: non-zeros in each row (nbody,)
|
B_rownnz: body-dof: non-zeros in each row (nbody,)
|
||||||
B_rowadr: body-dof: address of each row in B_colind (nbody,)
|
B_rowadr: body-dof: address of each row in B_colind (nbody,)
|
||||||
B_colind: body-dof: column indices of non-zeros (nB,)
|
B_colind: body-dof: column indices of non-zeros (nB,)
|
||||||
|
M_rownnz: inertia: non-zeros in each row (nv,)
|
||||||
|
M_rowadr: inertia: address of each row in M_colind (nv,)
|
||||||
|
M_colind: inertia: column indices of non-zeros (nM,)
|
||||||
|
mapM2M: index mapping from M (legacy) to M (CSR) (nM,)
|
||||||
C_rownnz: reduced dof-dof: non-zeros in each row (nv,)
|
C_rownnz: reduced dof-dof: non-zeros in each row (nv,)
|
||||||
C_rowadr: reduced dof-dof: address of each row in C_colind (nv,)
|
C_rowadr: reduced dof-dof: address of each row in C_colind (nv,)
|
||||||
C_colind: reduced dof-dof: column indices of non-zeros (nC,)
|
C_colind: reduced dof-dof: column indices of non-zeros (nC,)
|
||||||
@@ -1445,6 +1449,10 @@ class Data(PyTreeNode):
|
|||||||
B_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
B_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
B_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
B_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
B_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
B_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
|
M_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
|
M_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
|
M_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
|
mapM2M: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
C_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
C_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
C_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
C_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
C_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
C_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
|
||||||
|
|||||||
@@ -5526,6 +5526,38 @@ STRUCTS: Mapping[str, StructDecl] = dict([
|
|||||||
doc='body-dof: column indices of non-zeros',
|
doc='body-dof: column indices of non-zeros',
|
||||||
array_extent=('nB',),
|
array_extent=('nB',),
|
||||||
),
|
),
|
||||||
|
StructFieldDecl(
|
||||||
|
name='M_rownnz',
|
||||||
|
type=PointerType(
|
||||||
|
inner_type=ValueType(name='int'),
|
||||||
|
),
|
||||||
|
doc='inertia: non-zeros in each row',
|
||||||
|
array_extent=('nv',),
|
||||||
|
),
|
||||||
|
StructFieldDecl(
|
||||||
|
name='M_rowadr',
|
||||||
|
type=PointerType(
|
||||||
|
inner_type=ValueType(name='int'),
|
||||||
|
),
|
||||||
|
doc='inertia: address of each row in M_colind',
|
||||||
|
array_extent=('nv',),
|
||||||
|
),
|
||||||
|
StructFieldDecl(
|
||||||
|
name='M_colind',
|
||||||
|
type=PointerType(
|
||||||
|
inner_type=ValueType(name='int'),
|
||||||
|
),
|
||||||
|
doc='inertia: column indices of non-zeros',
|
||||||
|
array_extent=('nM',),
|
||||||
|
),
|
||||||
|
StructFieldDecl(
|
||||||
|
name='mapM2M',
|
||||||
|
type=PointerType(
|
||||||
|
inner_type=ValueType(name='int'),
|
||||||
|
),
|
||||||
|
doc='index mapping from M (legacy) to M (CSR)',
|
||||||
|
array_extent=('nM',),
|
||||||
|
),
|
||||||
StructFieldDecl(
|
StructFieldDecl(
|
||||||
name='C_rownnz',
|
name='C_rownnz',
|
||||||
type=PointerType(
|
type=PointerType(
|
||||||
|
|||||||
@@ -4942,6 +4942,10 @@ public unsafe struct mjData_ {
|
|||||||
public int* B_rownnz;
|
public int* B_rownnz;
|
||||||
public int* B_rowadr;
|
public int* B_rowadr;
|
||||||
public int* B_colind;
|
public int* B_colind;
|
||||||
|
public int* M_rownnz;
|
||||||
|
public int* M_rowadr;
|
||||||
|
public int* M_colind;
|
||||||
|
public int* mapM2M;
|
||||||
public int* C_rownnz;
|
public int* C_rownnz;
|
||||||
public int* C_rowadr;
|
public int* C_rowadr;
|
||||||
public int* C_colind;
|
public int* C_colind;
|
||||||
|
|||||||
Reference in New Issue
Block a user