针对3D NumPy数组的高效且模块化运算实现方案问询

问题背景 #

给定以下输入：

• img：形状为(200, 200, 256)的3D NumPy数组，代表高光谱图像，包含200×200个像素，每个像素对应256个等间隔波长的整数反射率值
• reference_spectrum：包含256个整数反射率值的NumPy数组（等价于单个像素的光谱向量）
• dist：接收两个等长NumPy数组，返回单个实数的距离函数

需要生成形状为(200, 200)的2D NumPy数组result，满足result[i, j] = dist(img[i, j], reference_spectrum)。

现有方案的问题 #

嵌套循环实现 #

最初的代码实现：

result = np.zeros((img.shape[0], img.shape[1]))
for i in range(img.shape[0]):
    for j in range(img.shape[1]):
        result[i, j] = dist(img[i, j], reference_spectrum)

缺点：Python层级的嵌套循环效率极低，处理大尺寸图像时速度很慢，且代码冗余不简洁。

特定距离的向量化实现 #

针对曼哈顿距离（$dist(a, b) = \sum_{i=0}^{n} |a_i-b_i|$）的优化代码：

temp = np.zeros(img.shape)
for k in range(img.shape[2]):
    temp[:, :, k] = abs(img[:, :, k] - reference_spectrum[k])
result = np.sum(temp, axis=2)

缺点：仅适配特定距离函数，无法通用到任意自定义dist函数，扩展性差。

通用高效解决方案 #

方案1：使用np.apply_along_axis（通用性优先） #

np.apply_along_axis可以沿着指定轴遍历数组的子元素，并应用自定义函数，完美平衡通用性与效率：

import numpy as np

# 示例自定义距离函数（可替换为任意符合要求的dist）
def dist(a, b):
    return np.sum(np.abs(a - b))  # 曼哈顿距离
    # return np.sqrt(np.sum((a - b)**2))  # 欧氏距离
    # 可自行实现余弦距离、相关系数等任意距离/相似度函数

# 核心计算：沿着波段轴（第2轴）对每个像素应用dist
result = np.apply_along_axis(
    func1d=lambda pixel: dist(pixel, reference_spectrum),
    axis=2,
    arr=img
)

说明：

• axis=2指定遍历方向为高光谱的波段维度，每个被处理的子数组就是单个像素的256维光谱向量
• 匿名函数将每个像素向量与参考光谱传入自定义dist函数，自动生成对应位置的结果值
• 相比嵌套循环，该方法内部基于NumPy的优化机制执行，效率提升明显，同时完全支持任意符合接口的dist函数。

方案2：向量化改造距离函数（性能优先） #

如果追求极致性能，可以将dist函数改造成完全向量化的形式，避免Python层级的循环：

# 向量化的曼哈顿距离示例
def vectorized_manhattan_dist(img, ref):
    # 通过广播将参考光谱扩展为与img匹配的形状，再逐元素计算后求和
    return np.sum(np.abs(img - ref[np.newaxis, np.newaxis, :]), axis=2)

# 向量化的欧氏距离示例
def vectorized_euclidean_dist(img, ref):
    return np.sqrt(np.sum((img - ref[np.newaxis, np.newaxis, :])**2, axis=2))

# 调用生成结果
result = vectorized_manhattan_dist(img, reference_spectrum)

说明：

• 利用NumPy的广播机制，将参考光谱扩展为(1,1,256)形状，与(200,200,256)的img逐元素运算
• 这种方式性能最优，但需要针对不同距离函数做对应的向量化改造，通用性稍弱。

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