Map canonical links to their models (#736)

Signed-off-by: Ashton Larkin <ashton@openrobotics.org>

src/systems/physics/CanonicalLinkModelTracker.hh

@@ -0,0 +1,116 @@
+/*
+ * Copyright (C) 2021 Open Source Robotics Foundation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+*/
+#ifndef IGNITION_GAZEBO_SYSTEMS_PHYSICS_CANONICAL_LINK_MODEL_TRACKER_HH_
+#define IGNITION_GAZEBO_SYSTEMS_PHYSICS_CANONICAL_LINK_MODEL_TRACKER_HH_
+
+#include <set>
+#include <unordered_map>
+
+#include "ignition/gazebo/Entity.hh"
+#include "ignition/gazebo/EntityComponentManager.hh"
+#include "ignition/gazebo/components/CanonicalLink.hh"
+#include "ignition/gazebo/components/Model.hh"
+#include "ignition/gazebo/config.hh"
+
+namespace ignition::gazebo
+{
+inline namespace IGNITION_GAZEBO_VERSION_NAMESPACE {
+namespace systems::physics_system
+{
+  /// \brief Helper class that keeps track of which models have a particular
+  /// canonical link. This is useful in the physics system for updating model
+  /// poses - if a canonical link moved in the most recent physics step, then
+  /// all of the models that have this canonical link should be updated. It's
+  /// important to preserve topological ordering of the models in case there's
+  /// a nested model that shares the same canonical link (in a case like this,
+  /// the parent model pose needs to be updated before updating the child model
+  /// pose - see the documentation that explains how model pose updates are
+  /// calculated in PhysicsPrivate::UpdateSim to understand why nested model
+  /// poses need to be updated in topological order).
+  ///
+  /// It's possible to loop through all of the models and to update poses if the
+  /// model moved using something like EntityComponentManager::Each, but the
+  /// performance of this approach is worse than using just the moved canonical
+  /// links to determine which model poses should be updated (consider the case
+  /// where there are a lot of non-static models in a world, but only a few move
+  /// frequently - if using EntityComponentManager::Each, we still need to check
+  /// every single non-static model after a physics update to make sure that the
+  /// model did not move. If we instead use the updated canonical link
+  /// information, then we can skip iterating over/checking the models that
+  /// don't need to be updated).
+  class CanonicalLinkModelTracker
+  {
+    /// \brief Save mappings for new models and their canonical links
+    /// \param[in] _ecm EntityComponentManager
+    public: void AddNewModels(const EntityComponentManager &_ecm);
+
+    /// \brief Get a topological ordering of models that have a particular
+    /// canonical link
+    /// \param[in] _canonicalLink The canonical link
+    /// \return The models that have this link as their canonical link, in
+    /// topological order
+    public: const std::set<Entity> &CanonicalLinkModels(
+                const Entity _canonicalLink) const;
+
+    /// \brief Remove a link from the mapping. This method should be called when
+    /// a link is removed from simulation
+    /// \param[in] _link The link to remove
+    public: void RemoveLink(const Entity &_link);
+
+    /// \brief A mapping of canonical links to the models that have this
+    /// canonical link. The key is the canonical link entity, and the value is
+    /// the model entities that have this canonical link. The models in the
+    /// value are in topological order
+    private: std::unordered_map<Entity, std::set<Entity>> linkModelMap;
+
+    /// \brief An empty set of models that is returned from the
+    /// CanonicalLinkModels method for links that map to no models
+    private: const std::set<Entity> emptyModelOrdering{};
+  };
+
+  void CanonicalLinkModelTracker::AddNewModels(
+      const EntityComponentManager &_ecm)
+  {
+    _ecm.EachNew<components::Model, components::ModelCanonicalLink>(
+        [this](const Entity &_model, const components::Model *,
+          const components::ModelCanonicalLink *_canonicalLinkComp)
+        {
+          this->linkModelMap[_canonicalLinkComp->Data()].insert(_model);
+          return true;
+        });
+  }
+
+  const std::set<Entity> &CanonicalLinkModelTracker::CanonicalLinkModels(
+      const Entity _canonicalLink) const
+  {
+    auto it = this->linkModelMap.find(_canonicalLink);
+    if (it != this->linkModelMap.end())
+      return it->second;
+
+    // if an invalid entity was given, it maps to no models
+    return this->emptyModelOrdering;
+  }
+
+  void CanonicalLinkModelTracker::RemoveLink(const Entity &_link)
+  {
+    this->linkModelMap.erase(_link);
+  }
+}
+}
+}
+
+#endif

src/systems/physics/Physics.cc

@@ -25,6 +25,7 @@
 #include <algorithm>
 #include <iostream>
 #include <deque>
+#include <map>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -119,6 +120,7 @@
 #include "ignition/gazebo/components/ThreadPitch.hh"
 #include "ignition/gazebo/components/World.hh"
 
+#include "CanonicalLinkModelTracker.hh"
 #include "EntityFeatureMap.hh"
 
 using namespace ignition;
@@ -193,7 +195,10 @@ class ignition::gazebo::systems::PhysicsPrivate
   /// not write this data to ForwardStep::Output. If not, _ecm is used to get
   /// this updated link pose data).
   /// \return A map of gazebo link entities to their updated pose data.
-  public: std::unordered_map<Entity, physics::FrameData3d> ChangedLinks(
+  /// std::map is used because canonical links must be in topological order
+  /// to ensure that nested models with multiple canonical links are updated
+  /// properly (models must be updated in topological order).
+  public: std::map<Entity, physics::FrameData3d> ChangedLinks(
               EntityComponentManager &_ecm,
               const ignition::physics::ForwardStep::Output &_updatedLinks);
 
@@ -203,7 +208,7 @@ class ignition::gazebo::systems::PhysicsPrivate
   /// most recent physics step. The key is the entity of the link, and the
   /// value is the updated frame data corresponding to that entity.
   public: void UpdateSim(EntityComponentManager &_ecm,
-              const std::unordered_map<
+              const std::map<
                 Entity, physics::FrameData3d> &_linkFrameData);
 
   /// \brief Update collision components from physics simulation
@@ -237,6 +242,11 @@ class ignition::gazebo::systems::PhysicsPrivate
   /// after a physics step.
   public: std::unordered_map<Entity, ignition::math::Pose3d> linkWorldPoses;
 
+  /// \brief Keep a mapping of canonical links to models that have this
+  /// canonical link. Useful for updating model poses efficiently after a
+  /// physics step
+  public: CanonicalLinkModelTracker canonicalLinkModelTracker;
+
   /// \brief Keep track of non-static model world poses. Since non-static
   /// models may not move on a given iteration, we want to keep track of the
   /// most recent model world pose change that took place.
@@ -1178,6 +1188,7 @@ void PhysicsPrivate::RemovePhysicsEntities(const EntityComponentManager &_ecm)
             this->topLevelModelMap.erase(childLink);
             this->staticEntities.erase(childLink);
             this->linkWorldPoses.erase(childLink);
+            this->canonicalLinkModelTracker.RemoveLink(childLink);
           }
 
           for (const auto &childJoint :
@@ -1835,13 +1846,13 @@ ignition::math::Pose3d PhysicsPrivate::RelativePose(const Entity &_from,
 }
 
 //////////////////////////////////////////////////
-std::unordered_map<Entity, physics::FrameData3d> PhysicsPrivate::ChangedLinks(
+std::map<Entity, physics::FrameData3d> PhysicsPrivate::ChangedLinks(
     EntityComponentManager &_ecm,
     const ignition::physics::ForwardStep::Output &_updatedLinks)
 {
   IGN_PROFILE("Links Frame Data");
 
-  std::unordered_map<Entity, physics::FrameData3d> linkFrameData;
+  std::map<Entity, physics::FrameData3d> linkFrameData;
 
   // Check to see if the physics engine gave a list of changed poses. If not, we
   // will iterate through all of the links via the ECM to see which ones changed
@@ -1912,83 +1923,91 @@ std::unordered_map<Entity, physics::FrameData3d> PhysicsPrivate::ChangedLinks(
 
 //////////////////////////////////////////////////
 void PhysicsPrivate::UpdateSim(EntityComponentManager &_ecm,
-    const std::unordered_map<Entity, physics::FrameData3d> &_linkFrameData)
+    const std::map<Entity, physics::FrameData3d> &_linkFrameData)
 {
   IGN_PROFILE("PhysicsPrivate::UpdateSim");
 
   IGN_PROFILE_BEGIN("Models");
 
-  _ecm.Each<components::Model, components::ModelCanonicalLink>(
-      [&](const Entity &_entity, components::Model *,
-          components::ModelCanonicalLink *_canonicalLink) -> bool
-      {
-        // If the model's canonical link did not move, we don't need to update
-        // the model's pose
-        auto linkFrameIt = _linkFrameData.find(_canonicalLink->Data());
-        if (linkFrameIt == _linkFrameData.end())
-          return true;
-
-        std::optional<math::Pose3d> parentWorldPose;
+  // make sure we have an up-to-date mapping of canonical links to their models
+  this->canonicalLinkModelTracker.AddNewModels(_ecm);
 
-        // If this model is nested, we assume the pose of the parent model has
-        // already been updated. We expect to find the updated pose in
-        // this->modelWorldPoses. If not found, this must not be nested, so
-        // this model's pose component would reflect it's absolute pose.
-        auto parentModelPoseIt =
-          this->modelWorldPoses.find(
-              _ecm.Component<components::ParentEntity>(_entity)->Data());
-        if (parentModelPoseIt != this->modelWorldPoses.end())
-        {
-          parentWorldPose = parentModelPoseIt->second;
-        }
+  for (const auto &[linkEntity, frameData] : _linkFrameData)
+  {
+    // get a topological ordering of the models that have linkEntity as the
+    // model's canonical link. If linkEntity isn't a canonical link for any
+    // models, canonicalLinkModels will be empty
+    auto canonicalLinkModels =
+      this->canonicalLinkModelTracker.CanonicalLinkModels(linkEntity);
+
+    // Update poses for all of the models that have this changed canonical link
+    // (linkEntity). Since we have the models in topological order and
+    // _linkFrameData stores links in topological order thanks to the ordering
+    // of std::map (entity IDs are created in ascending order), this should
+    // properly handle pose updates for nested models
+    for (auto &model : canonicalLinkModels)
+    {
+      std::optional<math::Pose3d> parentWorldPose;
+
+      // If this model is nested, the pose of the parent model has already
+      // been updated since we iterate through the modified links in
+      // topological order. We expect to find the updated pose in
+      // this->modelWorldPoses. If not found, this must not be nested, so this
+      // model's pose component would reflect it's absolute pose.
+      auto parentModelPoseIt =
+        this->modelWorldPoses.find(
+            _ecm.Component<components::ParentEntity>(model)->Data());
+      if (parentModelPoseIt != this->modelWorldPoses.end())
+      {
+        parentWorldPose = parentModelPoseIt->second;
+      }
 
-        // Given the following frame names:
-        // W: World/inertial frame
-        // P: Parent frame (this could be a parent model or the World frame)
-        // M: This model's frame
-        // L: The frame of this model's canonical link
-        //
-        // And the following quantities:
-        // (See http://sdformat.org/tutorials?tut=specify_pose for pose
-        // convention)
-        // parentWorldPose (X_WP): Pose of the parent frame w.r.t the world
-        // linkPoseFromModel (X_ML): Pose of the canonical link frame w.r.t the
-        // model frame
-        // linkWorldPose (X_WL): Pose of the canonical link w.r.t the world
-        // modelWorldPose (X_WM): Pose of this model w.r.t the world
-        //
-        // The Pose component of this model entity stores the pose of M w.r.t P
-        // (X_PM) and is calculated as
-        //   X_PM = (X_WP)^-1 * X_WM
-        //
-        // And X_WM is calculated from X_WL, which is obtained from physics as:
-        //   X_WM = X_WL * (X_ML)^-1
-        auto linkPoseFromModel =
-            this->RelativePose(_entity, linkFrameIt->first, _ecm);
-        const auto &linkWorldPose = linkFrameIt->second.pose;
-        const auto &modelWorldPose =
-            math::eigen3::convert(linkWorldPose) * linkPoseFromModel.Inverse();
-
-        this->modelWorldPoses[_entity] = modelWorldPose;
-
-        // update model's pose
-        auto modelPose = _ecm.Component<components::Pose>(_entity);
-        if (parentWorldPose)
-        {
-          *modelPose =
-              components::Pose(parentWorldPose->Inverse() * modelWorldPose);
-        }
-        else
-        {
-          // This is a non-nested model and parentWorldPose would be identity
-          // because it would be the pose of the parent (world) w.r.t the world.
-          *modelPose = components::Pose(modelWorldPose);
-        }
+      // Given the following frame names:
+      // W: World/inertial frame
+      // P: Parent frame (this could be a parent model or the World frame)
+      // M: This model's frame
+      // L: The frame of this model's canonical link
+      //
+      // And the following quantities:
+      // (See http://sdformat.org/tutorials?tut=specify_pose for pose
+      // convention)
+      // parentWorldPose (X_WP): Pose of the parent frame w.r.t the world
+      // linkPoseFromModel (X_ML): Pose of the canonical link frame w.r.t the
+      // model frame
+      // linkWorldPose (X_WL): Pose of the canonical link w.r.t the world
+      // modelWorldPose (X_WM): Pose of this model w.r.t the world
+      //
+      // The Pose component of this model entity stores the pose of M w.r.t P
+      // (X_PM) and is calculated as
+      //   X_PM = (X_WP)^-1 * X_WM
+      //
+      // And X_WM is calculated from X_WL, which is obtained from physics as:
+      //   X_WM = X_WL * (X_ML)^-1
+      auto linkPoseFromModel = this->RelativePose(model, linkEntity, _ecm);
+      const auto &linkWorldPose = frameData.pose;
+      const auto &modelWorldPose =
+          math::eigen3::convert(linkWorldPose) * linkPoseFromModel.Inverse();
+
+      this->modelWorldPoses[model] = modelWorldPose;
+
+      // update model's pose
+      auto modelPose = _ecm.Component<components::Pose>(model);
+      if (parentWorldPose)
+      {
+        *modelPose =
+            components::Pose(parentWorldPose->Inverse() * modelWorldPose);
+      }
+      else
+      {
+        // This is a non-nested model and parentWorldPose would be identity
+        // because it would be the pose of the parent (world) w.r.t the world.
+        *modelPose = components::Pose(modelWorldPose);
+      }
 
-        _ecm.SetChanged(_entity, components::Pose::typeId,
-                        ComponentState::PeriodicChange);
-        return true;
-      });
+      _ecm.SetChanged(model, components::Pose::typeId,
+                      ComponentState::PeriodicChange);
+    }
+  }
   IGN_PROFILE_END();
 
   // Link poses, velocities...