Add a show_hidden option to visualize (#360)

cubed/core/array.py

@@ -163,7 +163,12 @@ def rechunk(self: T_ChunkedArray, chunks) -> T_ChunkedArray:
         return rechunk(self, chunks)
 
     def visualize(
-        self, filename="cubed", format=None, optimize_graph=True, optimize_function=None
+        self,
+        filename="cubed",
+        format=None,
+        optimize_graph=True,
+        optimize_function=None,
+        show_hidden=False,
     ):
         """Produce a visualization of the computation graph for this array.
 
@@ -177,6 +182,12 @@ def visualize(
         optimize_graph : bool, optional
             If True, the graph is optimized before rendering.  Otherwise,
             the graph is displayed as is. Default is True.
+        optimize_function : callable, optional
+            Function to optimize the computation graph. Default is to use a
+            function that performs map fusion only.
+        show_hidden : bool
+            If True, show arrays and operations that are marked as hidden in the
+            visualization. Default is False.
 
         Returns
         -------
@@ -190,6 +201,7 @@ def visualize(
             format=format,
             optimize_graph=optimize_graph,
             optimize_function=optimize_function,
+            show_hidden=show_hidden,
         )
 
     def __getitem__(self: T_ChunkedArray, key, /) -> T_ChunkedArray:
@@ -246,7 +258,12 @@ def compute(
 
 
 def visualize(
-    *arrays, filename="cubed", format=None, optimize_graph=True, optimize_function=None
+    *arrays,
+    filename="cubed",
+    format=None,
+    optimize_graph=True,
+    optimize_function=None,
+    show_hidden=False,
 ):
     """Produce a visualization of the computation graph for multiple arrays.
 
@@ -262,6 +279,12 @@ def visualize(
     optimize_graph : bool, optional
         If True, the graph is optimized before rendering.  Otherwise,
         the graph is displayed as is. Default is True.
+    optimize_function : callable, optional
+        Function to optimize the computation graph. Default is to use a
+        function that performs map fusion only.
+    show_hidden : bool
+        If True, show arrays and operations that are marked as hidden in the
+        visualization. Default is False.
 
     Returns
     -------
@@ -275,6 +298,7 @@ def visualize(
         format=format,
         optimize_graph=optimize_graph,
         optimize_function=optimize_function,
+        show_hidden=show_hidden,
     )
 
 

cubed/core/plan.py

@@ -250,17 +250,18 @@ def visualize(
         rankdir="TB",
         optimize_graph=True,
         optimize_function=None,
+        show_hidden=False,
     ):
         dag = self._finalize_dag(optimize_graph, optimize_function)
         dag = dag.copy()  # make a copy since we mutate the DAG below
 
         # remove edges from create-arrays output node to avoid cluttering the diagram
         dag.remove_edges_from(list(dag.out_edges("arrays")))
 
-        # remove hidden nodes
-        dag.remove_nodes_from(
-            list(n for n, d in dag.nodes(data=True) if d.get("hidden", False))
-        )
+        if not show_hidden:
+            dag.remove_nodes_from(
+                list(n for n, d in dag.nodes(data=True) if d.get("hidden", False))
+            )
 
         dag.graph["graph"] = {
             "rankdir": rankdir,

cubed/tests/test_core.py

@@ -421,8 +421,8 @@ def test_compute_multiple_different_specs(tmp_path):
 
 
 def test_visualize(tmp_path):
-    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
-    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2))
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=xp.float64, chunks=(2, 2))
+    b = cubed.random.random((3, 3), chunks=(2, 2))
     c = xp.add(a, b)
     d = c * 2
     e = c * 3
@@ -438,6 +438,9 @@ def test_visualize(tmp_path):
     e.visualize(filename=tmp_path / "e")
     assert (tmp_path / "e.svg").exists()
 
+    e.visualize(filename=tmp_path / "e-hidden", show_hidden=True)
+    assert (tmp_path / "e-hidden.svg").exists()
+
     e.visualize(filename=tmp_path / "e", format="png")
     assert (tmp_path / "e.png").exists()