отображение более крупных размерных массивов

Question

Если у меня есть массив 1d a и я хочу отобразить функцию f для каждого элемента, я могу сделать

>>> import numpy as np 
>>> a = np.arange(5)
>>> def f(x):
...     return 3*x + x**2 #whatever 
>>> np.fromiter(map(f,a),float)
array([ 0.,  4., 10., 18., 28.])

Я бы хотел сделать что-то аналогичное с более сложными массивами. Один пример расчета: составить парные массивы 3х3 с умножением матрицы

>>> a = np.arange(5*2*3**2).reshape(5,2,3,3)
>>> def f(x):
...     return np.matmul(x[0],x[1])
# is there a smarter way? 
>>> np.array([f(x) for x in a])
array([[[   42,    45,    48],
        [  150,   162,   174],
        [  258,   279,   300]],

       [[ 1716,  1773,  1830],
        [ 1986,  2052,  2118],
        [ 2256,  2331,  2406]],

       [[ 5334,  5445,  5556],
        [ 5766,  5886,  6006],
        [ 6198,  6327,  6456]],

       [[10896, 11061, 11226],
        [11490, 11664, 11838],
        [12084, 12267, 12450]],

       [[18402, 18621, 18840],
        [19158, 19386, 19614],
        [19914, 20151, 20388]]])

Другим примером вычисления будет преобразование каждого вектора в массиве векторов путем умножения матриц

>>> a = np.arange(3*5).reshape(5,3)
>>> def f(x):
...     M = np.arange(3*3).reshape(3,3)
...     return np.dot(M,x)
>>> np.array([f(x) for x in a])
array([[  5,  14,  23],
       [ 14,  50,  86],
       [ 23,  86, 149],
       [ 32, 122, 212],
       [ 41, 158, 275]])

Есть ли хороший способ сделать такие вычисления с np.fromiter подобным подходом? Какой самый питонный способ сделать эти операции с numpy? Есть ли подход, который обрабатывает каждую примерную проблему здесь как np.specialnumpything(map(f,a))?

Answer 1

Это так же легко реализовать с помощью трансляции. А именно:

a = np.arange(5)
a*3 + a**2

array([ 0,  4, 10, 18, 28])

a = np.arange(5*2*3**2).reshape(5,2,3,3)

a[:, 0] @ a[:, 1]

array([[[   42,    45,    48],
        [  150,   162,   174],
        [  258,   279,   300]],

       [[ 1716,  1773,  1830],
        [ 1986,  2052,  2118],
        [ 2256,  2331,  2406]],

       [[ 5334,  5445,  5556],
        [ 5766,  5886,  6006],
        [ 6198,  6327,  6456]],

       [[10896, 11061, 11226],
        [11490, 11664, 11838],
        [12084, 12267, 12450]],

       [[18402, 18621, 18840],
        [19158, 19386, 19614],
        [19914, 20151, 20388]]])


a = np.arange(3*5).reshape(5,3)
M = np.arange(3*3).reshape(3,3)

M.dot(a.T).T

array([[  5,  14,  23],
       [ 14,  50,  86],
       [ 23,  86, 149],
       [ 32, 122, 212],
       [ 41, 158, 275]])

np.einsum('kj, ij -> ik', M, a)

array([[  5,  14,  23],
       [ 14,  50,  86],
       [ 23,  86, 149],
       [ 32, 122, 212],
       [ 41, 158, 275]])

Answer 2

Я бы использовал встроенный numpy.nditer, который может быть тем, что вы ищете:

https://docs.scipy.org/doc/numpy/reference/arrays.nditer.html

из примеров:

>>> a = np.arange(6).reshape(2,3)
>>> a
>>> array([[0, 1, 2],
           [3, 4, 5]])
>>> with np.nditer(a, op_flags=['readwrite']) as it:
    ...    for x in it:
    ...        x[...] = 2 * x
    ...
>>> a
    array([[ 0,  2,  4],
           [ 6,  8, 10]])