Repository: cubed-dev/cubed
Pull request #550 (2024-08-12 13:53:02)

More lenient dtype support
Some array API implementations are more lenient in the dtypes they allow for certain operations - e.g. `sum` can be used on arrays of `bool` (with return type `int64`). The array API doesn't rule out this behaviour - it's just not a part of the spec.

This is needed to improve Xarray integration, see https://github.com/cubed-dev/cubed-xarray/issues/8


Patch:
diff --git a/cubed/array_api/array_object.py b/cubed/array_api/array_object.py
index b3e8abe9..d7363241 100644
--- a/cubed/array_api/array_object.py
+++ b/cubed/array_api/array_object.py
@@ -399,12 +399,18 @@ class Array(CoreArray):
     # Utility methods
 
     def _check_allowed_dtypes(self, other, dtype_category, op):
-        if self.dtype not in _dtype_categories[dtype_category]:
+        if (
+            dtype_category != "all"
+            and self.dtype not in _dtype_categories[dtype_category]
+        ):
             raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
         if isinstance(other, (int, complex, float, bool)):
             other = self._promote_scalar(other)
         elif isinstance(other, CoreArray):
-            if other.dtype not in _dtype_categories[dtype_category]:
+            if (
+                dtype_category != "all"
+                and other.dtype not in _dtype_categories[dtype_category]
+            ):
                 raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
         else:
             return NotImplemented
diff --git a/cubed/array_api/statistical_functions.py b/cubed/array_api/statistical_functions.py
index 78ff6ae2..7ee6525e 100644
--- a/cubed/array_api/statistical_functions.py
+++ b/cubed/array_api/statistical_functions.py
@@ -1,6 +1,7 @@
 import math
 
 from cubed.array_api.dtypes import (
+    _boolean_dtypes,
     _numeric_dtypes,
     _real_floating_dtypes,
     _real_numeric_dtypes,
@@ -124,10 +125,13 @@ def min(x, /, *, axis=None, keepdims=False, use_new_impl=True, split_every=None)
 def prod(
     x, /, *, axis=None, dtype=None, keepdims=False, use_new_impl=True, split_every=None
 ):
-    if x.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in prod")
+    # boolean is allowed by numpy
+    if x.dtype not in _numeric_dtypes and x.dtype not in _boolean_dtypes:
+        raise TypeError("Only numeric or boolean dtypes are allowed in prod")
     if dtype is None:
-        if x.dtype in _signed_integer_dtypes:
+        if x.dtype in _boolean_dtypes:
+            dtype = int64
+        elif x.dtype in _signed_integer_dtypes:
             dtype = int64
         elif x.dtype in _unsigned_integer_dtypes:
             dtype = uint64
@@ -153,10 +157,13 @@ def prod(
 def sum(
     x, /, *, axis=None, dtype=None, keepdims=False, use_new_impl=True, split_every=None
 ):
-    if x.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in sum")
+    # boolean is allowed by numpy
+    if x.dtype not in _numeric_dtypes and x.dtype not in _boolean_dtypes:
+        raise TypeError("Only numeric or boolean dtypes are allowed in sum")
     if dtype is None:
-        if x.dtype in _signed_integer_dtypes:
+        if x.dtype in _boolean_dtypes:
+            dtype = int64
+        elif x.dtype in _signed_integer_dtypes:
             dtype = int64
         elif x.dtype in _unsigned_integer_dtypes:
             dtype = uint64
diff --git a/cubed/diagnostics/timeline.py b/cubed/diagnostics/timeline.py
index 1c0abc89..64de7551 100644
--- a/cubed/diagnostics/timeline.py
+++ b/cubed/diagnostics/timeline.py
@@ -1,6 +1,6 @@
-import os
 import time
 from dataclasses import asdict
+from pathlib import Path
 from typing import Optional
 
 import matplotlib.patches as mpatches
@@ -16,7 +16,7 @@ pylab.switch_backend("Agg")
 
 
 class TimelineVisualizationCallback(Callback):
-    def __init__(self, format: Optional[str] = None) -> None:
+    def __init__(self, format: Optional[str] = "svg") -> None:
         self.format = format
 
     def on_compute_start(self, event):
@@ -27,75 +27,62 @@ class TimelineVisualizationCallback(Callback):
         self.stats.append(asdict(event))
 
     def on_compute_end(self, event):
-        end_tstamp = time.time()
-        dst = f"history/{event.compute_id}"
-        format = self.format
-        create_timeline(self.stats, self.start_tstamp, end_tstamp, dst, format)
+        self.end_tstamp = time.time()
 
-
-# copy of lithops function of the same name, and modified for different field names
-def create_timeline(stats, start_tstamp, end_tstamp, dst=None, format=None):
-    stats_df = pd.DataFrame(stats)
-
-    stats_df = stats_df.sort_values(by=["task_create_tstamp", "name"], ascending=True)
-
-    total_calls = len(stats_df)
-
-    palette = sns.color_palette("deep", 6)
-
-    fig = pylab.figure(figsize=(10, 6))
-    ax = fig.add_subplot(1, 1, 1)
-
-    y = np.arange(total_calls)
-    point_size = 10
-
-    fields = [
-        ("task create", stats_df.task_create_tstamp - start_tstamp),
-        ("function start", stats_df.function_start_tstamp - start_tstamp),
-        ("function end", stats_df.function_end_tstamp - start_tstamp),
-        ("task result", stats_df.task_result_tstamp - start_tstamp),
-    ]
-
-    patches = []
-    for f_i, (field_name, val) in enumerate(fields):
-        ax.scatter(val, y, c=[palette[f_i]], edgecolor="none", s=point_size, alpha=0.8)
-        patches.append(mpatches.Patch(color=palette[f_i], label=field_name))
-
-    ax.set_xlabel("Execution Time (sec)")
-    ax.set_ylabel("Function Call")
-
-    legend = pylab.legend(handles=patches, loc="upper right", frameon=True)
-    legend.get_frame().set_facecolor("#FFFFFF")
-
-    yplot_step = int(np.max([1, total_calls / 20]))
-    y_ticks = np.arange(total_calls // yplot_step + 2) * yplot_step
-    ax.set_yticks(y_ticks)
-    ax.set_ylim(-0.02 * total_calls, total_calls * 1.02)
-    for y in y_ticks:
-        ax.axhline(y, c="k", alpha=0.1, linewidth=1)
-
-    max_seconds = np.max(end_tstamp - start_tstamp) * 1.25
-    xplot_step = max(int(max_seconds / 8), 1)
-    x_ticks = np.arange(max_seconds // xplot_step + 2) * xplot_step
-    ax.set_xlim(0, max_seconds)
-
-    ax.set_xticks(x_ticks)
-    for x in x_ticks:
-        ax.axvline(x, c="k", alpha=0.2, linewidth=0.8)
-
-    ax.grid(False)
-    fig.tight_layout()
-
-    if format is None:
-        format = "svg"
-
-    if dst is None:
-        os.makedirs("plots", exist_ok=True)
-        dst = os.path.join(
-            os.getcwd(), "plots", "{}_{}".format(int(time.time()), f"timeline.{format}")
+        stats_df = pd.DataFrame(self.stats)
+        stats_df = stats_df.sort_values(
+            by=["task_create_tstamp", "name"], ascending=True
         )
-    else:
-        dst = os.path.expanduser(dst) if "~" in dst else dst
-        dst = "{}/{}".format(os.path.realpath(dst), f"timeline.{format}")
-
-    fig.savefig(dst)
+        total_calls = len(stats_df)
+        palette = sns.color_palette("deep", 6)
+
+        fig = pylab.figure(figsize=(10, 6))
+        ax = fig.add_subplot(1, 1, 1)
+
+        y = np.arange(total_calls)
+        point_size = 10
+
+        fields = [
+            ("task create", stats_df.task_create_tstamp - self.start_tstamp),
+            ("function start", stats_df.function_start_tstamp - self.start_tstamp),
+            ("function end", stats_df.function_end_tstamp - self.start_tstamp),
+            ("task result", stats_df.task_result_tstamp - self.start_tstamp),
+        ]
+
+        patches = []
+        for f_i, (field_name, val) in enumerate(fields):
+            ax.scatter(
+                val, y, c=[palette[f_i]], edgecolor="none", s=point_size, alpha=0.8
+            )
+            patches.append(mpatches.Patch(color=palette[f_i], label=field_name))
+
+        ax.set_xlabel("Execution Time (sec)")
+        ax.set_ylabel("Function Call")
+
+        legend = pylab.legend(handles=patches, loc="upper right", frameon=True)
+        legend.get_frame().set_facecolor("#FFFFFF")
+
+        yplot_step = int(np.max([1, total_calls / 20]))
+        y_ticks = np.arange(total_calls // yplot_step + 2) * yplot_step
+        ax.set_yticks(y_ticks)
+        ax.set_ylim(-0.02 * total_calls, total_calls * 1.02)
+        for y in y_ticks:
+            ax.axhline(y, c="k", alpha=0.1, linewidth=1)
+
+        max_seconds = np.max(self.end_tstamp - self.start_tstamp) * 1.25
+        xplot_step = max(int(max_seconds / 8), 1)
+        x_ticks = np.arange(max_seconds // xplot_step + 2) * xplot_step
+        ax.set_xlim(0, max_seconds)
+
+        ax.set_xticks(x_ticks)
+        for x in x_ticks:
+            ax.axvline(x, c="k", alpha=0.2, linewidth=0.8)
+
+        ax.grid(False)
+        fig.tight_layout()
+
+        self.dst = Path(f"history/{event.compute_id}")
+        self.dst.mkdir(parents=True, exist_ok=True)
+        self.dst = self.dst / f"timeline.{self.format}"
+
+        fig.savefig(self.dst)
diff --git a/docs/user-guide/diagnostics.md b/docs/user-guide/diagnostics.md
index 3e013323..278a6f1a 100644
--- a/docs/user-guide/diagnostics.md
+++ b/docs/user-guide/diagnostics.md
@@ -41,7 +41,12 @@ Array `c` is coloured orange, which means it is materialized as a Zarr array. Ar
 
 Similarly, the operation that produces `c` is shown in a lilac colour to signify that it runs tasks to produce the output. Operations `op-001` and `op-002` don't run any tasks since `a` and `b` are just small constant arrays.
 
-## Progress bar
+
+## Callbacks
+
+You can pass callbacks to functions that call `compute`, such as {py:func}`store <cubed.store>` or {py:func}`to_zarr <cubed.to_zarr>`.
+
+### Progress bar
 
 You can display a progress bar to track your computation by passing callbacks to {py:meth}`compute() <cubed.Array.compute()>`:
 
@@ -53,14 +58,37 @@ You can display a progress bar to track your computation by passing callbacks to
   op-003 add    4/4 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00
 ```
 
+The two current progress bar choice are:
+- `from cubed.diagnostics.rich import RichProgressBar`
+- `from cubed.diagnostics.tqdm import TqdmProgressBar`
+
+
 This will work in Jupyter notebooks, and for all executors.
 
-You can also pass callbacks to functions that call `compute`, such as {py:func}`store <cubed.store>` or {py:func}`to_zarr <cubed.to_zarr>`.
 
-## History and timeline visualization
 
-The history and timeline visualization callbacks can be used to find out how long tasks took to run, and how much memory they used.
+### History
+The history callback can be used to understand how long tasks took to run, and how much memory they used. The history callback will write [`events.csv`, `plan.csv` and `stats.csv`] to a new directory under the current directory with the schema `history/compute-{id}`.
+
 
+```ipython
+>>> from cubed.diagnostics.history import HistoryCallback
+>>> hist = HistoryCallback()
+>>> c.compute(callbacks=[hist])
+```
+
+
+### Timeline
 The timeline visualization is useful to determine how much time was spent in worker startup, as well as how much stragglers affected the overall time of the computation. (Ideally, we want vertical lines on this plot, which would represent perfect horizontal scaling.)
 
+The timeline callback will write a graphic `timeline.svg` to a directory with the schema `history/compute-{id}`.
+
+
+```ipython
+>>> from cubed.diagnostics.timeline import TimelineVisualizationCallback
+>>> timeline_viz = TimelineVisualizationCallback()
+>>> c.compute(callbacks=[timeline_viz])
+```
+
+### Examples in use
 See the [examples](https://github.com/cubed-dev/cubed/blob/main/examples/README.md) for more information about how to use them.


Test patch:
diff --git a/cubed/tests/test_types.py b/cubed/tests/test_types.py
new file mode 100644
index 00000000..b21eca6e
--- /dev/null
+++ b/cubed/tests/test_types.py
@@ -0,0 +1,10 @@
+from numpy.testing import assert_array_equal
+
+import cubed.array_api as xp
+
+
+# This is less strict than the spec, but is supported by implementations like NumPy
+def test_prod_sum_bool():
+    a = xp.ones((2,), dtype=xp.bool)
+    assert_array_equal(xp.prod(a).compute(), xp.asarray([1], dtype=xp.int64))
+    assert_array_equal(xp.sum(a).compute(), xp.asarray([2], dtype=xp.int64))