Как применить функцию, которая возвращает вектор к каждому элементу массива numpy (и получить массив с большей размерностью)
Давайте напишем это прямо в коде
Примечание: я отредактировал маппер (в оригинальном примере, например, x -> (x, 2 * x, 3 * x)) для общей функции черного ящика, которая вызывает проблемы.
import numpy as np
def blackbox_fn(x): #I can't be changed!
assert np.array(x).shape == (), "I'm a fussy little function!"
return np.array([x, 2*x, 3*x])
# let's have 2d array
arr2d = np.array(list(range(4)), dtype=np.uint8).reshape(2, 2)
# each element should be mapped to vector
def mapper(x, blackbox_fn):
# there is some 3rdparty non-trivial function, returning np.array
# in examples returns np.array((x, 2 * x, 3 * x))
# but still this 3rdparty function operates only on scalar values
return vectorized_blackbox_fn(x)
Так что для ввода 2d массива
array([[0, 1],
[2, 3]], dtype=uint8)
Я хотел бы получить массив 3d
array([[[0, 0, 0],
[1, 2, 3]],
[[2, 4, 6],
[3, 6, 9]]], dtype=uint8)
Я могу написать наивный алгоритм, используя для цикла
# result should be 3d array, last dimension is same as mapper result size
arr3d = np.empty(arr2d.shape + (3,), dtype=np.uint8)
for y in range(arr2d.shape[1]):
for x in xrange(arr2d.shape[0]):
arr3d[x, y] = mapper(arr2d[x, y])
Но это кажется довольно медленным для больших массивов. Я знаю, что есть np.vectorize, но используя
np.vectorize(mapper)(arr2d)
не работает, из-за
ValueError: setting an array element with a sequence.
(кажется, что векторизация не может изменить размерность) Есть ли какое-нибудь лучшее (тупое идиоматическое и более быстрое) решение?
3 ответа
np.vectorize с новой опцией подписи может справиться с этим. Это не улучшает скорость, но облегчает учет размеров.
In [159]: def blackbox_fn(x): #I can't be changed!
...: assert np.array(x).shape == (), "I'm a fussy little function!"
...: return np.array([x, 2*x, 3*x])
...:
Документация для signature немного загадочно Я работал с этим раньше, поэтому сделал хорошее первое предположение:
In [161]: f = np.vectorize(blackbox_fn, signature='()->(n)')
In [162]: f(np.ones((2,2)))
Out[162]:
array([[[ 1., 2., 3.],
[ 1., 2., 3.]],
[[ 1., 2., 3.],
[ 1., 2., 3.]]])
С вашим массивом:
In [163]: arr2d = np.array(list(range(4)), dtype=np.uint8).reshape(2, 2)
In [164]: f(arr2d)
Out[164]:
array([[[0, 0, 0],
[1, 2, 3]],
[[2, 4, 6],
[3, 6, 9]]])
In [165]: _.dtype
Out[165]: dtype('int32')
dtype не сохраняется, потому что ваш blackbox_fn не сохраняет его. По умолчанию vectorize делает тестовый расчет с первым элементом и использует его dtype определить результат dtype. Можно указать возврат dtype с помощью otypes параметр.
Он может обрабатывать массивы, кроме 2d:
In [166]: f(np.arange(3))
Out[166]:
array([[0, 0, 0],
[1, 2, 3],
[2, 4, 6]])
In [167]: f(3)
Out[167]: array([3, 6, 9])
С signaturevectorize использует итерацию уровня Python. Без подписи он использует np.frompyfunc, с немного лучшей производительностью. Но пока blackbox_fn должен быть вызван для элемента ввода, мы не можем значительно увеличить скорость (не более 2х).
np.frompyfunc возвращает массив dtype объекта:
In [168]: fpy = np.frompyfunc(blackbox_fn, 1,1)
In [169]: fpy(1)
Out[169]: array([1, 2, 3])
In [170]: fpy(np.arange(3))
Out[170]: array([array([0, 0, 0]), array([1, 2, 3]), array([2, 4, 6])], dtype=object)
In [171]: np.stack(_)
Out[171]:
array([[0, 0, 0],
[1, 2, 3],
[2, 4, 6]])
In [172]: fpy(arr2d)
Out[172]:
array([[array([0, 0, 0]), array([1, 2, 3])],
[array([2, 4, 6]), array([3, 6, 9])]], dtype=object)
stack не может удалить вложенность массива в этом 2d случае:
In [173]: np.stack(_)
Out[173]:
array([[array([0, 0, 0]), array([1, 2, 3])],
[array([2, 4, 6]), array([3, 6, 9])]], dtype=object)
но мы можем разложить это и сложить. Это нужно reshape:
In [174]: np.stack(__.ravel())
Out[174]:
array([[0, 0, 0],
[1, 2, 3],
[2, 4, 6],
[3, 6, 9]])
Скоростные тесты:
In [175]: timeit f(np.arange(1000))
14.7 ms ± 322 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [176]: timeit fpy(np.arange(1000))
4.57 ms ± 161 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [177]: timeit np.stack(fpy(np.arange(1000).ravel()))
6.71 ms ± 207 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [178]: timeit np.array([blackbox_fn(i) for i in np.arange(1000)])
6.44 ms ± 235 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Если ваша функция вернет список вместо любого массива, это может упростить повторную сборку и, возможно, даже быстрее
def foo(x):
return [x, 2*x, 3*x]
или играть с frompyfunc параметры;
def foo(x):
return x, 2*x, 3*x # return a tuple
In [204]: np.stack(np.frompyfunc(foo, 1,3)(arr2d),2)
Out[204]:
array([[[0, 0, 0],
[1, 2, 3]],
[[2, 4, 6],
[3, 6, 9]]], dtype=object)
10-кратная скорость - я удивлен
In [212]: foo1 = np.frompyfunc(foo, 1,3)
In [213]: timeit np.stack(foo1(np.arange(1000)),1)
428 µs ± 17.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
Возможно, я ошибаюсь, но понимание делает свою работу:
a = np.array([[0, 1],
[2, 3]])
np.array([[[j, j*2, j*3] for j in i] for i in a ])
#[[[0 0 0]
# [1 2 3]]
#
# [[2 4 6]
# [3 6 9]]]
Вы можете использовать базовое вещание NumPy для таких "внешних продуктов".
np.arange(3)[:, None] * np.arange(2)
# array([[0, 0],
# [0, 1],
# [0, 2]])
В вашем случае это будет
def mapper(x):
return (np.arange(3)[:, None, None] * x).transpose((1, 2, 0))
Обратите внимание .transpose() нужен только если вам нужно, чтобы новая ось была в конце.
И это почти в 3 раза быстрее, чем сложение 3 отдельных умножений:
def mapper(x):
return (np.arange(3)[:, None, None] * x).transpose((1, 2, 0))
def mapper2(x):
return np.stack((x, 2 * x, 3 * x), axis = -1)
a = np.arange(30000).reshape(-1, 30)
%timeit mapper(a) # 48.2 µs ± 417 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit mapper2(a) # 137 µs ± 3.57 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)