Векторизация этой Numpy двойной петли - PullRequest
Новая
       8

Векторизация этой Numpy двойной петли

5 голосов
/ 28 ноября 2011

Как я могу векторизовать следующий двойной цикл?

У меня есть одна матрица N by A и одна матрица N by B, где A и B могут различаться, а N намного меньше, чем A и B. Iхотите создать матрицу A by B следующим образом, но в идеале без циклов: 

import numpy as np

def foo(arr):
    # can be anything - just an example so that the code runs
    return np.sum(arr)

num_a = 12
num_b = 8
num_dimensions = 3

a = np.random.rand(num_dimensions, num_a)
b = np.random.rand(num_dimensions, num_b)

# this is the loop I want to eliminate:
output = np.zeros( (num_a, num_b) )
for i in xrange(num_a):
    for j in xrange(num_b):
       output[i,j] = foo(a[:,i] - b[:,j])

Есть идеи?

1 Ответ

7 голосов
/ 28 ноября 2011

Первая векторизация foo(), т.е. изменить foo() таким образом, чтобы он мог корректно работать с массивом формы (N, A, B), возвращая массив формы (A, B). Этот шаг обычно трудный. Как это сделать целиком, зависит от того, что делает foo(). Для данного примера это очень легко сделать: 

def foo(arr):
    return np.sum(arr, axis=0)

Теперь используйте правила вещания , чтобы создать массив (N, A, B), содержащий все различия векторов, и передать его в foo(): 

foo(a[:, :, np.newaxis] - b[:, np.newaxis])
...