From 4343c4830189173a39462452b459f05136a96574 Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Tue, 16 Jan 2024 12:33:17 +0000
Subject: [PATCH] Add generated doc for primitive and core ops to design page

---
 cubed/core/ops.py            |  6 ++++-
 cubed/primitive/blockwise.py |  2 +-
 cubed/primitive/rechunk.py   |  2 +-
 docs/design.md               | 44 ++++++++++++++++++++----------------
 4 files changed, 31 insertions(+), 23 deletions(-)

diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 9c341d2e..9368f8d4 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -296,6 +296,7 @@ def blockwise(
 
 
 def elemwise(func, *args: "Array", dtype=None) -> "Array":
+    """Apply a function elementwise to array arguments, respecting broadcasting."""
     shapes = [arg.shape for arg in args]
     out_ndim = len(np.broadcast_shapes(*shapes))
     expr_inds = tuple(range(out_ndim))[::-1]
@@ -310,6 +311,7 @@ def elemwise(func, *args: "Array", dtype=None) -> "Array":
 
 
 def index(x, key):
+    "Subset an array, along one or more axes."
     if not isinstance(key, tuple):
         key = (key,)
 
@@ -586,7 +588,7 @@ def map_direct(
     func, *args: "Array", shape, dtype, chunks, extra_projected_mem, spec=None, **kwargs
 ) -> "Array":
     """
-    Map a function across blocks of a new array, using side-input arrays to read directly from.
+    Apply a function across blocks of a new array, reading directly from side inputs (not necessarily in a blockwise fashion).
 
     Parameters
     ----------
@@ -736,6 +738,7 @@ def reduction(
     keepdims=False,
     extra_func_kwargs=None,
 ) -> "Array":
+    """Apply a function to reduce an array along one or more axes."""
     if combine_func is None:
         combine_func = func
     if axis is None:
@@ -826,6 +829,7 @@ def reduction(
 
 
 def arg_reduction(x, /, arg_func, axis=None, *, keepdims=False):
+    """A reduction that returns the array indexes, not the values."""
     dtype = np.int64  # index data type
     intermediate_dtype = [("i", dtype), ("v", x.dtype)]
 
diff --git a/cubed/primitive/blockwise.py b/cubed/primitive/blockwise.py
index a99fb149..7bbfa8f1 100644
--- a/cubed/primitive/blockwise.py
+++ b/cubed/primitive/blockwise.py
@@ -108,7 +108,7 @@ def blockwise(
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
     **kwargs,
 ):
-    """Apply a function across blocks from multiple source Zarr arrays.
+    """Apply a function to multiple blocks from multiple inputs, expressed using concise indexing rules.
 
     Parameters
     ----------
diff --git a/cubed/primitive/rechunk.py b/cubed/primitive/rechunk.py
index abd66c03..5c91ca2d 100644
--- a/cubed/primitive/rechunk.py
+++ b/cubed/primitive/rechunk.py
@@ -28,7 +28,7 @@ def rechunk(
     target_store: T_Store,
     temp_store: Optional[T_Store] = None,
 ) -> List[CubedPipeline]:
-    """Rechunk a Zarr array to have target_chunks.
+    """Change the chunking of an array, without changing its shape or dtype.
 
     Parameters
     ----------
diff --git a/docs/design.md b/docs/design.md
index e62f1c00..c969ef95 100644
--- a/docs/design.md
+++ b/docs/design.md
@@ -21,31 +21,35 @@ Cubed uses external runtimes for computation. It follows the Rechunker model (an
 
 There are two primitive operations on arrays:
 
-<dl>
-  <dt><code>blockwise</code></dt>
-  <dd>Applies a function to multiple blocks from multiple inputs, expressed using concise indexing rules.</dd>
-  <dt><code>rechunk</code></dt>
-  <dd>Changes the chunking of an array, without changing its shape or dtype.</dd>
-</dl>
+```{eval-rst}
+.. currentmodule:: cubed.primitive.blockwise
+.. autosummary::
+    :nosignatures:
+
+    blockwise
+.. currentmodule:: cubed.primitive.rechunk
+.. autosummary::
+    :nosignatures:
+
+    rechunk
+```
 
 ## Core operations
 
 These are built on top of the primitive operations, and provide functions that are needed to implement all array operations.
 
-<dl>
-  <dt><code>elemwise</code></dt>
-  <dd>Applies a function elementwise to its arguments, respecting broadcasting.</dd>
-  <dt><code>map_blocks</code></dt>
-  <dd>Applies a function to corresponding blocks from multiple inputs.</dd>
-  <dt><code>map_direct</code></dt>
-  <dd>Applies a function across blocks of a new array, reading directly from side inputs (not necessarily in a blockwise fashion).</dd>
-  <dt><code>index</code> (<code>__getitem__</code>)</dt>
-  <dd>Subsets an array, along one or more axes.</dd>
-  <dt><code>reduction</code></dt>
-  <dd>Applies a function to reduce an array along one or more axes.</dd>
-  <dt><code>arg_reduction</code></dt>
-  <dd>A reduction that returns the array indexes, not the values.</dd>
-</dl>
+```{eval-rst}
+.. currentmodule:: cubed.core.ops
+.. autosummary::
+    :nosignatures:
+
+    elemwise
+    map_blocks
+    map_direct
+    index
+    reduction
+    arg_reduction
+```
 
 ## Array API