Adds misc algorithms over graph and forest

Author: sherm1

multibody/topology/link_joint_graph.cc

@@ -704,6 +704,154 @@ 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<BodyIndex> LinkJointGraph::FindPathFromWorld(
+    BodyIndex link_index) const {
+  ThrowIfForestNotBuiltYet(__func__);
+  const SpanningForest::Mobod* mobod =
+      &forest().mobods()[link_to_mobod(link_index)];
+  std::vector<BodyIndex> 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] = BodyIndex(0);
+  return path;
+}
+
+BodyIndex LinkJointGraph::FindFirstCommonAncestor(BodyIndex link1_index,
+                                                  BodyIndex 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<BodyIndex> LinkJointGraph::FindSubtreeLinks(
+    BodyIndex 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<BodyIndex>> LinkJointGraph::GetSubgraphsOfWeldedLinks()
+    const {
+  ThrowIfForestNotBuiltYet(__func__);
+
+  std::vector<std::set<BodyIndex>> 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<BodyIndex>(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<BodyIndex>{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<BodyIndex>> 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<BodyIndex>> subgraphs{CalcLinksWeldedTo(BodyIndex(0))};
+  for (BodyIndex index : subgraphs.back())
+    visited[index_to_ordinal(index)] = true;
+
+  std::vector<std::set<BodyIndex>> 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<BodyIndex> welded_links = CalcLinksWeldedTo(link.index());
+    for (BodyIndex index : welded_links)
+      visited[index_to_ordinal(index)] = true;
+    auto* which_list = ssize(welded_links) == 1 ? &singletons : &subgraphs;
+    which_list->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<BodyIndex> LinkJointGraph::GetLinksWeldedTo(
+    BodyIndex 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<BodyIndex>{link_index};
+  const std::vector<BodyIndex>& welded_links =
+      link_composites(*composite_index).links;
+  return std::set<BodyIndex>(welded_links.cbegin(), welded_links.cend());
+}
+
+// Without a Forest we don't have LinkComposites available so recursively
+// chase Weld joints instead.
+std::set<BodyIndex> LinkJointGraph::CalcLinksWeldedTo(
+    BodyIndex link_index) const {
+  std::set<BodyIndex> result;
+  AppendLinksWeldedTo(link_index, &result);
+  return result;
+}
+
+void LinkJointGraph::AppendLinksWeldedTo(BodyIndex link_index,
+                                         std::set<BodyIndex>* result) const {
+  DRAKE_DEMAND(result != nullptr);
+  DRAKE_DEMAND(link_index.is_valid());
+  DRAKE_THROW_UNLESS(has_link(link_index));
+  DRAKE_DEMAND(result->count(link_index) == 0);
+
+  const Link& link = link_by_index(link_index);
+
+  // A Link is always considered welded to itself.
+  result->insert(link_index);
+  for (auto joint_index : link.joints()) {
+    const Joint& joint = joint_by_index(joint_index);
+    if (joint.traits_index() != weld_joint_traits_index()) continue;
+    const BodyIndex welded_link_index = joint.other_link_index(link_index);
+    // Don't reprocess if we already did the other end.
+    if (result->count(welded_link_index) == 0)
+      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,
                                          BodyIndex link_index) const {
   throw std::logic_error(fmt::format(

multibody/topology/link_joint_graph.h

@@ -407,11 +407,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; }
 
@@ -499,6 +499,79 @@ 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 called prior to modeling
+  @see link_composites(), SpanningForest::FindPathFromWorld() */
+  std::vector<BodyIndex> FindPathFromWorld(BodyIndex 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. */
+  BodyIndex FindFirstCommonAncestor(BodyIndex link1_index,
+                                    BodyIndex 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 BodyIndex. Computational cost
+  is O(ℓ) where ℓ is the number of Links following the subtree. */
+  std::vector<BodyIndex> FindSubtreeLinks(BodyIndex 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 BodyIndex). 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<BodyIndex>> 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<BodyIndex>> 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 */
+  std::set<BodyIndex> GetLinksWeldedTo(BodyIndex 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. */
+  std::set<BodyIndex> CalcLinksWeldedTo(BodyIndex link_index) const;
+
   // FYI Debugging APIs (including Graphviz-related) are defined in
   // link_joint_graph_debug.cc.
 
@@ -508,6 +581,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;
@@ -690,6 +764,12 @@ class LinkJointGraph {
   void AddUnmodeledJointToComposite(JointOrdinal unmodeled_joint_ordinal,
                                     LinkCompositeIndex which);
 
+  // This is the implementation for CalcLinksWeldedTo().
+  void AppendLinksWeldedTo(BodyIndex link_index,
+                           std::set<BodyIndex>* result) const;
+
+  void ThrowIfForestNotBuiltYet(const char* func) const;
+
   [[noreturn]] void ThrowLinkWasRemoved(const char* func,
                                         BodyIndex link_index) const;