跨 3D numpy 数组的多个轴应用函数

Question

我有形状为(2, n, m) 的 3-D numpy 数组，n,m 是任意的，我们称之为X。我有两个函数F = F(x: np.ndarray, y: np.ndarray) 和G = G(x: np.ndarray, y: np.ndarray)，其中x 和y 是长度为m 的一维数组，这些函数也返回形状为m 的一维数组。我想对X 应用转换，在第二个轴上操作，使用F 和G

for idx in range(n):
    X[0][idx], X[1][idx] = F(X[0][idx], X[1][idx]), G(X[0][idx], X[1][idx])

示例

假设m = n = 2，然后让

import numpy as np

X = np.arange(8).reshape(2,2,2)


In [18]: X
Out[18]: 
array([[[0, 1],
        [2, 3]],

       [[4, 5],
        [6, 7]]])

假设 F 和 G 由（虚拟函数这些不是真实的）定义

def F(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    return x - y

def G(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    return y - x

然后将转换应用于 X

for idx in range(2):
    X[0][idx], X[1][idx] = F(X[0][idx], X[1][idx]), G(X[0][idx], X[1][idx])

我明白了

In [23]: X
Out[23]: 
array([[[-4, -4], 
        [-4, -4]],

       [[ 4,  4],
        [ 4,  4]]])

示例 2

考虑F = np.cross、G = -F 和X =np.arange(12).reshape(2,2,3)。

for idx in range(2):
    X[0][idx], X[1][idx] = np.cross(X[0][idx], X[1][idx]), - np.cross(X[0[idx], X[1][idx])

那么，

In [49]: X
Out[49]: 
array([[[ -6,  12,  -6],
        [ -6,  12,  -6]],

       [[  6, -12,   6],
        [  6, -12,   6]]])

我重申 F 和 G 不是真正的函数，它应该适用于任意定义。我想要一个不使用我上面写的 for 循环的解决方案

Answer 1

你可以使用np.apply_along_axis:

>>> np.apply_along_axis(lambda x: [F(x[0], x[1]), G(x[0], x[1])], 0, X)
array([[[-4, -4],
        [-4, -4]],

       [[ 4,  4],
        [ 4,  4]]])
>>> 

这会将函数应用于 numpy 数组的每个子数组，就像原始 for 循环一样，它使用 lambda 将函数映射到值。

编辑：

你可以试试zip和map：

>>> np.array([*zip(*map(lambda x: [np.cross(x[0], x[1]), -np.cross(x[0], x[1])], zip(*X)))])
array([[[ -6,  12,  -6],
        [ -6,  12,  -6]],

       [[  6, -12,   6],
        [  6, -12,   6]]])
>>>