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