Generally, I'm trying to split a distance matrix into K folds. Specifically, for the 3 x 3 case, my distance matrix might look like this:
full = np.array([
[0, 0, 3],
[1, 0, 1],
[2, 1, 0]
])
I also have a list of randomly generated assignments, the length of which is equal to the sum over all elements in the distance matrix. For the K = 3 case, it might look like this:
assignments = np.array([0, 1, 0, 2, 1, 1, 0, 0])
I want to create K = 3 new 3 x 3 matrices of zeros, in which the values of the distance matrix are "distributed" according to the assignments list. Code is more precise than words, so here's my current attempt:
def assign(full, assignments):
folds = [np.zeros(full.shape) for _ in xrange(np.max(assignments) + 1)]
rows, cols = full.shape
a = 0
for r in xrange(rows):
for c in xrange(cols):
for i in xrange(full[r, c]):
folds[assignments[a]][r, c] += 1
a += 1
return folds
This works (slowly), and in this example,
folds = assign(full, assignments)
for f in folds:
print f
returns
[[ 0. 0. 2.]
[ 0. 0. 0.]
[ 1. 1. 0.]]
[[ 0. 0. 1.]
[ 0. 0. 1.]
[ 1. 0. 0.]]
[[ 0. 0. 0.]
[ 1. 0. 0.]
[ 0. 0. 0.]]
as desired. However, my current attempt is very slow, especially for the N x N case for N large. How can I improve the speed of this function? Is there some numpy magic that I should be using here?
One idea I had was converting to a sparse matrix and looping over nonzero entries. This would only help a bit, however,