python结合map()的两种逻辑

Question

在我的代码中，map() 的函数 f(a,b) 将两个 1x3 numpy 数组作为输入。由于它是别人从 C 代码实现的，所以我无法更改 f()。

第二个输入 y 始终是 Nx3 numpy 数组。第一个输入 x 有两种情况。

在一种情况下，它是一个 1x3 numpy 数组，因此我这样做

unwrap = partial(f, x)
result = map(unwrap, y)

在另一种情况下，它是一个 Nx3 numpy 数组，那么我做

unwrap = f
result = map(unwrap, x, y)

有没有办法将这两种情况结合在一起？

Answer 1

np.broadcast_to 可以“ reshape ”A 以匹配 B；然后你可以一起迭代两者。它使用跨步，因此内存使用并没有实际增加。

In [370]: def f(a,b):
     ...:     assert(a.shape==(1,3))
     ...:     assert(b.shape==(1,3))
     ...:     return a+b
     ...: 
In [371]: B=np.arange(12).reshape(4,3)
In [372]: A=np.arange(3).reshape(1,3)
In [373]: np.broadcast_to(A, B.shape)   # (1,3) to (4,3)
Out[373]: 
array([[0, 1, 2],
       [0, 1, 2],
       [0, 1, 2],
       [0, 1, 2]])
In [374]: np.broadcast_to(B, B.shape)   # no change with (4,3)
Out[374]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

我通常使用列表理解而不是映射:

In [375]: [f(np.atleast_2d(a),np.atleast_2d(b)) for a,b in zip(np.broadcast_to(A,B.shape),B)]
Out[375]: 
[array([[0, 2, 4]]),
 array([[3, 5, 7]]),
 array([[ 6,  8, 10]]),
 array([[ 9, 11, 13]])]

In [376]: [f(np.atleast_2d(a),np.atleast_2d(b)) for a,b in zip(np.broadcast_to(B,B.shape),B)]
Out[376]: 
[array([[0, 2, 4]]),
 array([[ 6,  8, 10]]),
 array([[12, 14, 16]]),
 array([[18, 20, 22]])]

二维数组的迭代产生一维数组的列表，因此需要 np.atleast_2d 来满足我的 f 断言。如果 f 也接受 (3,) 输入，我将不需要它。

或者使用 map:

In [377]: map(lambda a,b: f(np.atleast_2d(a),np.atleast_2d(b)), np.broadcast_to(B,B.shape),B)
Out[377]: <map at 0xb14f4c6c>
In [378]: list(_)
Out[378]: 
[array([[0, 2, 4]]),
 array([[ 6,  8, 10]]),
 array([[12, 14, 16]]),
 array([[18, 20, 22]])]
In [379]: map(lambda a,b: f(np.atleast_2d(a),np.atleast_2d(b)), np.broadcast_to(A,B.shape),B)
Out[379]: <map at 0xb0871a8c>
In [380]: list(_)
Out[380]: 
[array([[0, 2, 4]]),
 array([[3, 5, 7]]),
 array([[ 6,  8, 10]]),
 array([[ 9, 11, 13]])]

np.vectorize 和 np.frompyfunc 也处理这种广播，但它们是为采用标量而非一维数组的函数设计的。

使用 broadcast_arrays 我可以平等对待两个数组:

In [386]: map(lambda a,b: f(np.atleast_2d(a),np.atleast_2d(b)), *np.broadcast_arrays(B,A))
Out[386]: <map at 0xb69851ac>
In [387]: list(_)
Out[387]: 
[array([[0, 2, 4]]),
 array([[3, 5, 7]]),
 array([[ 6,  8, 10]]),
 array([[ 9, 11, 13]])]

更一般地说，A 和B 可以是生成所需(N,3) 数组的任何内容。我可以通过生成 (N,1,3) 数组来摆脱 atleast_2d 的使用:

In [397]: map(f, *np.broadcast_arrays(np.arange(3)[None,None,:], np.arange(0,40,10)[:,None,None]))
Out[397]: <map at 0xb08b562c>
In [398]: list(_)
Out[398]: 
[array([[0, 1, 2]]),
 array([[10, 11, 12]]),
 array([[20, 21, 22]]),
 array([[30, 31, 32]])]