Improve detection of array names in visualize (#835)

* Find arrays in Xarray datasets
* Fix bug where cubed_xarray wasn't being considered
* Filter out variable names starting with _
* Show literal value of scalar arrays

Author: tomwhite

cubed/array_api/creation_functions.py

@@ -72,7 +72,8 @@ def asarray(
     name = gensym()
     target = virtual_in_memory(a, chunks=chunksize)
 
-    plan = Plan._new(name, "asarray", target)
+    scalar_value = str(a) if a.shape == () else None
+    plan = Plan._new(name, "asarray", target, scalar_value=scalar_value)
     return Array(name, target, spec, plan)
 
 

cubed/core/ops.py

@@ -395,6 +395,7 @@ def blockwise(
         op.target_array,
         op,
         False,
+        None,
         *source_arrays,
     )
     from cubed.array_api import Array
@@ -555,6 +556,7 @@ def _general_blockwise(
         op.target_array,
         op,
         False,
+        None,
         *source_arrays,
     )
     from cubed.array_api import Array
@@ -975,6 +977,7 @@ def rechunk(x, chunks, *, target_store=None, min_mem=None, use_new_impl=True):
             op.target_array,
             op,
             False,
+            None,
             x,
         )
         return Array(name, op.target_array, spec, plan)
@@ -987,6 +990,7 @@ def rechunk(x, chunks, *, target_store=None, min_mem=None, use_new_impl=True):
             op1.target_array,
             op1,
             False,
+            None,
             x,
         )
         x_int = Array(name_int, op1.target_array, spec, plan1)
@@ -998,6 +1002,7 @@ def rechunk(x, chunks, *, target_store=None, min_mem=None, use_new_impl=True):
             op2.target_array,
             op2,
             False,
+            None,
             x_int,
         )
         return Array(name, op2.target_array, spec, plan2)

cubed/core/plan.py

@@ -20,6 +20,7 @@
 from cubed.storage.zarr import LazyZarrArray, open_if_lazy_zarr_array
 from cubed.utils import (
     chunk_memory,
+    extract_array_names_from_stack_summaries,
     extract_stack_summaries,
     is_local_path,
     itemsize,
@@ -90,6 +91,7 @@ def _new(
         target,
         primitive_op=None,
         hidden=False,
+        scalar_value=None,
         *source_arrays,
     ):
         # create an empty DAG or combine from sources
@@ -137,6 +139,7 @@ def _new(
                     type="array",
                     target=target,
                     hidden=hidden,
+                    scalar_value=scalar_value,
                 )
                 dag.add_edge(op_name_unique, name)
         else:
@@ -582,21 +585,17 @@ def visualize(
         dag.graph["node"] = {"fontname": "helvetica", "shape": "box", "fontsize": "10"}
 
         # do an initial pass to extract array variable names from stack summaries
-        array_display_names = {}
+        stacks = []
         for _, d in dag.nodes(data=True):
             if "stack_summaries" in d:
                 stack_summaries = d["stack_summaries"]
-                first_cubed_i = min(
-                    i for i, s in enumerate(stack_summaries) if s.is_cubed()
-                )
-                caller_summary = stack_summaries[first_cubed_i - 1]
-                array_display_names.update(caller_summary.array_names_to_variable_names)
+                stacks.append(stack_summaries)
         # add current stack info
-        frame = inspect.currentframe().f_back  # go back one in the stack
+        # go back one in the stack to the caller of 'visualize'
+        frame = inspect.currentframe().f_back
         stack_summaries = extract_stack_summaries(frame, limit=10)
-        first_cubed_i = min(i for i, s in enumerate(stack_summaries) if s.is_cubed())
-        caller_summary = stack_summaries[first_cubed_i - 1]
-        array_display_names.update(caller_summary.array_names_to_variable_names)
+        stacks.append(stack_summaries)
+        array_display_names = extract_array_names_from_stack_summaries(stacks)
 
         # now set node attributes with visualization info
         for n, d in dag.nodes(data=True):
@@ -680,6 +679,10 @@ def visualize(
                     var_name = array_display_names[n]
                     label = f"{n}\n{var_name}"
                     tooltip += f"variable: {var_name}\n"
+                if "scalar_value" in d and d["scalar_value"] is not None:
+                    scalar_value = d["scalar_value"]
+                    label += f"\n{scalar_value}"
+                    tooltip += f"value: {scalar_value}\n"
                 tooltip += f"shape: {target.shape}\n"
                 tooltip += f"chunks: {target.chunks}\n"
                 tooltip += f"dtype: {target.dtype}\n"

cubed/utils.py

@@ -6,14 +6,15 @@
 import sys
 import sysconfig
 import traceback
-from collections.abc import Iterator
+from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
 from functools import partial
 from itertools import islice
 from math import prod
 from operator import add, mul
 from pathlib import Path
 from posixpath import join
+from types import FrameType
 from typing import Dict, Optional, Tuple, Union, cast
 from urllib.parse import quote, unquote, urlsplit, urlunsplit
 
@@ -169,22 +170,24 @@ class StackSummary:
     module: str
     array_names_to_variable_names: dict
 
-    def is_cubed(self):
+    def is_cubed(self) -> bool:
         """Return True if this stack frame is a Cubed call."""
-        return self.module.startswith("cubed.") and not self.module.startswith(
-            "cubed.tests."
-        )
+        return (
+            self.module.startswith("cubed.") or self.module.startswith("cubed_xarray.")
+        ) and not self.module.startswith("cubed.tests.")
 
-    def is_on_python_lib_path(self):
+    def is_on_python_lib_path(self) -> bool:
         """Return True if this stack frame is from a library on Python's library path."""
         python_lib_path = sysconfig.get_path("purelib")
 
         return self.filename.startswith(python_lib_path)
 
 
-def extract_stack_summaries(frame, limit=None):
+def extract_stack_summaries(
+    frame: Optional[FrameType], limit: Optional[int] = None
+) -> list[StackSummary]:
     """Like Python's ``StackSummary.extract``, but returns module information."""
-    frame_gen = traceback.walk_stack(frame)
+    frame_gen: Iterable[tuple[FrameType, int]] = traceback.walk_stack(frame)
 
     # from StackSummary.extract
     if limit is None:
@@ -214,24 +217,46 @@ def extract_stack_summaries(frame, limit=None):
     return stack_summaries
 
 
-def extract_array_names(frame):
+def extract_array_names(frame: FrameType) -> dict[str, str]:
     """Look for Cubed arrays in local variables to create a mapping from (internally generated) array names to variable names."""
 
     from cubed import Array
 
-    array_names_to_variable_names = {}
-    for var, arr in frame.f_locals.items():
-        if isinstance(arr, Array):
-            array_names_to_variable_names[arr.name] = var
+    array_names_to_variable_names: dict[str, str] = {}
+    for name, obj in frame.f_locals.items():
+        if name.startswith("_"):
+            # ignore previous output variables like _, __, and _1 (IPython)
+            continue
+        if isinstance(obj, Array):
+            array_names_to_variable_names[obj.name] = name
+        elif (
+            type(obj).__module__.split(".")[0] == "xarray"
+            and obj.__class__.__name__ == "DataArray"
+        ):
+            if isinstance(obj.data, Array):
+                array_names_to_variable_names[obj.data.name] = name
         elif (
-            type(arr).__module__.split(".")[0] == "xarray"
-            and arr.__class__.__name__ == "DataArray"
+            type(obj).__module__.split(".")[0] == "xarray"
+            and obj.__class__.__name__ == "Dataset"
         ):
-            if isinstance(arr.data, Array):
-                array_names_to_variable_names[arr.data.name] = arr.name
+            for var in obj.data_vars:
+                da = obj[var]
+                if isinstance(da.data, Array):
+                    array_names_to_variable_names[da.data.name] = f"{name}.{var}"
     return array_names_to_variable_names
 
 
+def extract_array_names_from_stack_summaries(
+    stacks: list[list[StackSummary]],
+) -> dict[str, str]:
+    array_display_names: dict[str, str] = {}
+    for stack_summaries in stacks:
+        first_cubed_i = min(i for i, s in enumerate(stack_summaries) if s.is_cubed())
+        caller_summary = stack_summaries[first_cubed_i - 1]
+        array_display_names.update(caller_summary.array_names_to_variable_names)
+    return array_display_names
+
+
 def convert_to_bytes(size: Union[int, float, str]) -> int:
     """
     Converts the input data size to bytes.