From 16e65d522bc8ef47eb77918c85aad028c2bdd647 Mon Sep 17 00:00:00 2001
From: zachary62 <zhuang333@wisc.edu>
Date: Wed, 8 Jan 2025 00:02:32 +0000
Subject: [PATCH] simplify viz

---
 docs/viz.md | 95 ++++++++++++++---------------------------------------
 1 file changed, 25 insertions(+), 70 deletions(-)

diff --git a/docs/viz.md b/docs/viz.md
index 28dce6d..c78862d 100644
--- a/docs/viz.md
+++ b/docs/viz.md
@@ -7,68 +7,35 @@ nav_order: 3
 
 # Visualization
 
-Visualizing the flow and node structure can greatly help understanding. While we do **not** include built-in visualization tools, we provide an example of how to create one using **Mermaid**.
+Visualizing the nested graph can help understanding. While we **don’t** include built-in visualization tools, we provide an example using **Mermaid**.
 
 ### Example: Visualization of Node with Mermaid
 
 {% raw %}
 ```python
-# Generate Mermaid diagram code for a Flow or Node structure
 def build_mermaid(start):
-    node_ids, visited = {}, set()
-    lines = ["graph LR"]
-    counter = [1]  # Counter for unique node IDs
-
-    def get_id(node):
-        if node not in node_ids:
-            node_ids[node] = f"N{counter[0]}"
-            counter[0] += 1
-        return node_ids[node]
-
-    def connect(src, tgt):
-        lines.append(f"    {src} --> {tgt}")
-
-    def walk(node, parent_id=None):
+    ids, visited, lines = {}, set(), ["graph LR"]
+    ctr = 1
+    def get_id(n):
+        nonlocal ctr
+        return ids[n] if n in ids else (ids.setdefault(n, f"N{ctr}"), (ctr := ctr + 1))[0]
+    def link(a, b):
+        lines.append(f"    {a} --> {b}")
+    def walk(node, parent=None):
         if node in visited:
-            if parent_id:
-                connect(parent_id, get_id(node))
-            return
+            return parent and link(parent, get_id(node))
         visited.add(node)
-
         if isinstance(node, Flow):
-            # Handle Flow nodes
-            if parent_id and node.start:
-                connect(parent_id, get_id(node.start))
-
+            node.start and parent and link(parent, get_id(node.start))
             lines.append(f"\n    subgraph sub_flow_{get_id(node)}[{type(node).__name__}]")
-
-            if node.start:
-                walk(node.start)
-
-            for succ_node in node.successors.values():
-                if node.start:
-                    walk(succ_node, get_id(node.start))
-                else:
-                    if parent_id:
-                        connect(parent_id, get_id(succ_node))
-                    walk(succ_node)
-
+            node.start and walk(node.start)
+            for nxt in node.successors.values():
+                node.start and walk(nxt, get_id(node.start)) or (parent and link(parent, get_id(nxt))) or walk(nxt)
             lines.append("    end\n")
-
         else:
-            # Handle simple Nodes
-            curr_id = get_id(node)
-            label = f'{curr_id}["{type(node).__name__}"]'
-
-            if parent_id:
-                lines.append(f"    {label}")
-                connect(parent_id, curr_id)
-            else:
-                lines.append(f"    {label}")
-
-            for succ_node in node.successors.values():
-                walk(succ_node, curr_id)
-
+            lines.append(f"    {(nid := get_id(node))}['{type(node).__name__}']")
+            parent and link(parent, nid)
+            [walk(nxt, nid) for nxt in node.successors.values()]
     walk(start)
     return "\n".join(lines)
 ```
@@ -76,40 +43,28 @@ def build_mermaid(start):
 
 ### Usage Example
 
-Here, we define some example Nodes and Flows to generate a Mermaid diagram:
+Here, we define some example Nodes and Flows:
 
 ```python
-class DataPrepBatchNode(BatchNode):
-    pass
-
-class ValidateDataNode(Node):
-    pass
-
-class FeatureExtractionNode(Node):
-    pass
-
-class TrainModelNode(Node):
-    pass
-
-class EvaluateModelNode(Node):
-    pass
-
-class ModelFlow(Flow):
-    pass
+class DataPrepBatchNode(BatchNode): pass
+class ValidateDataNode(Node): pass
+class FeatureExtractionNode(Node): pass
+class TrainModelNode(Node): pass
+class EvaluateModelNode(Node): pass
+class ModelFlow(Flow): pass
 
 feature_node = FeatureExtractionNode()
 train_node = TrainModelNode()
 evaluate_node = EvaluateModelNode()
 feature_node >> train_node >> evaluate_node
 model_flow = ModelFlow(start=feature_node)
-
 data_prep_node = DataPrepBatchNode()
 validate_node = ValidateDataNode()
 data_prep_node >> validate_node >> model_flow
 build_mermaid(start=data_prep_node)
 ```
 
-The above code produces a Mermaid diagram (e.g., use the [Mermaid Live Editor](https://mermaid.live/) to render it):
+The code generates a Mermaid diagram:
 
 ```mermaid
 graph LR