Adds misc algorithms over graph and forest

Author: sherm1

multibody/topology/link_joint_graph.cc

@@ -704,6 +704,169 @@ bool LinkJointGraph::link_is_static(const Link& link) const {
       ForestBuildingOptions::kStatic);
 }
 
+/* Runs through the Mobods in the model but records the (active) Link
+indexes rather than the Mobod indexes. */
+std::vector<LinkIndex> LinkJointGraph::FindPathFromWorld(
+    LinkIndex link_index) const {
+  ThrowIfForestNotBuiltYet(__func__);
+  const SpanningForest::Mobod* mobod =
+      &forest().mobods()[link_to_mobod(link_index)];
+  std::vector<LinkIndex> path(mobod->level() + 1);
+  while (mobod->inboard().is_valid()) {
+    const Link& link = links(mobod->link_ordinal());
+    path[mobod->level()] = link.index();  // Active Link if composite.
+    mobod = &forest().mobods(mobod->inboard());
+  }
+  DRAKE_DEMAND(mobod->is_world());
+  path[0] = LinkIndex(0);
+  return path;
+}
+
+LinkIndex LinkJointGraph::FindFirstCommonAncestor(LinkIndex link1_index,
+                                                  LinkIndex link2_index) const {
+  ThrowIfForestNotBuiltYet(__func__);
+  const MobodIndex mobod_ancestor = forest().FindFirstCommonAncestor(
+      link_to_mobod(link1_index), link_to_mobod(link2_index));
+  const Link& ancestor_link =
+      links(forest().mobod_to_link_ordinal(mobod_ancestor));
+  return ancestor_link.index();
+}
+
+std::vector<LinkIndex> LinkJointGraph::FindSubtreeLinks(
+    LinkIndex link_index) const {
+  ThrowIfForestNotBuiltYet(__func__);
+  const MobodIndex root_mobod_index = link_to_mobod(link_index);
+  return forest().FindSubtreeLinks(root_mobod_index);
+}
+
+// Our link_composites collection doesn't include lone Links that aren't welded
+// to anything. The return from this function must include every Link, with
+// the World link in the first set (even if nothing is welded to it).
+std::vector<std::set<LinkIndex>> LinkJointGraph::GetSubgraphsOfWeldedLinks()
+    const {
+  ThrowIfForestNotBuiltYet(__func__);
+
+  std::vector<std::set<LinkIndex>> subgraphs;
+
+  // First, collect all the precomputed Link Composites. World is always
+  // the first one, even if nothing is welded to it.
+  for (const LinkComposite& composite : link_composites()) {
+    subgraphs.emplace_back(
+        std::set<LinkIndex>(composite.links.cbegin(), composite.links.cend()));
+  }
+
+  // Finally, make one-Link subgraphs for Links that aren't in any composite.
+  for (const Link& link : links()) {
+    if (link.composite().has_value()) continue;
+    subgraphs.emplace_back(std::set<LinkIndex>{link.index()});
+  }
+
+  return subgraphs;
+}
+
+// Strategy here is to make repeated use of CalcLinksWeldedTo(), separating
+// the singleton sets from the actually-welded sets, and then move the
+// singletons to the end to match what GetSubgraphsOfWeldedLinks() does.
+std::vector<std::set<LinkIndex>> LinkJointGraph::CalcSubgraphsOfWeldedLinks()
+    const {
+  // Work with ordinals rather than indexes.
+  std::vector<bool> visited(num_user_links(), false);
+
+  // World always comes first, even if it is alone.
+  std::vector<std::set<LinkIndex>> subgraphs{CalcLinksWeldedTo(LinkIndex(0))};
+  for (LinkIndex index : subgraphs[0]) visited[index_to_ordinal(index)] = true;
+
+  std::vector<std::set<LinkIndex>> singletons;
+  // If a Forest was already built, there may be shadow links added to
+  // the graph -- don't process those here.
+  for (LinkOrdinal link_ordinal(1); link_ordinal < num_user_links();
+       ++link_ordinal) {
+    const Link& link = links(link_ordinal);
+    if (link.is_shadow() || visited[link_ordinal]) continue;
+    std::set<LinkIndex> welded_links = CalcLinksWeldedTo(link.index());
+    for (LinkIndex index : welded_links)
+      visited[index_to_ordinal(index)] = true;
+    if (ssize(welded_links) == 1) {
+      singletons.emplace_back(std::move(welded_links));
+    } else {
+      subgraphs.emplace_back(std::move(welded_links));
+    }
+  }
+
+  // Now move all the singletons onto the end of the subgraphs list.
+  for (auto& singleton : singletons)
+    subgraphs.emplace_back(std::move(singleton));
+
+  return subgraphs;
+}
+
+// If the Link isn't part of a LinkComposite just return the Link. Otherwise,
+// return all the Links in its LinkComposite.
+std::set<LinkIndex> LinkJointGraph::GetLinksWeldedTo(
+    LinkIndex link_index) const {
+  ThrowIfForestNotBuiltYet(__func__);
+  DRAKE_DEMAND(link_index.is_valid());
+  DRAKE_THROW_UNLESS(has_link(link_index));
+  const Link& link = link_by_index(link_index);
+  const std::optional<LinkCompositeIndex> composite_index = link.composite();
+  if (!composite_index.has_value()) return std::set<LinkIndex>{link_index};
+  const std::vector<LinkIndex>& welded_links =
+      link_composites(*composite_index).links;
+  return std::set<LinkIndex>(welded_links.cbegin(), welded_links.cend());
+}
+
+// Without a Forest we don't have LinkComposites available so recursively
+// chase Weld joints instead.
+std::set<LinkIndex> LinkJointGraph::CalcLinksWeldedTo(
+    LinkIndex link_index) const {
+  std::set<LinkIndex> result;
+  AppendLinksWeldedTo(link_index, &result);
+  return result;
+}
+
+void LinkJointGraph::AppendLinksWeldedTo(LinkIndex link_index,
+                                         std::set<LinkIndex>* result) const {
+  DRAKE_DEMAND(result != nullptr);
+  DRAKE_DEMAND(link_index.is_valid());
+  DRAKE_THROW_UNLESS(has_link(link_index));
+  DRAKE_DEMAND(!result->contains(link_index));
+
+  const Link& link = link_by_index(link_index);
+
+  // A Link is always considered welded to itself.
+  result->insert(link_index);
+
+  // For World we have to look for static links and pretend they are welded to
+  // World. (Links might have been explicitly flagged as static or part of a
+  // static model instance.)
+  if (link.is_world()) {
+    for (const Link& maybe_static : links()) {
+      if (result->contains(maybe_static.index())) continue;
+      if (link_is_static(maybe_static))
+        AppendLinksWeldedTo(maybe_static.index(), &*result);
+    }
+  }
+
+  // Now run through all the actual joints, looking for welds.
+  for (auto joint_index : link.joints()) {
+    const Joint& joint = joint_by_index(joint_index);
+    if (joint.traits_index() != weld_joint_traits_index()) continue;
+    const LinkIndex welded_link_index = joint.other_link_index(link_index);
+    // Don't reprocess if we already did the other end.
+    if (!result->contains(welded_link_index))
+      AppendLinksWeldedTo(welded_link_index, &*result);
+  }
+}
+
+void LinkJointGraph::ThrowIfForestNotBuiltYet(const char* func) const {
+  if (!forest_is_valid()) {
+    throw std::logic_error(
+        fmt::format("{}(): no SpanningForest available. Call BuildForest() "
+                    "before calling this function.",
+                    func));
+  }
+}
+
 void LinkJointGraph::ThrowLinkWasRemoved(const char* func,
                                          LinkIndex link_index) const {
   throw std::logic_error(fmt::format(

multibody/topology/link_joint_graph.h

@@ -417,11 +417,11 @@ class LinkJointGraph {
   [[nodiscard]] inline const LoopConstraint& loop_constraints(
       LoopConstraintIndex constraint_index) const;
 
-  /** Links with this index or higher are "ephemeral" (added during
+  /** Links with this ordinal or higher are "ephemeral" (added during
   forest-building). See the class comment for more information. */
   [[nodiscard]] int num_user_links() const { return data_.num_user_links; }
 
-  /** Joints with this index or higher are "ephemeral" (added during
+  /** Joints with this ordinal or higher are "ephemeral" (added during
   forest-building). See the class comment for more information. */
   [[nodiscard]] int num_user_joints() const { return data_.num_user_joints; }
 
@@ -509,6 +509,86 @@ class LinkJointGraph {
   void ChangeJointType(JointIndex existing_joint_index,
                        const std::string& name_of_new_type);
 
+  /** After the Forest is built, returns the sequence of Links from World to the
+  given Link L in the Forest. In case the Forest was built with a single Mobod
+  for each Link Composite (Links connected by Weld joints), we only report the
+  "active Link" for each Link Composite so that there is only one Link returned
+  for each level in Link L's tree. World is always the active Link for its
+  composite so will always be the first entry in the result. However, if L is
+  part of a Link Composite the final entry will be L's composite's active link,
+  which might not be L. Cost is O(ℓ) where ℓ is Link L's level in the
+  SpanningForest.
+  @@throws std::exception if the SpanningForest hasn't been built yet.
+  @see link_composites(), SpanningForest::FindPathFromWorld() */
+  std::vector<LinkIndex> FindPathFromWorld(LinkIndex link_index) const;
+
+  /** After the Forest is built, finds the first Link common to the paths
+  towards World from each of two Links in the SpanningForest. Returns World
+  immediately if the Links are on different trees of the forest. Otherwise the
+  cost is O(ℓ) where ℓ is the length of the longer path from one of the Links
+  to the ancestor.
+  @throws std::exception if the SpanningForest hasn't been built yet.
+  @see SpanningForest::FindFirstCommonAncestor()
+  @see SpanningForest::FindPathsToFirstCommonAncestor() */
+  LinkIndex FindFirstCommonAncestor(LinkIndex link1_index,
+                                    LinkIndex link2_index) const;
+
+  /** After the Forest is built, finds all the Links following the Forest
+  subtree whose root mobilized body is the one followed by the given Link. That
+  is, we look up the given Link's Mobod B and return all the Links that follow
+  B or any other Mobod in the subtree rooted at B. The Links following B
+  come first, and the rest follow the depth-first ordering of the Mobods.
+  In particular, the result is _not_ sorted by LinkIndex. Computational cost
+  is O(ℓ) where ℓ is the number of Links following the subtree.
+  @throws std::exception if the SpanningForest hasn't been built yet.
+  @see SpanningForest::FindSubtreeLinks() */
+  std::vector<LinkIndex> FindSubtreeLinks(LinkIndex link_index) const;
+
+  /** After the Forest is built, this method can be called to return a
+  partitioning of the LinkJointGraph into subgraphs such that (a) every Link is
+  in one and only one subgraph, and (b) two Links are in the same subgraph iff
+  there is a path between them which consists only of Weld joints.
+  Each subgraph of welded Links is represented as a set of
+  Link indexes (using LinkIndex). By definition, these subgraphs will be
+  disconnected by any non-Weld joint between two Links. A few notes:
+    - The maximum number of returned subgraphs equals the number of Links in
+      the graph. This corresponds to a graph with no Weld joints.
+    - The World Link is included in a set of welded Links, and this set is
+      element zero in the returned vector. The other subgraphs are in no
+      particular order.
+    - The minimum number of subgraphs is one. This corresponds to a graph with
+      all Links welded to World.
+
+  @throws std::exception if the SpanningForest hasn't been built yet.
+  @see CalcSubgraphsOfWeldedLinks() if you haven't built a Forest yet */
+  std::vector<std::set<LinkIndex>> GetSubgraphsOfWeldedLinks() const;
+
+  /** This much-slower method does not depend on the SpanningForest having
+  already been built. It is a fallback for when there is no Forest.
+  @see GetSubgraphsOfWeldedLinks() if you already have a Forest built */
+  std::vector<std::set<LinkIndex>> CalcSubgraphsOfWeldedLinks() const;
+
+  /** After the Forest is built, returns all Links that are transitively welded,
+  or rigidly affixed, to `link_index`, per these two definitions:
+    1. A Link is always considered welded to itself.
+    2. Two unique Links are considered welded together exclusively by the
+       presence of weld joints, not by other constructs such as constraints.
+
+  Therefore, if `link_index` is a valid index to a Link in this graph, then the
+  return vector will always contain at least one entry storing `link_index`.
+  This is fast because we just need to sort the already-calculated
+  LinkComposite the given `link_index` is part of (if any).
+
+  @throws std::exception if the SpanningForest hasn't been built yet or
+                         `link_index` is out of range
+  @see CalcLinksWeldedTo() if you haven't built a Forest yet */
+  std::set<LinkIndex> GetLinksWeldedTo(LinkIndex link_index) const;
+
+  /** This slower method does not depend on the SpanningForest having
+  already been built. It is a fallback for when there is no Forest.
+  @see GetLinksWeldedTo() if you already have a Forest built */
+  std::set<LinkIndex> CalcLinksWeldedTo(LinkIndex link_index) const;
+
   // FYI Debugging APIs (including Graphviz-related) are defined in
   // link_joint_graph_debug.cc.
 
@@ -518,6 +598,7 @@ class LinkJointGraph {
   The result is in the "dot" language, see https://graphviz.org. If you
   write it to some file foo.dot, you can generate a viewable png (for
   example) using the command `dot -Tpng foo.dot >foo.png`.
+  @see SpanningForest::GenerateGraphvizString()
   @see MakeGraphvizFiles() for an easier way to get pngs. */
   std::string GenerateGraphvizString(
       std::string_view label, bool include_ephemeral_elements = true) const;
@@ -700,6 +781,12 @@ class LinkJointGraph {
   void AddUnmodeledJointToComposite(JointOrdinal unmodeled_joint_ordinal,
                                     LinkCompositeIndex which);
 
+  // This is the implementation for CalcLinksWeldedTo().
+  void AppendLinksWeldedTo(LinkIndex link_index,
+                           std::set<LinkIndex>* result) const;
+
+  void ThrowIfForestNotBuiltYet(const char* func) const;
+
   [[noreturn]] void ThrowLinkWasRemoved(const char* func,
                                         LinkIndex link_index) const;
 

multibody/topology/spanning_forest.cc

@@ -1008,6 +1008,99 @@ void SpanningForest::GrowCompositeMobod(
   }
 }
 
+std::vector<MobodIndex> SpanningForest::FindPathFromWorld(
+    MobodIndex index) const {
+  const Mobod* mobod = &mobods(index);
+  std::vector<MobodIndex> path(mobod->level() + 1);
+  while (mobod->inboard().is_valid()) {
+    path[mobod->level()] = mobod->index();
+    mobod = &mobods(mobod->inboard());
+  }
+  DRAKE_DEMAND(mobod->is_world());
+  path[0] = MobodIndex(0);
+  return path;
+}
+
+MobodIndex SpanningForest::FindFirstCommonAncestor(
+    MobodIndex mobod1_index, MobodIndex mobod2_index) const {
+  // A body is its own first common ancestor.
+  if (mobod1_index == mobod2_index) return mobod1_index;
+
+  // If either body is World, that's the first common ancestor.
+  if (mobod1_index == world_mobod_index() ||
+      mobod2_index == world_mobod_index()) {
+    return world_mobod_index();
+  }
+
+  const Mobod* branch1 = &mobods(mobod1_index);
+  const Mobod* branch2 = &mobods(mobod2_index);
+
+  // If they are in different trees, World is the ancestor.
+  if (branch1->tree() != branch2->tree()) return world_mobod().index();
+
+  // Get down to a common level, then go down both branches.
+  while (branch1->level() > branch2->level())
+    branch1 = &mobods(branch1->inboard());
+  while (branch2->level() > branch1->level())
+    branch2 = &mobods(branch2->inboard());
+
+  // Both branches are at the same level now.
+  while (branch1->index() != branch2->index()) {
+    branch1 = &mobods(branch1->inboard());
+    branch2 = &mobods(branch2->inboard());
+  }
+
+  return branch1->index();  // Same as branch2->index().
+}
+
+MobodIndex SpanningForest::FindPathsToFirstCommonAncestor(
+    MobodIndex mobod1_index, MobodIndex mobod2_index,
+    std::vector<MobodIndex>* path1, std::vector<MobodIndex>* path2) const {
+  DRAKE_DEMAND(path1 != nullptr && path2 != nullptr);
+  path1->clear();
+  path2->clear();
+  // A body is its own first common ancestor.
+  if (mobod1_index == mobod2_index) return mobod1_index;
+
+  const Mobod* branch1 = &mobods(mobod1_index);
+  const Mobod* branch2 = &mobods(mobod2_index);
+
+  // Get down to a common level, then go down both branches.
+  while (branch1->level() > branch2->level()) {
+    path1->push_back(branch1->index());
+    branch1 = &mobods(branch1->inboard());
+  }
+  while (branch2->level() > branch1->level()) {
+    path2->push_back(branch2->index());
+    branch2 = &mobods(branch2->inboard());
+  }
+
+  // Both branches are at the same level now.
+  while (branch1->index() != branch2->index()) {
+    path1->push_back(branch1->index());
+    path2->push_back(branch2->index());
+    branch1 = &mobods(branch1->inboard());
+    branch2 = &mobods(branch2->inboard());
+  }
+
+  return branch1->index();  // Same as branch2->index().
+}
+
+std::vector<LinkIndex> SpanningForest::FindSubtreeLinks(
+    MobodIndex root_mobod_index) const {
+  const int num_subtree_mobods = mobods(root_mobod_index).num_subtree_mobods();
+  std::vector<LinkIndex> result;
+  result.reserve(num_subtree_mobods);  // Will be at least this big.
+  for (int i = 0; i < num_subtree_mobods; ++i) {
+    const Mobod& subtree_mobod = mobods(MobodIndex(root_mobod_index + i));
+    for (LinkOrdinal ordinal : subtree_mobod.follower_link_ordinals()) {
+      const Link& link = links(ordinal);
+      result.push_back(link.index());
+    }
+  }
+  return result;
+}
+
 SpanningForest::Data::Data() = default;
 SpanningForest::Data::Data(const Data&) = default;
 SpanningForest::Data::Data(Data&&) = default;