Making nested 'for' loops more pythonic

Question

I'm relatively new to python and am wondering how to make the following more efficient by avoiding explicit nested 'for' loops and using python's implicit looping instead. I'm working with image data, and in this case trying to speed up my k-means algorithm. Here's a sample of what I'm trying to do:

# shape of image will be something like 140, 150, 3
num_sets, rows_per_set, num_columns = image_values.shape

for set in range(0, num_sets):
    for row in range(0, rows_per_set):
        pos = np.argmin(calc_euclidean(rgb_[set][row], means_list)
        buckets[pos].append(image_values[set][row])

What I have today works great but I'd like to make it more efficient.

Feedback and recommendations are greatly appreciated.

What do you mean by Python's implicit looping? Do you mean the vectorized operations of numpy data structures? — juanpa.arrivillaga
– juanpa.arrivillaga, Commented Mar 28, 2017 at 2:23
unsure of what you mean by python's implicit looping, but you can avoid multiple for loop by using xrange and ` itertools.product` like this: for set, row in itertools.product(xrange(0, num_sets), xrange(0, rows_per_set)) — hemraj
– hemraj, Commented Mar 28, 2017 at 3:21
Hi... thanks for the comments. In this case I am referring to vectorized operations. The idea I had in mind was to make the processing of data run much faster. The 'for' loops are sort of brute force, and I was looking for recommendations on vectorizing to make it more efficient. — James
– James, Commented Mar 28, 2017 at 20:38

Paul Panzer · Accepted Answer · 2017-03-28 03:06:23Z

Here is a vectorised solution. I'm almost certain I got your dimensions muddled up (3 is not really the number of columns, is it?), but the principle should be recognisable anyway:

For demonstration I only collect the (flat) indices into set and row in the buckets.

import numpy as np

k = 6
rgb_=np.random.randint(0, 9, (140, 150, 3))
means_list = np.random.randint(0, 9, (k, 3))

# compute distance table; use some algebra to leverage highly optimised
# dot product
squared_dists = np.add.outer((rgb_*rgb_).sum(axis=-1),
                             (means_list*means_list).sum(axis=-1)) \
    - 2*np.dot(rgb_, means_list.T)
# find best cluster
best = np.argmin(squared_dists, axis=-1)

# find group sizes
counts = np.bincount(best.ravel())
# translate to block boundaries
bnds = np.cumsum(counts[:-1])
# group indices by best cluster; argpartition should be
# a bit cheaper than argsort
chunks = np.argpartition(best.ravel(), bnds)
# split into buckets
buckets = np.split(chunks, bnds)

# check

num_sets, rows_per_set, num_columns = rgb_.shape

def calc_euclidean(a, b):
    return ((a-b)**2).sum(axis=-1)

for set in range(0, num_sets):
    for row in range(0, rows_per_set):
        pos = np.argmin(calc_euclidean(rgb_[set][row], means_list))
        assert pos == best[set, row]
        assert rows_per_set*set+row in buckets[pos]

Collectives™ on Stack Overflow

Making nested 'for' loops more pythonic

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related