Самый быстрый способ заполнить матрицу функцией пар элементов в двух векторах numpy?

У меня есть два одномерных вектора numpy va и vb, которые используются для заполнения матрицы путем передачи всех парных комбинаций в функцию.

na = len(va)
nb = len(vb)
D = np.zeros((na, nb))
for i in range(na):
    for j in range(nb):
        D[i, j] = foo(va[i], vb[j])

В нынешнем виде этот фрагмент кода выполняется очень долго из-за того, что va и vb относительно велики (4626 и 737). Однако я надеюсь, что это можно улучшить из-за того, что аналогичная процедура выполняется с использованием метода cdist из scipy с очень хорошей производительностью.

D = cdist(va, vb, metric)

Я, очевидно, знаю, что у scipy есть преимущество запуска этого фрагмента кода на C, а не на python, но я надеюсь, что есть какая-то функция numpy, о которой я не знаю, которая может выполнить это быстро.


person Michael Aquilina    schedule 31.10.2014    source источник
comment
Используйте векторизованные функции, которые будут обрабатывать все элементы va и vb одновременно, используя функции Numpy outer... Или передавая сетку va и vb...   -  person Saullo G. P. Castro    schedule 31.10.2014
comment
Проблема в том, что эта функция является пользовательской, и изменение векторизации нетривиально. Я пробовал использовать np.meshgrid, а затем np.vectorize, но улучшения производительности были минимальными.   -  person Michael Aquilina    schedule 31.10.2014
comment
vectorize ничего не делает с внутренностями вашего foo. Это просто оболочка, которая в конечном итоге вызывает foo(a,b) с каждой парой скаляров. Это удобство, а не инструмент ускорения.   -  person hpaulj    schedule 31.10.2014


Ответы (3)


Одна из наименее известных функций numpy для того, что в документации называется процедурами функционального программирования — это np.frompyfunc. Это создает numpy ufunc из функции Python. Не какой-то другой объект, который точно имитирует пустой ufunc, а настоящий ufunc со всеми его наворотами. Хотя поведение во многих аспектах очень похоже на np.vectorize, оно имеет некоторые явные преимущества, которые, надеюсь, следует выделить в следующем коде:

In [2]: def f(a, b):
   ...:     return a + b
   ...:

In [3]: f_vec = np.vectorize(f)

In [4]: f_ufunc = np.frompyfunc(f, 2, 1)  # 2 inputs, 1 output

In [5]: a = np.random.rand(1000)

In [6]: b = np.random.rand(2000)

In [7]: %timeit np.add.outer(a, b)  # a baseline for comparison
100 loops, best of 3: 9.89 ms per loop

In [8]: %timeit f_vec(a[:, None], b)  # 50x slower than np.add
1 loops, best of 3: 488 ms per loop

In [9]: %timeit f_ufunc(a[:, None], b)  # ~20% faster than np.vectorize...
1 loops, best of 3: 425 ms per loop

In [10]: %timeit f_ufunc.outer(a, b)  # ...and you get to use ufunc methods
1 loops, best of 3: 427 ms per loop

Таким образом, хотя он все еще явно уступает правильно векторизованной реализации, он немного быстрее (зацикливание на C, но у вас все еще есть накладные расходы на вызов функции Python).

person Jaime    schedule 01.11.2014
comment
Похоже, это именно то, что я искал, я проверю это завтра и дам вам знать, как все прошло :) - person Michael Aquilina; 02.11.2014

cdist работает быстро, потому что он написан на высокооптимизированном коде C (как вы уже указали), и он поддерживает только небольшой предопределенный набор metric.

Поскольку вы хотите применить операцию в общем случае к любой заданной foo функции, у вас нет другого выбора, кроме как вызвать эту функцию na-раз-nb раз. Эта часть вряд ли поддается дальнейшей оптимизации.

Осталось оптимизировать циклы и индексацию. Некоторые предложения, чтобы попробовать:

  1. Используйте xrange вместо range (если в python2.x. в python3 диапазон уже похож на генератор)
  2. Используйте enumerate вместо диапазона + явное индексирование
  3. Используйте «магию» скорости Python, например cython или numba, чтобы ускорить процесс зацикливания.

Если вы можете сделать дополнительные предположения о foo, возможно, удастся еще больше ускорить его.

person shx2    schedule 31.10.2014

Как сказал @shx2, все зависит от того, что такое foo. Если вы можете выразить это с помощью numpy ufuncs, используйте метод outer:

In [11]: N = 400

In [12]: B = np.empty((N, N))

In [13]: x = np.random.random(N)

In [14]: y = np.random.random(N)

In [15]: %%timeit
for i in range(N):
   for j in range(N):
     B[i, j] = x[i] - y[j]
   ....: 
10 loops, best of 3: 87.2 ms per loop

In [16]: %timeit A = np.subtract.outer(x, y)   # <--- np.subtract is a ufunc
1000 loops, best of 3: 294 µs per loop

В противном случае вы можете опустить цикл до уровня cython. Продолжая тривиальный пример выше:

In [45]: %%cython
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
def foo(double[::1] x, double[::1] y, double[:, ::1] out):
    cdef int i, j
    for i in xrange(x.shape[0]):
        for j in xrange(y.shape[0]):
            out[i, j] = x[i] - y[j]
   ....: 

In [46]: foo(x, y, B)

In [47]: np.allclose(B, np.subtract.outer(x, y))
Out[47]: True

In [48]: %timeit foo(x, y, B)
10000 loops, best of 3: 149 µs per loop

Пример cython намеренно сделан слишком упрощенным: на самом деле вы можете добавить некоторые проверки формы/шага, выделить память внутри своей функции и т. д.

person ev-br    schedule 31.10.2014
comment
Ну, у меня есть два случая запуска этого кода. Один вычисляет расстояние Жаккара между двумя векторами. Я знаю, что у scipy есть функция для этого, но в нашем случае поведение немного отличается. Другой выполняет поиск в словаре D[i, j] = a.get((va[i], vb[j]), 1.0) - person Michael Aquilina; 31.10.2014
comment
Ну, это зависит и должно быть измерено для вашей конкретной установки, для которой %timeit является вашим другом. Если вы где-то застряли в пути, лучше задайте отдельный вопрос с подробностями. Удачи! - person ev-br; 31.10.2014