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.])

我想做一些与更复杂的数组类似的事情。一个示例计算是这样的：用矩阵乘法组成成对的3x3数组

>>> 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进行这些操作的最Python方式是什么？有没有一种可以np.specialnumpything(map(f,a))的方式处理每个示例问题的方法？

Answer 1

我会使用内置的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]])

Answer 2

这与广播一样容易实现。即：

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]])

映射到高维numpy数组

2 个答案: