编写python函数以在任意轴上操作

Question

很长一段时间以来，我一直在各种问题上努力解决这个问题，但从未设法找到好的解决方案。

基本上，如果我想编写一个函数来对给定的任意秩数组的任意轴执行操作（例如np.mean（A，axis = some_axis）的样式），一般的想法是如何做到的。

问题似乎总是归因于切片语法的僵化；如果我想访问第3个索引的第ith个切片，可以使用A [：，：，i]，但不能将其推广到第n个索引。

Answer 1

numpy函数使用几种方法来做到这一点：

• 移置轴以将目标轴移动到已知位置，通常是第一个或最后一个；并根据需要转置结果

• 重塑形状（以及转置）以减小问题尺寸。如果您将焦点放在第n个维度上，则（：n）维度在何处变平可能并不重要。他们只是在“兜风”。

• 构造一个索引元组。 idx = (slice(None), slice(None), j); A[idx]等效于A[:,:,j]。从一个大小合适的列表或数组开始，用切片填充，在其中进行修饰，然后转换为元组（元组是不可变的）。

• 使用indexing_tricks，np.r_等np.s_工具构建索引。

提供轴的学习代码。编译后的ufuncs并没有帮助，但是tensordot，take_along_axis，apply_along_axis，np.cross之类的函数是用Python编写的，并且使用了其中的一个或多个技巧

Answer 2

您不能对此进行概括。实际上，在这种情况下，示例numpy.mean(a, axis=axis_index)很好看。即使使用主要用C语言编写的numpy，也要遍历轴索引以知道在哪里计算平均值。看一下numpy.mean核心的reduction.c。即使它们在执行操作之前以有利的方式格式化数据，也始终需要用axis_index在所有轴上循环。