Repository: cubed-dev/cubed
Pull request #566 (2024-09-04 13:19:56)

Coerce args to `map_blocks` to arrays
This change makes this aspect of `map_blocks` more consistent with Dask.


Patch:
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index e589c4e2..36fd9dc0 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -578,6 +578,13 @@ def map_blocks(
 ) -> "Array":
     """Apply a function to corresponding blocks from multiple input arrays."""
 
+    from cubed.array_api.creation_functions import asarray
+
+    # Coerce all args to Cubed arrays
+    specs = [a.spec for a in args if hasattr(a, "spec")]
+    spec0 = specs[0] if len(specs) > 0 else spec
+    args = tuple(asarray(a, spec=spec0) for a in args)
+
     # Handle the case where an array is created by calling `map_blocks` with no input arrays
     if len(args) == 0:
         from cubed.array_api.creation_functions import empty_virtual_array
diff --git a/cubed/diagnostics/tqdm.py b/cubed/diagnostics/tqdm.py
index 03ef4ac9..7cbc9102 100644
--- a/cubed/diagnostics/tqdm.py
+++ b/cubed/diagnostics/tqdm.py
@@ -17,19 +17,31 @@ class TqdmProgressBar(Callback):
     def on_compute_start(self, event):
         from tqdm.auto import tqdm
 
+        # find the maximum display width so we can align bars below
+        max_op_display_name = (
+            max(
+                len(node["op_display_name"].replace("\n", " "))
+                for _, node in visit_nodes(event.dag, event.resume)
+            )
+            + 1  # for the colon
+        )
+
         self.pbars = {}
-        i = 0
-        for name, node in visit_nodes(event.dag, event.resume):
+        for i, (name, node) in enumerate(visit_nodes(event.dag, event.resume)):
             num_tasks = node["primitive_op"].num_tasks
-            op_display_name = node["op_display_name"].replace("\n", " ")
+            op_display_name = node["op_display_name"].replace("\n", " ") + ":"
+            # note double curlies to get literal { and } for tqdm bar format
+            bar_format = (
+                f"{{desc:{max_op_display_name}}} {{percentage:3.0f}}%|{{bar}}{{r_bar}}"
+            )
             self.pbars[name] = tqdm(
                 *self.args,
                 desc=op_display_name,
                 total=num_tasks,
                 position=i,
+                bar_format=bar_format,
                 **self.kwargs,
             )
-            i = i + 1
 
     def on_compute_end(self, event):
         for pbar in self.pbars.values():


Test patch:
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 4e092e96..a4070237 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -235,6 +235,13 @@ def test_map_blocks_with_different_block_shapes(spec):
     assert_array_equal(c.compute(), np.array([[[12, 13]]]))
 
 
+def test_map_blocks_with_non_cubed_array(spec):
+    a = xp.arange(10, dtype="int64", chunks=(2,), spec=spec)
+    b = np.array([1, 2], dtype="int64")  # numpy array will be coerced to cubed
+    c = cubed.map_blocks(nxp.add, a, b, dtype="int64")
+    assert_array_equal(c.compute(), np.array([1, 3, 3, 5, 5, 7, 7, 9, 9, 11]))
+
+
 def test_multiple_ops(spec, executor):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)