From f253f0defdaa204c64aeb0bc1bed80e9713b9008 Mon Sep 17 00:00:00 2001
From: Ian Chen <ichen@openrobotics.org>
Date: Tue, 10 Dec 2024 00:52:04 +0800
Subject: [PATCH] rescale inertia to match specified mass

Signed-off-by: Ian Chen <ichen@openrobotics.org>
---
 src/Link.cc      | 24 +++++++++++++++++++++---
 src/Link_TEST.cc | 46 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/src/Link.cc b/src/Link.cc
index 19c3569fe..aef87e48f 100644
--- a/src/Link.cc
+++ b/src/Link.cc
@@ -211,7 +211,6 @@ Errors Link::Load(ElementPtr _sdf, const ParserConfig &_config)
       // If auto is to true but user has still provided
       // inertial values
       if (inertialElem->HasElement("pose") ||
-          inertialElem->HasElement("mass") ||
           inertialElem->HasElement("inertia"))
       {
         Error err(
@@ -681,7 +680,27 @@ void Link::ResolveAutoInertials(sdf::Errors &_errors,
       totalInertia = totalInertia + collisionInertia;
     }
 
-    this->dataPtr->inertial = totalInertia;
+    // If mass is specified, scale inertia to match desired mass
+    auto inertialElem = this->dataPtr->sdf->GetElement("inertial");
+    if (inertialElem->HasElement("mass"))
+    {
+      double mass = inertialElem->Get<double>("mass");
+      const gz::math::MassMatrix3d &totalMassMatrix = totalInertia.MassMatrix();
+      // normalize to get the unit mass matrix
+      gz::math::MassMatrix3d unitMassMatrix(1.0,
+          totalMassMatrix.DiagonalMoments() / totalMassMatrix.Mass(),
+          totalMassMatrix.OffDiagonalMoments() / totalMassMatrix.Mass());
+      // scale the final inertia to match specified mass
+      this->dataPtr->inertial = gz::math::Inertiald(
+          gz::math::MassMatrix3d(mass,
+              mass * unitMassMatrix.DiagonalMoments(),
+              mass * unitMassMatrix.OffDiagonalMoments()),
+          totalInertia.Pose());
+    }
+    else
+    {
+      this->dataPtr->inertial = totalInertia;
+    }
 
     // If CalculateInertial() was called with SAVE_CALCULATION
     // configuration then set autoInertiaSaved to true
@@ -695,7 +714,6 @@ void Link::ResolveAutoInertials(sdf::Errors &_errors,
     {
       this->dataPtr->autoInertiaSaved = true;
       // Write calculated inertia values to //link/inertial element
-      auto inertialElem = this->dataPtr->sdf->GetElement("inertial");
       inertialElem->GetElement("pose")->GetValue()->Set<gz::math::Pose3d>(
         totalInertia.Pose());
       inertialElem->GetElement("mass")->GetValue()->Set<double>(
diff --git a/src/Link_TEST.cc b/src/Link_TEST.cc
index 7a9207dd0..8c81bd496 100644
--- a/src/Link_TEST.cc
+++ b/src/Link_TEST.cc
@@ -725,7 +725,7 @@ TEST(DOMLink, ResolveAutoInertialsWithMultipleCollisions)
 }
 
 /////////////////////////////////////////////////
-TEST(DOMLink, InertialValuesGivenWithAutoSetToTrue)
+TEST(DOMLink, ResolveAutoInertialsWithMass)
 {
   std::string sdf = "<?xml version=\"1.0\"?>"
   "<sdf version=\"1.11\">"
@@ -763,6 +763,50 @@ TEST(DOMLink, InertialValuesGivenWithAutoSetToTrue)
   root.ResolveAutoInertials(errors, sdfParserConfig);
   EXPECT_TRUE(errors.empty());
 
+  EXPECT_DOUBLE_EQ(4.0, link->Inertial().MassMatrix().Mass());
+  EXPECT_EQ(gz::math::Pose3d::Zero, link->Inertial().Pose());
+  EXPECT_EQ(gz::math::Vector3d(0.66666, 0.66666, 0.66666),
+    link->Inertial().MassMatrix().DiagonalMoments());
+}
+
+/////////////////////////////////////////////////
+TEST(DOMLink, InertialValuesGivenWithAutoSetToTrue)
+{
+  std::string sdf = "<?xml version=\"1.0\"?>"
+  "<sdf version=\"1.11\">"
+  "  <model name='compound_model'>"
+  "   <link name='compound_link'>"
+  "     <inertial auto='true'>"
+  "       <pose>1 1 1 2 2 2</pose>"
+  "       <inertia>"
+  "         <ixx>1</ixx>"
+  "         <iyy>1</iyy>"
+  "         <izz>1</izz>"
+  "       </inertia>"
+  "     </inertial>"
+  "     <collision name='box_collision'>"
+  "      <density>2.0</density>"
+  "      <geometry>"
+  "        <box>"
+  "          <size>1 1 1</size>"
+  "        </box>"
+  "        </geometry>"
+  "      </collision>"
+  "    </link>"
+  "   </model>"
+  "  </sdf>";
+
+  sdf::Root root;
+  const sdf::ParserConfig sdfParserConfig;
+  sdf::Errors errors = root.LoadSdfString(sdf, sdfParserConfig);
+  EXPECT_TRUE(errors.empty());
+  EXPECT_NE(nullptr, root.Element());
+
+  const sdf::Model *model = root.Model();
+  const sdf::Link *link = model->LinkByIndex(0);
+  root.ResolveAutoInertials(errors, sdfParserConfig);
+  EXPECT_TRUE(errors.empty());
+
   EXPECT_NE(4.0, link->Inertial().MassMatrix().Mass());
   EXPECT_EQ(gz::math::Pose3d::Zero, link->Inertial().Pose());
   EXPECT_EQ(gz::math::Vector3d(0.33333, 0.33333, 0.33333),