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.

Answer 1

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]