Add num_blocks

This adds the `num_blocks` function to compute the number of blocks in
each dimension given a known data shape and block shape.

kenjutsu/blocks.py

@@ -6,12 +6,91 @@
 
 
 import itertools
+import math
 import operator
 import warnings
 
 import kenjutsu.format
 
 
+def num_blocks(space_shape, block_shape):
+    """
+        Computes the number of blocks.
+
+        Takes an array with ``space_shape`` and ``block_shape`` for every
+        dimension. From this, it can compute slicings to use for cutting
+        each block out from the original array, HDF5 dataset, or other.
+
+
+        Args:
+            space_shape(tuple):            Shape of array to slice
+            block_shape(tuple):            Size of each block to take
+
+        Returns:
+            tuple:                         Number of blocks per dimension
+
+        Examples:
+
+            >>> num_blocks(
+            ...     (2, 3,), (1, 1,)
+            ... )  #doctest: +NORMALIZE_WHITESPACE
+            (2, 3)
+
+    """
+
+    try:
+        irange = xrange
+    except NameError:
+        irange = range
+
+    try:
+        from itertools import ifilter, imap
+    except ImportError:
+        ifilter, imap = filter, map
+
+    if not (len(space_shape) == len(block_shape)):
+        raise ValueError(
+            "The dimensions of `space_shape` and `block_shape` should be"
+            " the same."
+        )
+
+    vec_type = lambda t, a: imap(t, a)
+
+    vec_ceil = lambda a: imap(math.ceil, a)
+
+    vec_div = lambda a, b: imap(operator.truediv, a, b)
+    vec_mod = lambda a, b: imap(operator.mod, a, b)
+
+    vec_nonzero = lambda a: \
+            imap(lambda _: _[0], ifilter(lambda _: _[1], enumerate(a)))
+    vec_str = lambda a: imap(str, a)
+
+    uneven_block_division = tuple(vec_mod(space_shape, block_shape))
+
+    if any(uneven_block_division):
+        uneven_block_division_str = vec_nonzero(uneven_block_division)
+        uneven_block_division_str = vec_str(uneven_block_division_str)
+        uneven_block_division_str = ", ".join(uneven_block_division_str)
+
+        warnings.warn(
+            "Blocks will not evenly divide the array." +
+            " The following dimensions will be unevenly divided: %s." %
+            uneven_block_division_str,
+            RuntimeWarning
+        )
+
+    for each_dim in irange(len(space_shape)):
+        # Construct each block using the block size given. Allow to spill over.
+        if block_shape[each_dim] == -1:
+            block_shape = (block_shape[:each_dim] +
+                           space_shape[each_dim:each_dim+1] +
+                           block_shape[each_dim+1:])
+
+    n_blocks = vec_type(int, vec_ceil(vec_div(space_shape, block_shape)))
+
+    return tuple(n_blocks)
+
+
 def split_blocks(space_shape, block_shape, block_halo=None):
     """
         Return a list of slicings to cut each block out of an array or other.

tests/test_blocks.py

@@ -26,6 +26,47 @@ def setUp(self):
         pass
 
 
+    def test_num_blocks(self):
+        with self.assertRaises(ValueError) as e:
+            blocks.num_blocks((1,), (1, 2))
+
+        self.assertEqual(
+            str(e.exception),
+            "The dimensions of `space_shape` and `block_shape` should be the"
+            " same."
+        )
+
+        result = blocks.num_blocks((2,), (1,))
+        self.assertEqual(
+            result,
+            (2,)
+        )
+
+        result = blocks.num_blocks((2,), (-1,))
+        self.assertEqual(
+            result,
+            (1,)
+        )
+
+        result = blocks.num_blocks((2, 3), (1, 1,))
+        self.assertEqual(
+            result,
+            (2, 3)
+        )
+
+        result = blocks.num_blocks((2, 3), (1, 2,))
+        self.assertEqual(
+            result,
+            (2, 2)
+        )
+
+        result = blocks.num_blocks((10, 12), (3, 2,))
+        self.assertEqual(
+            result,
+            (4, 6)
+        )
+
+
     def test_split_blocks(self):
         with self.assertRaises(ValueError) as e:
             blocks.split_blocks((1,), (1, 2), (1, 2, 3))