# -*- coding: utf-8 -*-

from __future__ import division


__author__ = """John Kirkham"""
__email__ = "kirkhamj@janelia.hhmi.org"


from ._version import get_versions
__version__ = get_versions()['version']
del get_versions


import operator

import numpy
import dask.array

from . import _compat


def center_of_mass(input, labels=None, index=None):
    """
    Calculate the center of mass of the values of an array at labels.

    Parameters
    ----------
    input : ndarray
        Data from which to calculate center-of-mass. The masses can either
        be positive or negative.
    labels : ndarray, optional
        Labels for objects in `input`, as generated by `ndimage.label`.
        Only used with `index`.  Dimensions must be the same as `input`.
    index : int or sequence of ints, optional
        Labels for which to calculate centers-of-mass. If not specified,
        all labels greater than zero are used.  Only used with `labels`.

    Returns
    -------
    center_of_mass : tuple, or list of tuples
        Coordinates of centers-of-mass.
    """

    input = _compat._asarray(input)

    if labels is None:
        labels = (input != 0).astype(numpy.int64)
        index = None

    if index is None:
        labels = (labels > 0).astype(numpy.int64)
        index = dask.array.ones(tuple(), dtype=numpy.int64, chunks=tuple())

    labels = _compat._asarray(labels)
    index = _compat._asarray(index)

    # SciPy transposes these for some reason.
    # So we do the same thing here.
    # This only matters if index is some array.
    index = index.T

    if input.shape != labels.shape:
        raise ValueError("The input and labels arrays must be the same shape.")

    input_i = _compat._indices(input.shape, chunks=input.chunks)

    lbl_mtch = operator.eq(
        index[(Ellipsis,) + labels.ndim * (None,)],
        labels[index.ndim * (None,)]
    )

    input_mtch = lbl_mtch.astype(input.dtype) * input[index.ndim * (None,)]

    input_mtch_ind_wt = (
        input_mtch[index.ndim * (slice(None),) + (None,)] *
        input_i[index.ndim * (None,)]
    )

    input_mtch = input_mtch.astype(numpy.float64)
    input_mtch_ind_wt = input_mtch_ind_wt.astype(numpy.float64)

    com_lbl = input_mtch_ind_wt.sum(
        axis=tuple(range(1 + index.ndim, input_mtch_ind_wt.ndim))
    )
    input_mtch_sum = input_mtch.sum(
        axis=tuple(range(index.ndim, input_mtch.ndim))
    )
    com_lbl /= input_mtch_sum[index.ndim * (slice(None),) + (None,)]

    return com_lbl

@@ -1905,38 +1905,43 @@
 h = 
-lbl_mtch.astype(input.dtype) *
+dask.array.where(%0A        lbl_mtch,
  inp
@@ -1960,24 +1960,51 @@
 m * (None,)%5D
+, input.dtype.type(0)%0A    )
 %0A%0A    input_

Use Dask Array's where instead of multiplying