如何为ndarray操作函数添加轴控制参数并实现维度扩展？

我想把原本处理n维ndarray的函数改造为可接收(n+1)维数组，支持指定axis参数，沿该轴的所有n维切片执行运算。比如让仅支持2D数组的函数能处理3D批量数据，同时希望用装饰器（比如@along_axis）精简代码，实现类似np.apply_along_axis的功能，但要支持n维输入函数，还要尽量避免内存拷贝。

我参考动态切片实现了一个针对2D函数的装饰器示例：

import numpy as np
from functools import wraps

def along_axis(func):
    """
    包装一个处理2D数组的函数，使其能沿3D数组的指定轴对所有2D切片执行运算。输入2D数组时忽略axis参数。
    """
    @wraps(func)
    def wrapper(arr, *args, axis=0, **kwargs):
        if arr.ndim == 3:
            arr = np.moveaxis(arr,axis,-1)
            return np.array([func(arr[:,:,t],*args,**kwargs) for t in range(arr.shape[-1])])
        return func(arr,*args,**kwargs)
    return wrapper

将其应用到处理2D图像的radial_mean函数：

@along_axis
def radial_mean(image):
    """
    计算2D图像的径向均值。包装后可处理3D批量图像数据。
    """
    # 创建以图像中心为原点的径向网格
    X,Y = np.meshgrid(np.arange(image.shape[1]),np.arange(image.shape[0]))
    R = np.sqrt((X - image.shape[1]//2)**2 + (Y - image.shape[0]//2)**2)
    
    # 遍历所有径向距离
    r  = np.arange(int(R.max()))
    
    # 计算径向均值
    f = np.vectorize(lambda r : image[(R >= r-0.5) & (R < r+0.5)].mean())
    return f(r)

测试验证该方案能同时处理单张2D图像和批量3D数据：

images = np.random.rand(1000,16,16)
print(radial_mean(images, axis=0).shape)
print(radial_mean(images[0], axis=0).shape)

> (1000,11)
> (11,)

请问这个实现方案是否合理？有没有更优的实现方式？

现有方案的合理性 #

你的实现逻辑通顺，核心思路没问题：

• 用np.moveaxis将指定轴移到末尾，方便按索引切片遍历，该操作仅修改数组视图，不会产生内存拷贝，效率很高。
• 装饰器通过@wraps保留了原函数的元信息，符合Python装饰器的最佳实践。
• 兼容了2D单样本和3D批量样本的输入场景，满足了基本需求。

但方案存在几个可优化的点：

1. 装饰器通用性不足：当前仅支持2D→1D的函数，无法适配任意n维输入的函数（比如3D函数处理4D批量数据）。
2. 列表推导的内存开销：np.array([...])会先将所有切片的计算结果存储为Python列表，再转换为数组，批量数据量较大时会额外占用内存。
3. np.vectorize效率低下：vectorize本质是Python循环的封装，并非真正的向量化运算，处理大图像时速度较慢。

更优的实现方式 #

1. 通用化along_axis装饰器 #

改造装饰器，使其支持任意n维输入函数，自动识别输入数组维度是否为目标维度+1，动态处理切片：

import numpy as np
from functools import wraps

def along_axis(func):
    @wraps(func)
    def wrapper(arr, *args, axis=0, **kwargs):
        # 若输入数组维度比原函数处理的维度多1，则执行批量运算
        if arr.ndim > func.__code__.co_argcount:
            # 先推导原函数处理的目标维度
            target_ndim = arr.ndim - 1
            # 将指定轴移到末尾，仅修改视图无内存拷贝
            arr_moved = np.moveaxis(arr, axis, -1)
            batch_size = arr_moved.shape[-1]
            
            # 预分配结果数组：通过单切片计算确定结果形状，避免列表转数组的内存开销
            sample_result = func(arr_moved[..., 0], *args, **kwargs)
            result = np.empty((batch_size,) + sample_result.shape, dtype=sample_result.dtype)
            
            # 循环填充结果
            for i in range(batch_size):
                result[i] = func(arr_moved[..., i], *args, **kwargs)
            return result
        return func(arr, *args, **kwargs)
    return wrapper

这个版本的装饰器可适配任意维度的函数，预分配数组的方式也避免了列表推导带来的额外内存占用。

2. 优化radial_mean函数的效率 #

替换np.vectorize为真正的向量化运算，大幅提升计算速度：

@along_axis
def radial_mean(image):
    X,Y = np.meshgrid(np.arange(image.shape[1]),np.arange(image.shape[0]))
    R = np.sqrt((X - image.shape[1]//2)**2 + (Y - image.shape[0]//2)**2)
    r_max = int(R.max())
    r = np.arange(r_max)
    
    # 向量化生成掩码并计算均值
    R_expanded = R[np.newaxis, ...]
    r_expanded = r[:, np.newaxis, np.newaxis]
    masks = (R_expanded >= r_expanded - 0.5) & (R_expanded < r_expanded + 0.5)
    
    # 用sum/count代替mean，避免空切片的除以0警告
    sums = np.sum(image * masks, axis=(1,2))
    counts = np.sum(masks, axis=(1,2))
    means = np.divide(sums, counts, where=counts!=0)
    means[counts==0] = 0  # 空切片结果可根据需求设为0或NaN
    return means

该版本完全基于NumPy向量化运算实现，避免了Python循环的开销，处理大图像时速度提升明显。

3. 进一步减少内存开销 #

• 保留np.moveaxis的视图操作，避免不必要的内存拷贝。
• 预分配结果数组时显式指定dtype，避免类型转换的额外开销。

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