如何用Numba实现三维数组中每个像素的极值计算？

问题原因 #

Numba的nopython模式对np.max/np.min的axis参数支持有限，不接受负整数索引（如axis=-1），这会导致编译时抛出TypingError。

解决方案 #

方案1：替换为固定正轴索引（适用于维度固定的场景）

如果你的数组维度固定（比如始终是(H, W, 3)的BGR图像），直接将axis=-1替换为对应的正索引（比如axis=2）即可：

import numba as nb
import numpy as np

@nb.njit(cache=True, fastmath=True)
def max_per_cell(arr):
    # 针对3维图像数组，最后一维索引为2
    return np.max(arr, axis=2)

@nb.njit(cache=True, fastmath=True)
def min_per_cell(arr):
    return np.min(arr, axis=2)

img = np.random.random((3, 4, 3))
max_per_cell(img)
min_per_cell(img)

方案2：动态计算正轴索引（适用于维度可变的场景）

如果数组维度可能变化，通过arr.ndim - 1将负索引转换为正索引，保证代码通用性：

import numba as nb
import numpy as np

@nb.njit(cache=True, fastmath=True)
def max_per_cell(arr):
    # 将axis=-1转为对应的正索引
    target_axis = arr.ndim - 1
    return np.max(arr, axis=target_axis)

@nb.njit(cache=True, fastmath=True)
def min_per_cell(arr):
    target_axis = arr.ndim - 1
    return np.min(arr, axis=target_axis)

img = np.random.random((3, 4, 3))
max_per_cell(img)
min_per_cell(img)

方案3：使用reduce方法替代np.max/min

Numba对np.maximum.reduce和np.minimum.reduce的支持更完善，性能表现与np.max/np.min相当，也是可靠的替代方案：

import numba as nb
import numpy as np

@nb.njit(cache=True, fastmath=True)
def max_per_cell(arr):
    return np.maximum.reduce(arr, axis=arr.ndim - 1)

@nb.njit(cache=True, fastmath=True)
def min_per_cell(arr):
    return np.minimum.reduce(arr, axis=arr.ndim - 1)

img = np.random.random((3, 4, 3))
max_per_cell(img)
min_per_cell(img)

性能说明 #

上述三种方案的性能差异极小，都能满足HSL/BGR转换的性能需求。如果是固定维度的图像数组，方案1的性能略优（省去了维度计算步骤）。

内容的提问来源于stack exchange，提问作者Ξένη Γήινος