如何实现将2D Numpy数组映射为3D数组的向量化函数？

嘿，我完全懂你遇到的困扰——np.vectorize看起来应该能帮你把每个元素映射成一维数组进而扩展维度，但实际输出却拆成了多个独立数组，根本不是你想要的形状。别担心，咱们来一步步解决这个问题，还有更高效的方法替代vectorize哦！

为什么np.vectorize达不到预期？ #

np.vectorize本质上是Python循环的包装器，它默认会把函数返回的多个值（或者一个多维数组）拆分成独立的输出数组，而不是将它们合并成更高维度的结果。如果不指定otypes参数，它会自动推断输出类型，这就导致你看到的是元组形式的多个数组，而非合并后的高维数组。

方法1：优先使用Numpy原生向量化操作（最推荐） #

如果你的映射逻辑可以用Numpy的内置操作实现，这绝对是最优解——因为它是C语言底层实现的，效率比vectorize高得多，还能完美处理维度扩展。

示例1：实现foo的需求（每个元素复制3次） #

import numpy as np

a = np.array([[1,2],[3,4]])
# 给原数组增加一个末尾维度，再和长度为3的数组广播
y = np.expand_dims(a, axis=-1) * np.ones(3, dtype=a.dtype)
print(y)
# 输出：
# [[[1 1 1]
#   [2 2 2]]
#  [[3 3 3]
#   [4 4 4]]]

示例2：实现bar的需求（每个元素生成[x, 2x, 5]） #

a = np.array([[1,2],[3,4]])
# 分别构造三个对应数组，再沿最后一个轴堆叠
y = np.stack([a, 2*a, np.full_like(a, 5)], axis=-1)
print(y)
# 输出：
# [[[1 2 5]
#   [2 4 5]]
#  [[3 6 5]
#   [4 8 5]]]

这种方法的通用性极强：不管输入是几维数组，只要你能构造出对应的各个维度的数组，用np.stack沿新轴堆叠就能得到原维度数+新维度数的结果。比如输入是5维数组，每个元素映射成3维数组，最终就能得到7维数组。

方法2：用np.vectorize配合类型指定（适合复杂自定义函数） #

如果你的映射逻辑非常复杂，没法用原生向量化操作实现（比如包含条件判断、Python原生逻辑），可以给np.vectorize指定otypes为object，让它返回包含数组的对象数组，再合并成高维数组。

示例：改造foo函数 #

import numpy as np

@np.vectorize(otypes=[object])
def foo(x):
    return np.array([x, x, x])

a = np.array([[1,2],[3,4]])
# 得到(2,2)的object数组，每个元素是(3,)的数组
result_obj = foo(a)
# 展平后堆叠，再重塑成目标形状
y = np.stack(result_obj.ravel(), axis=0).reshape(a.shape + (3,))
print(y)

示例：改造bar函数 #

@np.vectorize(otypes=[object])
def bar(x):
    return np.array([x, 2*x, 5])

a = np.array([[1,2],[3,4]])
result_obj = bar(a)
y = np.stack(result_obj.ravel(), axis=0).reshape(a.shape + (3,))
print(y)

⚠️ 注意：这种方法本质还是Python循环，效率远低于原生向量化操作，只适合处理逻辑复杂、数据量不大的场景。

方法3：用np.apply_along_axis（适合简单自定义映射） #

你也可以用np.apply_along_axis来处理，先将原数组展平，处理后再重塑形状：

import numpy as np

def bar_func(x):
    # x是单个元素（展平后的输入）
    return np.array([x, 2*x, 5])

a = np.array([[1,2],[3,4]])
# 展平数组后沿轴0处理，再转置重塑
result = np.apply_along_axis(bar_func, 0, a.ravel())
y = result.T.reshape(a.shape + (3,))
print(y)

这个方法同样是Python循环包装，效率不如原生向量化，仅供参考。

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