Rewrite center_of_mass using labeled_comprehension (#78)

* Add a custom kernel for `center_of_mass`

Create a custom kernel of `center_of_mass` to use with
`labeled_comprehension`. In the kernel take the ROI's intensity and 1-D
positions along with a few arguments to be bound before calling
`labeled_comprehension` (like the shape of the original image and the
results expected type). Use these to unravel the 1-D positions back into
N-D positions and computed the center of mass for the given ROI to pass
back in an expected form. As `labeled_comprehension` can only handle
scalars being returned for each ROI, use a structured type to pass back
each coordinate of the center of mass in a different field. The
structured type uses a simple naming scheme based on a string holding
the index corresponding to that dimension.

* Use `labeled_comprehension` in `center_of_mass`

Rewrite `center_of_mass` to use `labeled_comprehension` along with a
custom kernel. This simplifies the Dask graph significantly, improves
the performance noticeably, and cleans up the code a bit as the whole
center of mass computation is now handled by the custom kernel function.
Only a little bit of bookkeeping is needed to bind a few extra arguments
(since `labeled_comprehension` does not handle keyword arguments) to the
custom kernel and to convert the custom structured array type (needed
since `labeled_comprehension` cannot handle results with a shape) into a
normal array before returning it user.

dask_image/ndmeasure/__init__.py

@@ -1,8 +1,5 @@
 # -*- coding: utf-8 -*-
 
-from __future__ import division
-
-
 __author__ = """John Kirkham"""
 __email__ = "kirkhamj@janelia.hhmi.org"
 
@@ -53,22 +50,22 @@ def center_of_mass(input, labels=None, index=None):
     # This only matters if index is some array.
     index = index.T
 
-    input_mtch_sum = sum(input, labels, index)
-
-    input_wt_mtch_sum = []
-    for i in _pycompat.irange(input.ndim):
-        sl = input.ndim * [None]
-        sl[i] = slice(None)
-        sl = tuple(sl)
-
-        input_i = dask.array.arange(input.shape[i], chunks=input.chunks[i])
-        input_wt = input * input_i[sl]
+    type_mapping = collections.OrderedDict([
+        (("%i" % i), input.dtype) for i in _pycompat.irange(input.ndim)
+    ])
+    out_dtype = numpy.dtype(list(type_mapping.items()))
 
-        input_wt_mtch_sum.append(sum(input_wt, labels, index))
+    default_1d = numpy.full((1,), numpy.nan, dtype=out_dtype)
 
-    input_wt_mtch_sum = dask.array.stack(input_wt_mtch_sum, axis=-1)
+    func = functools.partial(
+        _utils._center_of_mass, shape=input.shape, dtype=out_dtype
+    )
+    com_lbl = labeled_comprehension(
+        input, labels, index,
+        func, out_dtype, default_1d[0], pass_positions=True
+    )
 
-    com_lbl = input_wt_mtch_sum / input_mtch_sum[..., None]
+    com_lbl = dask.array.stack([com_lbl[k] for k in type_mapping], axis=-1)
 
     return com_lbl
 

dask_image/ndmeasure/_utils.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 
 
+from __future__ import division
+
 import operator
 import warnings
 
@@ -128,6 +130,26 @@ def _argmin(a, positions):
     return positions[numpy.argmin(a)]
 
 
+def _center_of_mass(a, positions, shape, dtype):
+    """
+    Find the center of mass for each ROI.
+
+    Package the result in a structured array with each field as an index.
+    """
+
+    result = numpy.empty((1,), dtype=dtype)
+
+    positions_nd = numpy.unravel_index(positions, shape)
+    a_sum = numpy.sum(a)
+
+    a_wt_i = numpy.empty_like(a)
+    for i, pos_nd_i in enumerate(positions_nd):
+        a_wt_sum_i = numpy.multiply(a, pos_nd_i, out=a_wt_i).sum()
+        result[("%i" % i)] = a_wt_sum_i / a_sum
+
+    return result[0]
+
+
 def _extrema(a, positions, dtype):
     """
     Find minimum and maximum as well as positions for both.