如何在NumPy中实现任意维度数组的函数通用化处理？

问题描述 #

现有不同维度的numpy网格数组，每个元素对应一组参数（数组最后一维为参数集合），需要将对应参数数量的函数应用到每个元素上。目前针对1D、2D、3D分别实现了循环逻辑，希望编写一个支持任意维度的通用函数apply_function，替代重复的维度分支代码。

现有分维度代码 #

import numpy as np

# 1D
grid1 = np.array([1, 2, 3])
foo1 = lambda x: x**2
f_grid1 = np.zeros_like(grid1)
for i in range(grid1.shape[0]):
    f_grid1[i] = foo1(grid1[i])
print(f_grid1)  # 输出 [1 4 9]

# 2D
grid2 = np.array(
    [
        [[1, 0.1], [1, 0.2], [1, 0.3]],
        [[2, 0.1], [2, 0.2], [2, 0.3]],
        [[3, 0.1], [3, 0.2], [3, 0.3]],
        [[4, 0.1], [4, 0.2], [4, 0.3]]
    ]
)
foo2 = lambda x, y: x**2 - y
f_grid2 = np.empty(grid2.shape[:-1], dtype=np.float64)
for i in range(grid2.shape[0]):
    for j in range(grid2.shape[1]):
        f_grid2[i, j] = foo2(grid2[i, j, 0], grid2[i, j, 1])
print(f_grid2)
# 输出
# [[ 0.9  0.8  0.7]
#  [ 3.9  3.8  3.7]
#  [ 8.9  8.8  8.7]
#  [15.9 15.8 15.7]]

# 3D
grid3 = np.array(
    [
        [ [[1, 0.1, 0.0001], [1, 0.2, 0.0001], [1, 0.3, 0.0001]],
          [[2, 0.1, 0.0001], [2, 0.2, 0.0001], [2, 0.3, 0.0001]],
          [[3, 0.1, 0.0001], [3, 0.2, 0.0001], [3, 0.3, 0.0001]] ],

        [ [[1, 0.1, 0.0002], [1, 0.2, 0.0002], [1, 0.3, 0.0002]],
          [[2, 0.1, 0.0002], [2, 0.2, 0.0002], [2, 0.3, 0.0002]],
          [[3, 0.1, 0.0002], [3, 0.2, 0.0002], [3, 0.3, 0.0002]] ],

        [ [[1, 0.1, 0.0003], [1, 0.2, 0.0003], [1, 0.3, 0.0003]],
          [[2, 0.1, 0.0003], [2, 0.2, 0.0003], [2, 0.3, 0.0003]],
          [[3, 0.1, 0.0003], [3, 0.2, 0.0003], [3, 0.3, 0.0003]] ],
    ]
)
foo3 = lambda x, y, z: x**2 - y + z
f_grid3 = np.empty(grid3.shape[:-1], dtype=np.float64)
for i in range(grid3.shape[0]):
    for j in range(grid3.shape[1]):
        for k in range(grid3.shape[2]):
            f_grid3[i, j, k] = foo3(grid3[i, j, k, 0], grid3[i, j, k, 1], grid3[i, j, k, 2])
print(f_grid3)
# 输出
# [[[0.9001 0.8001 0.7001]
#   [3.9001 3.8001 3.7001]
#   [8.9001 8.8001 8.7001]]
# 
#  [[0.9002 0.8002 0.7002]
#   [3.9002 3.8002 3.7002]
#   [8.9002 8.8002 8.7002]]
# 
#  [[0.9003 0.8003 0.7003]
#   [3.9003 3.8003 3.7003]
#   [8.9003 8.8003 8.7003]]]

期望的通用函数结构：

def apply_function(grid, func):
    shape = grid.shape
    if len(shape) == 1:
        grid = grid[np.newaxis].T

    out_grid = np.empty(grid.shape[:-1], dtype=np.float64)

    # 此处实现通用逻辑

通用实现方案 #

以下提供三种高效的通用实现方式，优先推荐前两种（基于numpy原生向量化/内置函数，性能优于手动循环）：

方案1：利用np.apply_along_axis（简洁直观） #

np.apply_along_axis可以沿着指定轴对数组的每个子数组应用函数，这里我们沿着最后一维的前序维度，将每个参数组传递给目标函数（通过*arr解包参数）。

import numpy as np

def apply_function(grid, func):
    shape = grid.shape
    # 处理1D特殊情况：将(3,)转为(3,1)，统一最后一维为参数维度
    if len(shape) == 1:
        grid = grid[:, np.newaxis]
    
    # 沿着最后一维，将每个参数数组解包后传入func
    out_grid = np.apply_along_axis(lambda arr: func(*arr), axis=-1, arr=grid)
    return out_grid.astype(np.float64)

测试验证：

# 测试1D
grid1 = np.array([1,2,3])
foo1 = lambda x: x**2
print(apply_function(grid1, foo1))  # 输出 [1. 4. 9.]

# 测试2D
grid2 = np.array(
    [
        [[1, 0.1], [1, 0.2], [1, 0.3]],
        [[2, 0.1], [2, 0.2], [2, 0.3]],
        [[3, 0.1], [3, 0.2], [3, 0.3]],
        [[4, 0.1], [4, 0.2], [4, 0.3]]
    ]
)
foo2 = lambda x,y: x**2 - y
print(apply_function(grid2, foo2))
# 输出与原代码一致

# 测试3D
grid3 = np.array(
    [
        [ [[1, 0.1, 0.0001], [1, 0.2, 0.0001], [1, 0.3, 0.0001]],
          [[2, 0.1, 0.0001], [2, 0.2, 0.0001], [2, 0.3, 0.0001]],
          [[3, 0.1, 0.0001], [3, 0.2, 0.0001], [3, 0.3, 0.0001]] ],

        [ [[1, 0.1, 0.0002], [1, 0.2, 0.0002], [1, 0.3, 0.0002]],
          [[2, 0.1, 0.0002], [2, 0.2, 0.0002], [2, 0.3, 0.0002]],
          [[3, 0.1, 0.0002], [3, 0.2, 0.0002], [3, 0.3, 0.0002]] ],

        [ [[1, 0.1, 0.0003], [1, 0.2, 0.0003], [1, 0.3, 0.0003]],
          [[2, 0.1, 0.0003], [2, 0.2, 0.0003], [2, 0.3, 0.0003]],
          [[3, 0.1, 0.0003], [3, 0.2, 0.0003], [3, 0.3, 0.0003]] ],
    ]
)
foo3 = lambda x,y,z: x**2 - y + z
print(apply_function(grid3, foo3))
# 输出与原代码一致

方案2：向量化参数拆分（性能最优） #

利用numpy的索引广播，将最后一维的每个参数单独拆分出来，直接传入函数进行向量化计算，完全避免循环，性能最佳。

import numpy as np

def apply_function(grid, func):
    shape = grid.shape
    if len(shape) == 1:
        grid = grid[:, np.newaxis]
    
    # 将最后一维的参数逐个拆分，作为func的位置参数
    params = [grid[..., i] for i in range(grid.shape[-1])]
    out_grid = func(*params)
    return out_grid.astype(np.float64)

说明：这种方式要求func本身支持numpy数组的向量化运算（大部分numpy原生操作和lambda表达式都支持），如果你的func是自定义的非向量化函数，可以先用np.vectorize包装，但性能会略有下降。

方案3：手动遍历所有维度（兼容非向量化函数） #

如果func不支持向量化运算，可以用np.ndenumerate遍历所有索引，逐个计算结果：

import numpy as np

def apply_function(grid, func):
    shape = grid.shape
    if len(shape) == 1:
        grid = grid[:, np.newaxis]
    
    out_shape = grid.shape[:-1]
    out_grid = np.empty(out_shape, dtype=np.float64)
    
    # 遍历每个非参数维度的索引
    for idx in np.ndindex(out_shape):
        # 取出当前索引对应的参数组，解包后传入func
        params = grid[idx]
        out_grid[idx] = func(*params)
    
    return out_grid

优缺点：兼容所有类型的函数，但性能低于前两种方案，适合处理复杂的非向量化逻辑。

内容的提问来源于stack exchange，提问作者Diego