From 1d47a91c19eced1e9de1d445205bb31e1d1cae9d Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sun, 19 Jan 2025 11:33:35 -0800
Subject: [PATCH] Add a test for mj_solveM2

PiperOrigin-RevId: 717278730
Change-Id: I7507e569620c34f250482eec8402c08c793f4e48
---
 test/engine/engine_core_smooth_test.cc | 46 ++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc
index e5f66c38..b63e57a3 100644
--- a/test/engine/engine_core_smooth_test.cc
+++ b/test/engine/engine_core_smooth_test.cc
@@ -545,6 +545,52 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
   mj_deleteModel(m);
 }
 
+TEST_F(CoreSmoothTest, SolveM2) {
+  const std::string xml_path = GetTestDataFilePath(kInertiaPath);
+  char error[1024];
+  mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error;
+
+  mjData* d = mj_makeData(m);
+  mj_forward(m, d);
+
+  int nv = m->nv;
+  int nC = m->nC;
+
+  // copy LD into LDs: CSR format
+  vector<mjtNum> LDs(nC);
+  for (int i=0; i < nC; i++) {
+    LDs[i] = d->qLD[d->mapM2C[i]];
+  }
+
+  // inverse square root of D from inertia LDL decomposition
+  vector<mjtNum> sqrtInvD(nv);
+  for (int i=0; i < nv; i++) {
+    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
+  }
+
+  // compare full solve and half solve
+  int n = 3;
+  vector<mjtNum> vec(nv*n);
+  vector<mjtNum> vec2(nv*n);
+  for (int i=0; i < nv*n; i++) vec[i] = vec2[i] = 2 + 3*i;
+  for (int i=0; i < nv*n; i+=3) vec[i] = vec2[i] = 0;
+  vector<mjtNum> res(nv*n);
+
+  mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n);
+  mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n,
+              d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+
+  // expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v)
+  for (int i=0; i < n; i++) {
+    EXPECT_FLOAT_EQ(mju_dot(vec2.data() + i*nv, vec.data() + i*nv, nv),
+                    mju_dot(res.data() + i*nv, res.data() + i*nv, nv));
+  }
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 TEST_F(CoreSmoothTest, FactorIs) {
   const std::string xml_path = GetTestDataFilePath(kInertiaPath);
   char error[1024];