Added buffer_size= option to apply_over_* functions for convenience.

This just gets passed along to the block size/shape chooser functions if
the block size/shape is not explicitly supplied to the apply* function.

Also renamed memory= to buffer_size as it is more self-explanatory.

Author: LTLA

src/delayedarray/DelayedArray.py

@@ -16,7 +16,7 @@
 from .extract_dense_array import extract_dense_array
 from .to_dense_array import to_dense_array
 from .extract_sparse_array import extract_sparse_array
-from .apply_over_blocks import apply_over_blocks, choose_block_shape_for_iteration
+from .apply_over_blocks import apply_over_blocks
 from .create_dask_array import create_dask_array
 from .chunk_shape import chunk_shape
 from .is_sparse import is_sparse
@@ -934,19 +934,18 @@ def _reduce_SparseNdarray(x: SparseNdarray, multipliers: List[int], axes: List[i
 
 
 def _reduce(x: DelayedArray, axes: List[int], operation: Callable, buffer_size: int):
-    block_shape = choose_block_shape_for_iteration(x, memory = buffer_size)
     multipliers = _create_offset_multipliers(x.shape, axes)
     if is_sparse(x):
         apply_over_blocks(
             x, 
             lambda position, block : _reduce_SparseNdarray(block, multipliers, axes, position, operation), 
-            block_shape=block_shape, 
-            allow_sparse=True
+            buffer_size=buffer_size,
+            allow_sparse=True,
         )
     else:
         apply_over_blocks(
             x, 
             lambda position, block : _reduce_ndarray(block, multipliers, axes, position, operation), 
-            block_shape=block_shape
+            buffer_size=buffer_size,
         )
     return

src/delayedarray/apply_over_blocks.py

@@ -11,7 +11,7 @@
 __license__ = "MIT"
 
 
-def choose_block_shape_for_iteration(x, memory: int = 10000000) -> Tuple[int, ...]:
+def choose_block_shape_for_iteration(x, buffer_size: int = 10000000) -> Tuple[int, ...]:
     """
     Choose the block dimensions for blockwise iteration through an array, see
     `~apply_over_blocks` for details.
@@ -21,7 +21,9 @@ def choose_block_shape_for_iteration(x, memory: int = 10000000) -> Tuple[int, ..
 
         dimension: Dimension to iterate over.
 
-        memory: Available memory in bytes, to hold a single block in memory.
+        buffer_size: 
+            Buffer_size in bytes, to hold a single block per iteration. Larger
+            values generally improve speed at the cost of memory.
 
     Returns:
         Dimensions of the blocks. All values are guaranteed to be positive,
@@ -32,7 +34,7 @@ def choose_block_shape_for_iteration(x, memory: int = 10000000) -> Tuple[int, ..
         if d == 0:
             return (*(max(1, d) for d in x.shape),)
 
-    num_elements = memory / x.dtype.itemsize
+    num_elements = buffer_size / x.dtype.itemsize
     chunk_dims = chunk_shape(x)
     block_size = 1
     for s in chunk_dims:
@@ -64,7 +66,7 @@ def choose_block_shape_for_iteration(x, memory: int = 10000000) -> Tuple[int, ..
     return (*block_dims,)
 
 
-def apply_over_blocks(x, fun: Callable, block_shape: Optional[Tuple] = None, allow_sparse: bool = False) -> list:
+def apply_over_blocks(x, fun: Callable, block_shape: Optional[Tuple] = None, allow_sparse: bool = False, buffer_size: int = 1e8) -> list:
     """
     Iterate over an array by blocks. We apply a user-provided function and
     collect the results before proceeding to the next block.
@@ -88,11 +90,16 @@ def apply_over_blocks(x, fun: Callable, block_shape: Optional[Tuple] = None, all
             ``x`` contains a sparse array, the block contents are instead
             represented by a :py:class:`~SparseNdarray.SparseNdarray`.
 
+        buffer_size: 
+            Buffer_size in bytes, to hold a single block per iteration. Larger
+            values generally improve speed at the cost of memory. Only used
+            if ``block_shape`` is not provided.
+
     Returns:
         List containing the output of ``fun`` on each block.
     """
     if block_shape is None:
-        block_shape = choose_block_shape_for_iteration(x)
+        block_shape = choose_block_shape_for_iteration(x, buffer_size = buffer_size)
 
     num_tasks_total = 1
     num_tasks_by_dim = []

src/delayedarray/apply_over_dimension.py

@@ -18,7 +18,7 @@ def guess_iteration_block_size(x, dimension, memory: int = 10000000) -> int:
     return choose_block_size_for_1d_iteration(x, dimension, memory)
 
 
-def choose_block_size_for_1d_iteration(x, dimension: int, memory: int = 10000000) -> int:
+def choose_block_size_for_1d_iteration(x, dimension: int, buffer_size: int = 10000000) -> int:
     """
     Choose a block size for iterating over an array on a certain dimension,
     see `~apply_over_dimension` for more details.
@@ -28,7 +28,9 @@ def choose_block_size_for_1d_iteration(x, dimension: int, memory: int = 10000000
 
         dimension: Dimension to iterate over.
 
-        memory: Available memory in bytes, to hold a single block in memory.
+        buffer_size: 
+            Buffer_size in bytes, to hold a single block per iteration. Larger
+            values generally improve speed at the cost of memory.
 
     Returns:
         Size of the block on the iteration dimension. This is guaranteed to be
@@ -46,7 +48,7 @@ def choose_block_size_for_1d_iteration(x, dimension: int, memory: int = 10000000
         if i != dimension:
             prod_other *= s
 
-    num_elements = memory / x.dtype.itemsize
+    num_elements = buffer_size / x.dtype.itemsize
     ideal = int(num_elements / prod_other)
     if ideal == 0:
         return 1
@@ -59,7 +61,7 @@ def choose_block_size_for_1d_iteration(x, dimension: int, memory: int = 10000000
     return int(ideal / curdim) * curdim
 
 
-def apply_over_dimension(x, dimension: int, fun: Callable, block_size: Optional[int] = None, allow_sparse: bool = False) -> list:
+def apply_over_dimension(x, dimension: int, fun: Callable, block_size: Optional[int] = None, allow_sparse: bool = False, buffer_size: int = 1e8) -> list:
     """
     Iterate over an array on a certain dimension. At each iteration, the block
     of observations consists of the full extent of all dimensions other than
@@ -87,11 +89,16 @@ def apply_over_dimension(x, dimension: int, fun: Callable, block_size: Optional[
             ``x`` contains a sparse array, the block contents are instead
             represented by a :py:class:`~SparseNdarray.SparseNdarray`.
 
+        buffer_size: 
+            Buffer_size in bytes, to hold a single block per iteration. Larger
+            values generally improve speed at the cost of memory. Only used
+            if ``block_size`` is not provided.
+
     Returns:
         List containing the output of ``fun`` on each block.
     """
     if block_size is None:
-        block_size = choose_block_size_for_1d_iteration(x, dimension)
+        block_size = choose_block_size_for_1d_iteration(x, dimension, buffer_size = buffer_size)
 
     limit = x.shape[dimension]
     tasks = math.ceil(limit / block_size)

tests/test_apply_over_blocks.py

@@ -27,21 +27,21 @@ def chunk_shape_ChunkyBoi(x: _ChunkyBoi):
 
 def test_choose_block_shape_for_iteration():
     x = np.random.rand(100, 10)
-    assert da.choose_block_shape_for_iteration(x, memory=200) == (2, 10)
-    assert da.choose_block_shape_for_iteration(x, memory=800) == (10, 10)
+    assert da.choose_block_shape_for_iteration(x, buffer_size=200) == (2, 10)
+    assert da.choose_block_shape_for_iteration(x, buffer_size=800) == (10, 10)
 
-    # Not enough memory. 
-    assert da.choose_block_shape_for_iteration(x, memory=0) == (1, 1)
-    assert da.choose_block_shape_for_iteration(x, memory=40) == (1, 5)
+    # Not enough buffer_size. 
+    assert da.choose_block_shape_for_iteration(x, buffer_size=0) == (1, 1)
+    assert da.choose_block_shape_for_iteration(x, buffer_size=40) == (1, 5)
 
     # Behaves correctly with empty objects.
     empty = np.random.rand(100, 0)
     assert da.choose_block_shape_for_iteration(empty) == (100, 1)
 
     x = _ChunkyBoi((100, 200), (20, 25))
-    assert da.choose_block_shape_for_iteration(x, memory=4000) == (20, 25)
-    assert da.choose_block_shape_for_iteration(x, memory=40000) == (100, 50)
-    assert da.choose_block_shape_for_iteration(x, memory=80000) == (100, 100)
+    assert da.choose_block_shape_for_iteration(x, buffer_size=4000) == (20, 25)
+    assert da.choose_block_shape_for_iteration(x, buffer_size=40000) == (100, 50)
+    assert da.choose_block_shape_for_iteration(x, buffer_size=80000) == (100, 100)
 
 
 def _dense_sum(position, block):

tests/test_apply_over_dimension.py

@@ -27,12 +27,12 @@ def chunk_shape_ChunkyBoi(x: _ChunkyBoi):
 
 def test_choose_block_size_for_1d_iteration():
     x = np.random.rand(100, 10)
-    assert da.choose_block_size_for_1d_iteration(x, 0, memory=800) == 10
-    assert da.choose_block_size_for_1d_iteration(x, 1, memory=800) == 1
+    assert da.choose_block_size_for_1d_iteration(x, 0, buffer_size=800) == 10
+    assert da.choose_block_size_for_1d_iteration(x, 1, buffer_size=800) == 1
 
-    # No memory.
-    assert da.choose_block_size_for_1d_iteration(x, 0, memory=0) == 1
-    assert da.choose_block_size_for_1d_iteration(x, 1, memory=0) == 1
+    # No buffer_size.
+    assert da.choose_block_size_for_1d_iteration(x, 0, buffer_size=0) == 1
+    assert da.choose_block_size_for_1d_iteration(x, 1, buffer_size=0) == 1
 
     # Behaves correctly with empty objects.
     empty = np.random.rand(100, 0)
@@ -41,10 +41,10 @@ def test_choose_block_size_for_1d_iteration():
 
     # Making a slightly more complex situation.
     x = _ChunkyBoi((100, 200), (20, 25))
-    assert da.choose_block_size_for_1d_iteration(x, 0, memory=4000) == 2
-    assert da.choose_block_size_for_1d_iteration(x, 1, memory=4000) == 5
-    assert da.choose_block_size_for_1d_iteration(x, 0, memory=40000) == 20
-    assert da.choose_block_size_for_1d_iteration(x, 1, memory=40000) == 50
+    assert da.choose_block_size_for_1d_iteration(x, 0, buffer_size=4000) == 2
+    assert da.choose_block_size_for_1d_iteration(x, 1, buffer_size=4000) == 5
+    assert da.choose_block_size_for_1d_iteration(x, 0, buffer_size=40000) == 20
+    assert da.choose_block_size_for_1d_iteration(x, 1, buffer_size=40000) == 50
 
 
 def _dense_sum(position, block):