Add utility functions for banded-then-dense symmetric matrices.

PiperOrigin-RevId: 528757637
Change-Id: I916d843140322c28e857100a6b305a041f704d53

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)

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<traits::mju_cholFactorBand>(
+      pymodule, [](Eigen::Ref<EigenVectorX> 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<traits::mju_cholSolveBand>(
+      pymodule,
+      [](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenVectorX> mat,
+         Eigen::Ref<const EigenVectorX> 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<traits::mju_band2Dense>(
+      pymodule,
+      [](Eigen::Ref<EigenArrayXX> res, Eigen::Ref<const EigenVectorX> 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<traits::mju_dense2Band>(pymodule, [](Eigen::Ref<EigenVectorX> res,
+                                           Eigen::Ref<const EigenArrayXX> 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<traits::mju_bandMulMatVec>(
+      pymodule,
+      [](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenArrayXX> mat,
+         Eigen::Ref<const EigenArrayXX> 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<traits::mju_bandDiag>(pymodule);
   Def<traits::mju_eig3>(pymodule);
   DEF_WITH_OMITTED_PY_ARGS(traits::mju_boxQP, "n")(
       pymodule,