Making nested 'for' loops more pythonic
I'm relatively new to python and I'm wondering how to make the following more efficient by avoiding explicit nested for loops and using the implicit python loop. I am working with image data, in which case I am trying to speed up my k-means algorithm. Here's an example 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 would like to make it more efficient.
Feedback and recommendations are appreciated.
+3
source to share
1 answer
Here's a vector solution. I'm pretty sure I've got your sizes confusing (3 is not really the number of columns, is it?), But the principle should be recognizable anyway:
For demonstration, I am collecting (flat) indexes into a set and a string in 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]
+1
source to share