在 i^{th} 轴之前和之后合并轴答案

【问题标题】：Merge axes before and after the i^{th} axis在 i^{th} 轴之前和之后合并轴
【发布时间】：2023-03-29 02:15:01
【问题描述】：

对于arr 中的任何一个axis，我想将一个numpy 数组arr 重塑为(before, at, after) 的形状。如何更快地做到这一点？

轴已归一化：0 <= axis < arr.ndim

程序：

import numpy as np
def f(arr, axis):
    shape = arr.shape
    before = int(np.product(shape[:axis]))
    at = shape[axis]
    return arr.reshape(before, at, -1)

测试：

a = np.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5)
print(f(a, 2).shape)

结果：

(6, 4, 5)

【问题讨论】：

% timeit (f(a, 2).shape) 8.14 µs ± 489 ns 每个循环（平均值±标准差，7 次运行，每次 100000 次循环）你真的对速度感兴趣吗？ ;)
现在是星期五晚上，我想我很无聊或强迫症哈哈。像arr += np.arange(arr.size, dtype=arr.dtype) 这样的东西在同样的意义上有点疯狂，但很多人都这样做。这只是 c++ 中的一个简单循环 ...
每次循环使用低于 5.08 µs ± 234 ns 的 reduceat（平均 ± 标准偏差，7 次运行，每次 100000 次循环），因此您将有更多时间在手 ;)

标签： python numpy

【解决方案1】：

shape 是一个元组，期望的结果也是一个元组。转换为/从数组使用np.prod 或其他一些数组函数需要时间。所以如果我们可以用纯 Python 代码做同样的事情，我们可能会节省时间。

例如shape:

In [309]: shape
Out[309]: (2, 3, 4, 5)
In [310]: np.prod(shape)
Out[310]: 120
In [311]: functools.reduce(operator.mul,shape)
Out[311]: 120

In [312]: timeit np.prod(shape)
13.6 µs ± 30.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [313]: timeit functools.reduce(operator.mul,shape)
647 ns ± 12.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

python 版本明显更快。我必须导入 functools 和 operator 才能得到 sum (Python3) 的乘法等价物。

或者获取新的形状元组：

In [314]: axis=2
In [315]: (functools.reduce(operator.mul,shape[:axis]),shape[axis],-1)
Out[315]: (6, 4, -1)
In [316]: timeit (functools.reduce(operator.mul,shape[:axis]),shape[axis],-1)
739 ns ± 30.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

比较提议的reduceat：

In [318]: tuple(np.multiply.reduceat(shape, (0, axis, axis+1)))
Out[318]: (6, 4, 5)
In [319]: timeit tuple(np.multiply.reduceat(shape, (0, axis, axis+1)))
11.3 µs ± 21.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

【讨论】：

带初始化器，它可以与axis=0一起工作。 (functools.reduce(operator.mul,shape[:axis], 1),shape[axis],-1)stackoverflow.com/questions/33945882/…
@hpaulj 我什至可以使用它;-)
@hamsteronwheels 仍然在左侧插入一个尴尬的 1。 OTOH，在某些情况下甚至可能是可取的

【解决方案2】：

如果你的轴真的在中间，你可以使用np.multiply.reduceat

 shape = (2, 3, 4, 5, 6)
 axis = 2
 np.multiply.reduceat(shape, (0, axis, axis+1))
 # array([ 6,  4, 30])
 axis = 3
 np.multiply.reduceat(shape, (0, axis, axis+1))
 # array([24,  5,  6])

但是，如果您想要第零个或最后一个轴，则必须是特殊情况。

【讨论】：