Numpy quirk: применить функцию ко всем парам двух массивов 1D, чтобы получить один 2D-массив

Скажем, у меня есть 2 одномерных (1D) массива numpy, a и b, с длиной n1 и n2 соответственно. У меня также есть функция F(x,y), которая принимает два значения. Теперь я хочу применить эту функцию к каждой паре значений из двух моих массивов 1D, поэтому результатом будет массив 2D numpy с формой n1, n2. Элементом i, j двумерного массива будет F(a[i], b[j]).

Я не смог найти способ сделать это без ужасного количества for-loops, и я уверен, что там намного проще (и быстрее!) сделать это в numpy.

Спасибо заранее!

Ответы

Ответ 1

Вы можете использовать numpy broadcasting для вычисления на двух массивах, превращая a в вертикальный 2D-массив с помощью newaxis:

In [11]: a = np.array([1, 2, 3]) # n1 = 3
    ...: b = np.array([4, 5]) # n2 = 2
    ...: #if function is c(i, j) = a(i) + b(j)*2:
    ...: c = a[:, None] + b*2

In [12]: c
Out[12]: 
array([[ 9, 11],
       [10, 12],
       [11, 13]])

Для сравнения:

In [28]: a = arange(100)

In [29]: b = arange(222)

In [30]: timeit r = np.array([[f(i, j) for j in b] for i in a])
10 loops, best of 3: 29.9 ms per loop

In [31]: timeit c = a[:, None] + b*2
10000 loops, best of 3: 71.6 us per loop

Ответ 2

Если F находится вне вашего контроля, вы можете автоматически его обернуть, чтобы быть "векторным", используя numpy.vectorize. Ниже представлен рабочий пример, где я определяю свой собственный F только для полноты. Этот подход имеет преимущество простоты, но если у вас есть контроль над F, переписывание его с небольшим вниманием к векторизации правильно может иметь огромные преимущества по скорости

import numpy

n1 = 100
n2 = 200

a = numpy.arange(n1)
b = numpy.arange(n2)

def F(x, y):
    return x + y

# Everything above this is setup, the answer to your question lies here:
fv = numpy.vectorize(F)
r = fv(a[:, numpy.newaxis], b)

На моем компьютере найдены следующие тайминги, показывающие цену, которую вы платите за "автоматическую" вектозацию:

%timeit fv(a[:, numpy.newaxis], b)
100 loops, best of 3: 3.58 ms per loop

%timeit F(a[:, numpy.newaxis], b)
10000 loops, best of 3: 38.3 µs per loop

Ответ 3

Вы можете использовать список для создания массива массивов:

import numpy as np

# Arrays
a = np.array([1, 2, 3]) # n1 = 3
b = np.array([4, 5]) # n2 = 2

# Your function (just an example)
def f(i, j):
    return i + j

result = np.array([[f(i, j)for j in b ]for i in a])
print result

Вывод:

[[5 6]
 [6 7]
 [7 8]]

Ответ 4

Могу ли я предложить, если ваш прецедент более ограничен продуктами, что вы используете внешний продукт?

например:.

import numpy

a = array([0, 1, 2])
b = array([0, 1, 2, 3])

numpy.outer(a,b)

возвращает

array([[0, 0, 0, 0],
       [0, 1, 2, 3],
       [0, 2, 4, 6]])

Затем вы можете применить другие преобразования:

numpy.outer(a,b) + 1

возвращает

array([[1, 1, 1, 1],
       [1, 2, 3, 4],
       [1, 3, 5, 7]])

Это намного быстрее:

>>> import timeit
>>> timeit.timeit('numpy.array([[i*j for i in a] for j in b])', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
31.79583477973938

>>> timeit.timeit('numpy.outer(a,b)', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
9.351550102233887
>>> timeit.timeit('numpy.outer(a,b)+1', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
12.308301210403442

Ответ 5

В качестве еще одной альтернативы, которая немного более расширяема, чем точечный продукт, менее чем за 1/5 - 1/9-е время использования вложенных списков, используйте numpy.newaxis (потребовалось немного больше, чтобы найти):

>>> import numpy
>>> a = numpy.array([0,1,2])
>>> b = numpy.array([0,1,2,3])

На этот раз, используя функцию питания:

>>> pow(a[:,numpy.newaxis], b)
array([[1, 0, 0, 0],
       [1, 1, 1, 1],
       [1, 2, 4, 8]])

По сравнению с альтернативой:

>>> numpy.array([[pow(i,j) for j in b] for i in a])
array([[1, 0, 0, 0],
       [1, 1, 1, 1],
       [1, 2, 4, 8]])

И сравнение времени:

>>> import timeit
>>> timeit.timeit('numpy.array([[pow(i,j) for i in a] for j in b])', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
31.943181037902832
>>> timeit.timeit('pow(a[:, numpy.newaxis], b)', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
5.985810041427612

>>> timeit.timeit('numpy.array([[pow(i,j) for i in a] for j in b])', 'import numpy; a=numpy.arange(10); b=numpy.arange(10)')
109.74687385559082
>>> timeit.timeit('pow(a[:, numpy.newaxis], b)', 'import numpy; a=numpy.arange(10); b=numpy.arange(10)')
11.989138126373291

Ответ 6

Если F() работает с широковещательными аргументами, определенно используйте это, как описывают другие. Альтернативой является использование np.fromfunction (function_on_an_int_grid будет лучшим именем.) Следующее просто отображает сетку int в вашу сетку a-b, затем в F():

import numpy as np

def func_allpairs( F, a, b ):
    """ -> array len(a) x len(b):
        [[ F( a0 b0 )  F( a0 b1 ) ... ]
         [ F( a1 b0 )  F( a1 b1 ) ... ]
         ...
        ]
    """
    def fab( i, j ):
        return F( a[i], b[j] )  # F scalar or vec, e.g. gradient

    return np.fromfunction( fab, (len(a), len(b)), dtype=int )  # -> fab( all pairs )


#...............................................................................
def F( x, y ):
    return x + 10*y

a = np.arange( 100 )
b = np.arange( 222 )
A = func_allpairs( F, a, b )
# %timeit: 1000 loops, best of 3: 241 µs per loop -- imac i5, np 1.9.3