Add ability to save graph as a dot graph (#114)

swri-robotics · Dec 31, 2024 · 49ef1fe · 49ef1fe
1 parent e1b1c1c
commit 49ef1fe
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 2 deletions.
diff --git a/descartes_light/core/include/descartes_light/solvers/ladder_graph/impl/ladder_graph.hpp b/descartes_light/core/include/descartes_light/solvers/ladder_graph/impl/ladder_graph.hpp
@@ -20,6 +20,7 @@
 
 #include <descartes_light/solvers/ladder_graph/ladder_graph.h>
 #include <iomanip>
+#include <fstream>
 
 namespace descartes_light
 {
@@ -128,6 +129,74 @@ void LadderGraph<FloatType>::clear()
   rungs_.clear();
 }
 
+template <typename FloatType>
+void LadderGraph<FloatType>::toDotGraph(std::string filepath,
+                                        const std::vector<unsigned>& path) const
+{
+  std::ofstream out(filepath);
+  if (!out)
+    return;
+
+  out << "digraph ladder {\n";
+  out << "  rankdir=LR;\n";
+  out << "  node [shape=box];\n";
+
+  // Create subgraphs for each rung
+  for (size_t rung_idx = 0; rung_idx < rungs_.size(); ++rung_idx)
+  {
+    out << "  subgraph cluster_" << rung_idx << " {\n";
+    out << "    label=\"Rung " << rung_idx << "\";\n";
+    out << "    style=dotted;\n";
+
+    const auto& rung = rungs_[rung_idx];
+    for (size_t node_idx = 0; node_idx < rung.nodes.size(); ++node_idx)
+    {
+      std::string node_name = "r" + std::to_string(rung_idx) + "n" + std::to_string(node_idx);
+
+      // Fix: Properly escape quotes and handle label closing
+      out << "    " << node_name << " [label=\"" << node_idx;
+      if (rung.nodes[node_idx].sample.cost >= 0)
+        out << "\\nCost: " << rung.nodes[node_idx].sample.cost;
+
+      out << "\\nState: ";
+      if (rung.nodes[node_idx].sample.state)
+        for (Eigen::Index i = 0; i < rung.nodes[node_idx].sample.state->values.size(); ++i)
+          out << std::fixed << std::setprecision(3) << rung.nodes[node_idx].sample.state->values[i] << " ";
+
+      if (!path.empty() && rung_idx < path.size() && path[rung_idx] == node_idx)
+        out << "\",color=green,style=filled];\n";
+      else
+        out << "\"];\n";
+    }
+    out << "  }\n";
+
+    // Create edges to next rung
+    if (rung_idx < rungs_.size() - 1)
+    {
+      for (size_t node_idx = 0; node_idx < rung.nodes.size(); ++node_idx)
+      {
+        const auto& node = rung.nodes[node_idx];
+        std::string from_node = "r" + std::to_string(rung_idx) + "n" + std::to_string(node_idx);
+
+        for (const auto& edge : node.edges)
+        {
+          std::string to_node = "r" + std::to_string(rung_idx + 1) + "n" + std::to_string(edge.idx);
+          // Fix: Properly format edge labels
+          if (!path.empty() && rung_idx + 1 < path.size() && path[rung_idx] == node_idx &&
+              path[rung_idx + 1] == edge.idx)
+            out << "  " << from_node << " -> " << to_node << " [label=\"" << edge.cost
+                << "\",color=green,penwidth=2.0];\n";
+          else
+            out << "  " << from_node << " -> " << to_node << " [label=\"" << edge.cost << "\"];\n";
+        }
+      }
+    }
+  }
+
+  out << "}\n";
+  out.close();
+}
+
 }  // namespace descartes_light
 
 #endif  // DESCARTES_LIGHT_IMPL_LADDER_GRAPH_HPP
diff --git a/descartes_light/core/include/descartes_light/solvers/ladder_graph/ladder_graph.h b/descartes_light/core/include/descartes_light/solvers/ladder_graph/ladder_graph.h
@@ -26,6 +26,7 @@ DESCARTES_IGNORE_WARNINGS_PUSH
 #include <vector>
 #include <Eigen/Geometry>
 #include <iomanip>
+#include <iosfwd>
 DESCARTES_IGNORE_WARNINGS_POP
 
 namespace descartes_core
@@ -157,6 +158,8 @@ class LadderGraph
    */
   void clear();
 
+  void toDotGraph(std::string filepath, const std::vector<unsigned>& path = std::vector<unsigned>()) const;
+
   friend std::ostream& operator<<(std::ostream& out, const LadderGraph<FloatType>& ladder_graph)
   {
     out << "\nRung\t(Nodes)\t|# Outgoing Edges|\n";
@@ -190,11 +193,16 @@ class LadderGraph
         Node<FloatType> node2 = rung2.nodes.front();
         typename State<FloatType>::ConstPtr state2 = node2.sample.state;
         out << "Rung # " << failed_id << "\n";
+        out << "\tNode 1\t|\tNode 2\t|\tDiff (Node2 - Node1)\n";
+        out << "\t---------------------------------------\n";
         for (Eigen::Index i = 0; i < state1->values.rows(); i++)
         {
-          out << std::setprecision(4) << std::fixed << "\t" << state1->values[i] << "\t|\t" << state2->values[i]
-              << "\n";
+            out << std::setprecision(4) << std::fixed 
+          << "\t" << state1->values[i] 
+          << "\t|\t" << state2->values[i] 
+          << "\t|\t" << state2->values[i] - state1->values[i] << "\n";
         }
+        out << "\tTotal L2 Norm: " << (state2->values - state1->values).norm() << "\n";
       }
     }
     out << "\n";