Curlew Curlew - 7 months ago 6
Python Question

Python - Apply a function over a labeled multidimensional array

I have a

numpy
array that is labelled using
scipy
connected component labelling.

import numpy
from scipy import ndimage

a = numpy.zeros((8,8), dtype=numpy.int)
a[1,1] = a[1,2] = a[2,1] = a[2,2] = a[3,1] = a[3,2] = 1
a[5,5] = a[5,6] = a[6,5] = a[6,6] = a[7,5] = a[7,6] = 1
lbl, numpatches = ndimage.label(a)


I want to apply a custom function (calculation of a specific value) over all labels within the labelled array.
Similar as for instance the ndimage algebra functions:

ndimage.sum(a,lbl,range(1,numpatches+1))


( Which in this case returns me the number of values for each label
[6,6]
. )

Is there a way to do this?

Answer

You can pass an arbitrary function to ndimage.labeled_comprehension, which is roughly equivalent to

[func(a[lbl == i]) for i in index]

Here is the labeled_comprehension-equivalent of ndimage.sum(a,lbl,range(1,numpatches+1)):

import numpy as np
from scipy import ndimage

a = np.zeros((8,8), dtype=np.int)
a[1,1] = a[1,2] = a[2,1] = a[2,2] = a[3,1] = a[3,2] = 1
a[5,5] = a[5,6] = a[6,5] = a[6,6] = a[7,5] = a[7,6] = 1 
lbl, numpatches = ndimage.label(a)

def func(x):
    return x.sum()

print(ndimage.labeled_comprehension(a, lbl, index=range(1, numpatches+1), 
                                    func=func, out_dtype='float', default=None))
# [6 6]