diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 48419452..0d564960 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -2748,7 +2748,7 @@ mju_cholSolve .. mujoco-include:: mju_cholSolve -Solve mat * res = vec, where mat is Cholesky-factorized +Solve (mat*mat') * res = vec, where mat is a Cholesky factor. .. _mju_cholUpdate: @@ -2759,6 +2759,101 @@ mju_cholUpdate Cholesky rank-one update: L*L' +/- x*x'; return rank. +.. _mju_cholFactorBand: + +mju_cholFactorBand +~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_cholFactorBand + +Band-dense Cholesky decomposition. +|br| Add ``diagadd + diagmul*mat_ii`` to diagonal before decomposition. +|br| Returns the minimum value of the factorized diagonal or 0 if rank-defficient. + + **Symmetric band-dense matrices** + + :ref:`mju_cholFactorBand` and subsequent functions containing the substring "band" operate on matrices which are a + generalization of symmetric `band matrices `_. *Symmetric band-dense* or + "arrowhead" matrices have non-zeros along proximal diagonal bands and dense blocks on the bottom rows and right + columns. These matrices have the property that Cholesky factorization creates no fill-in and can therefore be + performed efficiently in-place. Matrix structure is defined by three integers: + + - ``ntotal``: the number of rows (columns) of the symmetric matrix. + - ``nband``: the number of bands under (over) the diagonal, inclusive of the diagonal. + - ``ndense``: the number of dense rows (columns) at the bottom (right). + + The non-zeros are stored in memory as two contiguous row-major blocks, colored green and blue in the illustration + below. The first block has size ``nband x (ntotal-ndense)`` and contains the diagonal and the bands below it. The + second block has size ``ndense x ntotal`` and contains the dense part. Total required memory is the sum of the block + sizes. + + .. figure:: /images/APIreference/arrowhead.svg + :width: 750px + :align: left + + For example, consider an arrowhead matrix with ``nband = 3``, ``ndense = 2`` and ``ntotal = 8``. In this example, the + total memory required is ``3*(8-2) + 2*8 = 34`` mjtNum's, laid out as follows: + + .. code-block:: + + 0 1 2 + 3 4 5 + 6 7 8 + 9 10 11 + 12 13 14 + 15 16 17 + 18 19 20 21 22 23 24 25 + 26 27 28 29 30 31 32 33 + + + The diagonal elements are ``2, 5, 8, 11, 14, 17, 24, 33``. + |br| Elements ``0, 1, 3, 25`` are present in memory but never touched. + +.. _mju_cholSolveBand: + +mju_cholSolveBand +~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_cholSolveBand + +Solve (mat*mat')*res = vec where mat is a band-dense Cholesky factor. + +.. _mju_band2Dense: + +mju_band2Dense +~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_band2Dense + +Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0. + +.. _mju_dense2Band: + +mju_dense2Band +~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_dense2Band + +Convert dense matrix to banded matrix. + +.. _mju_bandMulMatVec: + +mju_bandMulMatVec +~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_bandMulMatVec + +Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0. + +.. _mju_bandDiag: + +mju_bandDiag +~~~~~~~~~~~~ + +.. mujoco-include:: mju_bandDiag + +Address of diagonal element i in band-dense matrix representation. + .. _mju_eig3: mju_eig3 diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 40d6fcd6..3ad43149 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -321,6 +321,51 @@ mju_ceil .. _Decompositions: +.. _mju_cholFactorBand: + +Band-dense Cholesky decomposition. +|br| Add ``diagadd + diagmul*mat_ii`` to diagonal before decomposition. +|br| Returns the minimum value of the factorized diagonal or 0 if rank-defficient. + + **Symmetric band-dense matrices** + + :ref:`mju_cholFactorBand` and subsequent functions containing the substring "band" operate on matrices which are a + generalization of symmetric `band matrices `_. *Symmetric band-dense* or + "arrowhead" matrices have non-zeros along proximal diagonal bands and dense blocks on the bottom rows and right + columns. These matrices have the property that Cholesky factorization creates no fill-in and can therefore be + performed efficiently in-place. Matrix structure is defined by three integers: + + - ``ntotal``: the number of rows (columns) of the symmetric matrix. + - ``nband``: the number of bands under (over) the diagonal, inclusive of the diagonal. + - ``ndense``: the number of dense rows (columns) at the bottom (right). + + The non-zeros are stored in memory as two contiguous row-major blocks, colored green and blue in the illustration + below. The first block has size ``nband x (ntotal-ndense)`` and contains the diagonal and the bands below it. The + second block has size ``ndense x ntotal`` and contains the dense part. Total required memory is the sum of the block + sizes. + + .. figure:: /images/APIreference/arrowhead.svg + :width: 750px + :align: left + + For example, consider an arrowhead matrix with ``nband = 3``, ``ndense = 2`` and ``ntotal = 8``. In this example, the + total memory required is ``3*(8-2) + 2*8 = 34`` mjtNum's, laid out as follows: + + .. code-block:: + + 0 1 2 + 3 4 5 + 6 7 8 + 9 10 11 + 12 13 14 + 15 16 17 + 18 19 20 21 22 23 24 25 + 26 27 28 29 30 31 32 33 + + + The diagonal elements are ``2, 5, 8, 11, 14, 17, 24, 33``. + |br| Elements ``0, 1, 3, 25`` are present in memory but never touched. + .. _mju_boxQP: Minimize :math:`\tfrac{1}{2} x^T H x + x^T g \quad \text{s.t.} \quad l \le x \le u`, return rank or -1 if failed. diff --git a/doc/changelog.rst b/doc/changelog.rst index 5fb33741..adf5e2ac 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -5,18 +5,26 @@ Changelog Upcoming version (not yet released) ----------------------------------- -Plugins -^^^^^^^ - .. youtube:: hqIMTNGaLF4 :align: right :width: 240px +Plugins +^^^^^^^ + - Added touch-grid sensor plugin. See `documentation `_ for details, and associated `touch_grid.xml `_ example model. The plugin includes `in-scene visualisation `_. -- Add ``mj_multiRay`` function for intersecting multiple rays emanating from a single point. This is significantly - faster than calling ``mj_ray`` multiple times. + +General +^^^^^^^ + +- Added :ref:`mjd_inverseFD` for finite-differenced inverse-dynamics derivatives. +- Added functions for operations on banded-then-dense "arrowhead" matrices. Such matrices are + common when doing direct trajectory optimization. + See :ref:`mju_cholFactorBand` documentation for details. +- Added :ref:`mj_multiRay` function for intersecting multiple rays emanating from a single point. + This is significantly faster than calling :ref:`mj_ray` multiple times. Version 2.3.5 (April 25, 2023) ------------------------------ diff --git a/doc/images/APIreference/arrowhead.svg b/doc/images/APIreference/arrowhead.svg new file mode 100644 index 00000000..6fa78ec0 --- /dev/null +++ b/doc/images/APIreference/arrowhead.svg @@ -0,0 +1 @@ + diff --git a/doc/includes/references.h b/doc/includes/references.h index 48ca2f6b..decaea18 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -2408,6 +2408,16 @@ void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[4], int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag); void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n); int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus); +mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense, + mjtNum diagadd, mjtNum diagmul); +void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec, + int ntotal, int nband, int ndense); +void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense, + mjtByte flg_sym); +void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense); +void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, + int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym); +int mju_bandDiag(int i, int ntotal, int nband, int ndense); int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]); int mju_boxQP(mjtNum* res, mjtNum* R, int* index, const mjtNum* H, const mjtNum* g, int n, const mjtNum* lower, const mjtNum* upper); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 6d1d384f..4180fd17 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1068,12 +1068,39 @@ MJAPI void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[ // Cholesky decomposition: mat = L*L'; return rank, decomposition performed in-place into mat. MJAPI int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag); -// Solve mat * res = vec, where mat is Cholesky-factorized +// Solve (mat*mat') * res = vec, where mat is a Cholesky factor. MJAPI void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n); // Cholesky rank-one update: L*L' +/- x*x'; return rank. MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus); +// Band-dense Cholesky decomposition. +// Returns minimum value in the factorized diagonal, or 0 if rank-deficient. +// mat has (ntotal-ndense) x nband + ndense x ntotal elements. +// The first (ntotal-ndense) x nband store the band part, left of diagonal, inclusive. +// The second ndense x ntotal store the band part as entire dense rows. +// Add diagadd+diagmul*mat_ii to diagonal before factorization. +MJAPI mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense, + mjtNum diagadd, mjtNum diagmul); + +// Solve (mat*mat')*res = vec where mat is a band-dense Cholesky factor. +MJAPI void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec, + int ntotal, int nband, int ndense); + +// Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0. +MJAPI void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense, + mjtByte flg_sym); + +// Convert dense matrix to banded matrix. +MJAPI void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense); + +// Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0. +MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, + int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym); + +// Address of diagonal element i in band-dense matrix representation. +MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense); + // Eigenvalue decomposition of symmetric 3x3 matrix. MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]); diff --git a/introspect/functions.py b/introspect/functions.py index fcdcd020..9c0932f3 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -6815,7 +6815,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ type=ValueType(name='int'), ), ), - doc='Solve mat * res = vec, where mat is Cholesky-factorized', + doc="Solve (mat*mat') * res = vec, where mat is a Cholesky factor.", )), ('mju_cholUpdate', FunctionDecl( @@ -6845,6 +6845,216 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc="Cholesky rank-one update: L*L' +/- x*x'; return rank.", )), + ('mju_cholFactorBand', + FunctionDecl( + name='mju_cholFactorBand', + return_type=ValueType(name='mjtNum'), + parameters=( + FunctionParameterDecl( + name='mat', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='ntotal', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='nband', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='ndense', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='diagadd', + type=ValueType(name='mjtNum'), + ), + FunctionParameterDecl( + name='diagmul', + type=ValueType(name='mjtNum'), + ), + ), + doc='Band-dense Cholesky decomposition. Returns minimum value in the factorized diagonal, or 0 if rank-deficient. mat has (ntotal-ndense) x nband + ndense x ntotal elements. The first (ntotal-ndense) x nband store the band part, left of diagonal, inclusive. The second ndense x ntotal store the band part as entire dense rows. Add diagadd+diagmul*mat_ii to diagonal before factorization.', # pylint: disable=line-too-long + )), + ('mju_cholSolveBand', + FunctionDecl( + name='mju_cholSolveBand', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='res', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='mat', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='vec', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='ntotal', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='nband', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='ndense', + type=ValueType(name='int'), + ), + ), + doc="Solve (mat*mat')*res = vec where mat is a band-dense Cholesky factor.", # pylint: disable=line-too-long + )), + ('mju_band2Dense', + FunctionDecl( + name='mju_band2Dense', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='res', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='mat', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='ntotal', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='nband', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='ndense', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='flg_sym', + type=ValueType(name='mjtByte'), + ), + ), + doc='Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0.', # pylint: disable=line-too-long + )), + ('mju_dense2Band', + FunctionDecl( + name='mju_dense2Band', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='res', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='mat', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='ntotal', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='nband', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='ndense', + type=ValueType(name='int'), + ), + ), + doc='Convert dense matrix to banded matrix.', + )), + ('mju_bandMulMatVec', + FunctionDecl( + name='mju_bandMulMatVec', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='res', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='mat', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='vec', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='ntotal', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='nband', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='ndense', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='nvec', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='flg_sym', + type=ValueType(name='mjtByte'), + ), + ), + doc='Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0.', # pylint: disable=line-too-long + )), + ('mju_bandDiag', + FunctionDecl( + name='mju_bandDiag', + return_type=ValueType(name='int'), + parameters=( + FunctionParameterDecl( + name='i', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='ntotal', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='nband', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='ndense', + type=ValueType(name='int'), + ), + ), + doc='Address of diagonal element i in band-dense matrix representation.', # pylint: disable=line-too-long + )), ('mju_eig3', FunctionDecl( name='mju_eig3', diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index e8a1ee45..437eec51 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -1075,6 +1075,36 @@ Euler integrator, semi-implicit in velocity. self.assertGreater(np.linalg.norm(ds_dv), eps) self.assertGreater(np.linalg.norm(ds_da), eps) + def test_banded(self): + n_total = 4 + n_band = 1 + n_dense = 1 + dense = np.array([[1.0, 0, 0, 0.1], + [0, 2.0, 0, 0.2], + [0, 0, 3.0, 0.3], + [0.1, 0.2, 0.3, 4.0]]) + band = np.zeros(n_band*(n_total-n_dense) + n_dense*n_total) + mujoco.mju_dense2Band(band, dense, n_total, n_band, n_dense) + for i in range(4): + index = mujoco.mju_bandDiag(i, n_total, n_band, n_dense) + self.assertEqual(band[index], i+1) + dense2 = np.zeros((n_total, n_total)) + flg_sym = 1 + mujoco.mju_band2Dense(dense2, band, n_total, n_band, n_dense, flg_sym) + np.testing.assert_array_equal(dense, dense2) + vec = np.array([[2.0], [2.0], [3.0], [4.0]]) + res = np.zeros_like(vec) + n_vec = 1 + mujoco.mju_bandMulMatVec(res, band, vec, + n_total, n_band, n_dense, n_vec, flg_sym) + np.testing.assert_array_equal(res, dense @ vec) + diag_add = 0 + diag_mul = 0 + mujoco.mju_cholFactorBand(band, n_total, n_band, n_dense, + diag_add, diag_mul) + mujoco.mju_cholSolveBand(res, band, vec, n_total, n_band, n_dense) + np.testing.assert_almost_equal(res, np.linalg.solve(dense, vec)) + def test_mju_box_qp(self): n = 5 res = np.zeros(n) diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 543c5f14..b0ed622e 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -991,6 +991,99 @@ PYBIND11_MODULE(_functions, pymodule) { return InterceptMjErrors(::mju_cholUpdate)( mat.data(), x.data(), mat.rows(), flg_plus); }); + Def( + pymodule, [](Eigen::Ref mat, int ntotal, int nband, + int ndense, mjtNum diagadd, mjtNum diagmul) { + int nMat = (ntotal - ndense) * nband + ndense * ntotal; + if (mat.size() != nMat) { + throw py::type_error( + "mat must have size (ntotal-ndense)*nband + ndense*ntotal"); + } + return InterceptMjErrors(::mju_cholFactorBand)( + mat.data(), ntotal, nband, ndense, diagadd, diagmul); + }); + Def( + pymodule, + [](Eigen::Ref res, Eigen::Ref mat, + Eigen::Ref vec, int ntotal, int nband, + int ndense) { + int nMat = (ntotal - ndense) * nband + ndense * ntotal; + if (mat.size() != nMat) { + throw py::type_error( + "mat must have (ntotal-ndense)*nband + " + "ndense*ntotal elements"); + } + if (res.size() != ntotal) { + throw py::type_error("size of res should equal ntotal"); + } + if (vec.size() != ntotal) { + throw py::type_error("size of vec should equal ntotal"); + } + return InterceptMjErrors(::mju_cholSolveBand)( + res.data(), mat.data(), vec.data(), ntotal, nband, ndense); + }); + Def( + pymodule, + [](Eigen::Ref res, Eigen::Ref mat, + int ntotal, int nband, int ndense, mjtByte flg_sym) { + int nMat = (ntotal - ndense) * nband + ndense * ntotal; + if (mat.size() != nMat) { + throw py::type_error( + "mat must have size (ntotal-ndense)*nband + ndense*ntotal"); + } + if (res.rows() != ntotal) { + throw py::type_error("res should have ntotal rows"); + } + if (res.cols() != ntotal) { + throw py::type_error("res should have ntotal columns"); + } + return InterceptMjErrors(::mju_band2Dense)( + res.data(), mat.data(), ntotal, nband, ndense, flg_sym); + }); + Def(pymodule, [](Eigen::Ref res, + Eigen::Ref mat, + int ntotal, int nband, int ndense) { + int nRes = (ntotal - ndense) * nband + ndense * ntotal; + if (res.size() != nRes) { + throw py::type_error( + "res must have size (ntotal-ndense)*nband + ndense*ntotal"); + } + if (mat.rows() != ntotal) { + throw py::type_error("mat should have ntotal rows"); + } + if (mat.cols() != ntotal) { + throw py::type_error("mat should have ntotal columns"); + } + return InterceptMjErrors(::mju_dense2Band)(res.data(), mat.data(), ntotal, + nband, ndense); + }); + Def( + pymodule, + [](Eigen::Ref res, Eigen::Ref mat, + Eigen::Ref vec, int ntotal, int nband, int ndense, + int nVec, mjtByte flg_sym) { + int nMat = (ntotal - ndense) * nband + ndense * ntotal; + if (mat.size() != nMat) { + throw py::type_error( + "mat must have size (ntotal-ndense)*nband + ndense*ntotal"); + } + if (res.rows() != ntotal) { + throw py::type_error("res should have ntotal rows"); + } + if (res.cols() != nVec) { + throw py::type_error("res should have nVec columns"); + } + if (vec.rows() != ntotal) { + throw py::type_error("vec should have ntotal rows"); + } + if (vec.cols() != nVec) { + throw py::type_error("vec should have nVec columns"); + } + return InterceptMjErrors(::mju_bandMulMatVec)(res.data(), mat.data(), + vec.data(), ntotal, nband, + ndense, nVec, flg_sym); + }); + Def(pymodule); Def(pymodule); DEF_WITH_OMITTED_PY_ARGS(traits::mju_boxQP, "n")( pymodule, diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 1d1c2b2b..697c4994 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -300,6 +300,288 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus, return rank; } +//---------------------------- banded Cholesky ----------------------------------------------------- + +// band-dense Cholesky decomposition +// returns minimum value in the factorized diagonal, or 0 if rank-deficient +// mat has (ntotal-ndense) x nband + ndense x ntotal elements +// the first (ntotal-ndense) x nband store the band part, left of diagonal, inclusive +// the second ndense x ntotal store the band part as entire dense rows +// add diagadd+diagmul*mat_ii to diagonal before factorization +mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense, + mjtNum diagadd, mjtNum diagmul) { + int nsparse = ntotal - ndense; + mjtNum mindiag = -1; + + // sparse part, including sparse-sparse and sparse-dense + for (int j=0; j0 ? mju_dot(mat+adr_jj-width_jj, mat+adr_jj-width_jj, width_jj) : 0; + mjtNum Ljj = diagadd + diagmul*mat[adr_jj] + mat[adr_jj] - left_ij; + + // update mindiag + if (Ljjj, sparse part + for (int i=j+1; i<=j+height; i++) { + // number of non-zeros left of (i,j) + int width_ij = mjMIN(j, nband-1-i+j); + + // address of (i,j) + int adr_ij = (i+1)*nband-1-i+j; + + // in-place computation of L(i,j) + left_ij = width_ij>0 ? mju_dot(mat+adr_jj-width_ij, mat+adr_ij-width_ij, width_ij) : 0; + mat[adr_ij] = scale * (mat[adr_ij] - left_ij); + } + + // compute L(i,j) for i>j, dense part + for (int i=nsparse; i0 ? mju_dot(mat+adr_jj-width_jj, mat+adr_ij-width_jj, width_jj) : 0; + mat[adr_ij] = scale * (mat[adr_ij] - left_ij); + } + + // save L(j,j) + mat[adr_jj] = Ljj; + } + + // dense part + for (int j=nsparse; jj + for (int i=j+1; i=nsparse; i--) { + for (int j=i+1; j=0; i--) { + // number of non-zeros below (i,i), sparse part + height = mjMIN(nsparse-1-i, nband-1); + + // sparse rows + for (int j=i+1; j<=i+height; j++) + res[i] -= mat[(j+1)*nband-1-(j-i)] * res[j]; + + // dense rows + for (int j=nsparse; j0 +MJAPI void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense, + mjtByte flg_sym); + +// convert dense matrix to banded matrix +MJAPI void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense); + +// multiply band-diagonal matrix with vector, include upper triangle if flg_sym>0 +MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, + int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym); + +// address of diagonal element i in band-dense matrix representation +MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense); + // sparse reverse-order LU factorization, no fill-in (assuming tree topology) // LU = L + U; original = (U+I) * L; scratch is size n void mju_factorLUSparse(mjtNum *LU, int n, int* scratch, diff --git a/test/engine/engine_util_solve_test.cc b/test/engine/engine_util_solve_test.cc index 0448e39a..6e79dc41 100644 --- a/test/engine/engine_util_solve_test.cc +++ b/test/engine/engine_util_solve_test.cc @@ -29,10 +29,16 @@ namespace mujoco { namespace { using ::testing::DoubleEq; +using ::testing::Pointwise; +using ::testing::DoubleNear; using ::std::string; using ::std::setw; using QCQP2Test = MujocoTest; +std::vector AsVector(const mjtNum* array, int n) { + return std::vector(array, array + n); +} + TEST_F(QCQP2Test, DegenerateAMatrix) { // A 2x2 matrix with determinant zero. const mjtNum Ain[9] { 6, -15, 2, -5 }; @@ -347,5 +353,291 @@ TEST_F(BoxQPTest, BoundedQPvariations) { mju_free(upper); } +// ------------------------- band matrices ------------------------------------- + +using BandMatrixTest = MujocoTest; + +// utility: random "arrowhead", banded-then-dense SPD matrix +// optional random vector and diagonal regularizer +void randomBanded(mjtNum* H, int nTotal, int nBand, int nDense, int seed, + mjtNum* vec = nullptr, mjtNum reg = 0) { + // make distribution using seed + std::mt19937_64 rng; + rng.seed(seed); + std::normal_distribution dist(0, 1); + + // allocate square root + mjtNum* sqrtH = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal*nTotal); + + // sample + for (int i=0; i < nTotal; i++) { + if (vec) vec[i] = dist(rng); + for (int j=0; j < nTotal; j++) { + sqrtH[nTotal*i+j] = dist(rng); + } + } + + // make SPD matrix H + mju_mulMatTMat(H, sqrtH, sqrtH, nTotal, nTotal, nTotal); + + // set zeros + int nSparse = nTotal-nDense; + for (int i=0; i < nSparse; i++) { + int nzeros = mjMAX(0, i + 1 - nBand); + for (int j=0; j < nzeros; j++) { + H[nTotal*i + j] = 0; + H[nTotal*j + i] = 0; + } + } + + // add regularizer to diagonal + for (int i=0; i < nTotal; i++) { + H[nTotal*i + i] += reg; + } + + mju_free(sqrtH); +} + + +// test banded-vector diagonal values +TEST_F(BandMatrixTest, Diagonal) { + int seed = 1; + int nTotal = 8; + for (int nBand : {1, 3}) { + for (int nDense : {0, 2}) { + // allocate + int nB = (nTotal-nDense)*nBand + nDense*nTotal; + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB); + int nH = nTotal*nTotal; + mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + + // make random banded SPD matrix, dense representation + randomBanded(H, nTotal, nBand, nDense, seed++); + + // convert to banded representation + mju_dense2Band(B, H, nTotal, nBand, nDense); + + // expect diagonals to be equal + for (int i=0; i < nTotal; i++) { + EXPECT_EQ(H[i*nTotal + i], B[mju_bandDiag(i, nTotal, nBand, nDense)]); + } + + mju_free(H); + mju_free(B); + } + } +} + +// test conversion of banded <-> dense +TEST_F(BandMatrixTest, Conversion) { + int seed = 1; + int nTotal = 8; + for (int nBand : {0, 1, 3}) { + for (int nDense : {0, 2}) { + // allocate + int nB = (nTotal-nDense)*nBand + nDense*nTotal; + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB); + int nH = nTotal*nTotal; + mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + mjtNum* H1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + + // make random banded SPD matrix, dense + randomBanded(H, nTotal, nBand, nDense, seed++); + + // convert to banded + mju_dense2Band(B, H, nTotal, nBand, nDense); + + // convert back to dense + mju_band2Dense(H1, B, nTotal, nBand, nDense, /*flg_sym=*/1); + + // expect exact equality + EXPECT_EQ(AsVector(H, nH), AsVector(H1, nH)); + + mju_free(H1); + mju_free(H); + mju_free(B); + } + } +} + +// test banded-vector multiplication +TEST_F(BandMatrixTest, Multiplication) { + int seed = 1; + int nTotal = 8; + for (int nBand : {0, 1, 3}) { + for (int nDense : {0, 2}) { + // allocate + int nB = (nTotal-nDense)*nBand + nDense*nTotal; + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB); + int nH = nTotal*nTotal; + mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + mjtNum* res1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + + // make random banded SPD matrix, dense + randomBanded(H, nTotal, nBand, nDense, seed++, vec); + + // multiply dense + mju_mulMatVec(res, H, vec, nTotal, nTotal); + + // convert to banded, multiply + mju_dense2Band(B, H, nTotal, nBand, nDense); + mju_bandMulMatVec(res1, B, vec, nTotal, nBand, nDense, + /*nVec=*/1, /*flg_sym=*/1); + + // expect numerical equality + mjtNum eps = 1e-12; + EXPECT_THAT(AsVector(res, nTotal), + Pointwise(DoubleNear(eps), AsVector(res1, nTotal))); + + mju_free(res1); + mju_free(res); + mju_free(vec); + mju_free(H); + mju_free(B); + } + } +} + +// test banded factorization and vector product with factor +TEST_F(BandMatrixTest, Factorization) { + int seed = 1; + int nTotal = 8; + for (int nBand : {1, 3}) { + for (int nDense : {0, 2}) { + for (mjtNum diagadd : {0.0, 1.0}) { + for (int diagmul : {0.0, 1.3}) { + // allocate + int nB = (nTotal-nDense)*nBand + nDense*nTotal; + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB); + int nH = nTotal*nTotal; + mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + mjtNum* H1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + mjtNum* res1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + + // make random banded matrix, dense representation + // add regularizer to ensure PD + randomBanded(H, nTotal, nBand, nDense, seed++, vec, /*reg=*/nTotal); + + // convert to banded + mju_dense2Band(B, H, nTotal, nBand, nDense); + + // apply diagadd and diagmul + for (int i=0; i < nTotal; i++) { + H[nTotal*i + i] += diagadd + diagmul*H[nTotal*i + i]; + } + + // in-place dense factorization + int rank = mju_cholFactor(H, nTotal, /*mindiag=*/0); + + // expect factorization to have succeeded + EXPECT_EQ(rank, nTotal); + + // banded factorization + mjtNum minDiag = mju_cholFactorBand(B, nTotal, nBand, nDense, + diagadd, diagmul); + + // expect factorization to have succeeded + EXPECT_GT(minDiag, 0); + + // convert back to dense, lower triangle only + mju_band2Dense(H1, B, nTotal, nBand, nDense, /*flg_sym=*/0); + + // zero upper triangle of H (unused) + for (int i=0; i < nTotal-1; i++) { + mju_zero(H + nTotal*i + i + 1, nTotal - i - 1); + } + + // expect numerical equality + mjtNum eps = 1e-12; + EXPECT_THAT(AsVector(H, nH), + Pointwise(DoubleNear(eps), AsVector(H1, nH))); + + // multiply dense + mju_mulMatVec(res, H, vec, nTotal, nTotal); + + // multiply sparse, only lower triangle + mju_bandMulMatVec(res1, B, vec, nTotal, nBand, nDense, + /*nVec=*/1, /*flg_sym=*/0); + + // expect numerical equality + EXPECT_THAT(AsVector(res, nTotal), + Pointwise(DoubleNear(eps), AsVector(res1, nTotal))); + + mju_free(res1); + mju_free(res); + mju_free(vec); + mju_free(H1); + mju_free(H); + mju_free(B); + } + } + } + } +} + +// test banded solve +TEST_F(BandMatrixTest, Solve) { + int seed = 1; + int nTotal = 8; + for (int nBand : {1, 3}) { + for (int nDense : {0, 2}) { + // allocate + int nB = (nTotal-nDense)*nBand + nDense*nTotal; + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*nB); + int nH = nTotal*nTotal; + mjtNum* H = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + mjtNum* H1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nH); + mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + mjtNum* res1 = (mjtNum*) mju_malloc(sizeof(mjtNum)*nTotal); + + // make random banded matrix, dense representation + // add regularizer to ensure PD + randomBanded(H, nTotal, nBand, nDense, seed++, vec, /*reg=*/nTotal); + + // convert to banded + mju_dense2Band(B, H, nTotal, nBand, nDense); + + // in-place dense factorization + int rank = mju_cholFactor(H, nTotal, /*mindiag=*/0); + + // expect factorization to have succeeded + EXPECT_EQ(rank, nTotal); + + // banded factorization + mjtNum minDiag = mju_cholFactorBand(B, nTotal, nBand, nDense, + /*diagadd=*/0, /*diagmul=*/0); + + // expect factorization to have succeeded + EXPECT_GT(minDiag, 0); + + // convert back to dense, lower triangle only + mju_band2Dense(H1, B, nTotal, nBand, nDense, /*flg_sym=*/0); + + // solve with dense + mju_cholSolve(res, H, vec, nTotal); + + // solve with banded + mju_cholSolveBand(res1, B, vec, nTotal, nBand, nDense); + + // expect numerical equality + mjtNum eps = 1e-12; + EXPECT_THAT(AsVector(res, nTotal), + Pointwise(DoubleNear(eps), AsVector(res1, nTotal))); + + mju_free(res1); + mju_free(res); + mju_free(vec); + mju_free(H1); + mju_free(H); + mju_free(B); + } + } +} + } // namespace } // namespace mujoco diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 0974f7f9..cac3cd88 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -3733,6 +3733,24 @@ public static unsafe extern void mju_cholSolve(double* res, double* mat, double* [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern int mju_cholUpdate(double* mat, double* x, int n, int flg_plus); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern double mju_cholFactorBand(double* mat, int ntotal, int nband, int ndense, double diagadd, double diagmul); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mju_cholSolveBand(double* res, double* mat, double* vec, int ntotal, int nband, int ndense); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mju_band2Dense(double* res, double* mat, int ntotal, int nband, int ndense, byte flg_sym); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mju_dense2Band(double* res, double* mat, int ntotal, int nband, int ndense); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mju_bandMulMatVec(double* res, double* mat, double* vec, int ntotal, int nband, int ndense, int nvec, byte flg_sym); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern int mju_bandDiag(int i, int ntotal, int nband, int ndense); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern int mju_eig3(double* eigval, double* eigvec, double* quat, double* mat);