N维numpy数组的任意N维重复

Question

假设您有以下 N 维数组

>>> import numpy as np
>>> Z = np.array(7*[6*[5*[4*[3*[range(2)]]]]])
>>> Z.ndim
6

请注意 N = 6，但我想在讨论中保持任意性。

---

然后我执行多轴操作——当然是主观的——有问题“折叠”（如here 和there 所述）维度，计算这些维度.

假设计算轴是

>>> axes = (0,2,5)

因此元组的长度属于 [1,N]。

---

正如您可能已经猜到的那样，我想让np.mean 的输出的形状与其输入的形状相同。例如。

>>> Y = np.mean(Z, axis=axes)
>>> Y.shape
(6L, 4L, 3L)

同时

>>> Z.shape
(7L, 6L, 5L, 4L, 3L, 2L)

---

我有一个自制的解决方案，如下

def nd_repeat(arr, der, axes):
    if not isinstance(axes, tuple):
        axes = (axes,)
    shape = list(arr.shape)
    for axis in axes:
        shape[axis] = 1
    return np.zeros_like(arr) + der.reshape(shape) 

顺便说一句，der 代表“派生”。

>>> nd_repeat(Z, Y, axes).shape
(7L, 6L, 5L, 4L, 3L, 2L)

---

什么是 numpy-builtin 方式来完成这个 N 维重复？

---

性能问题，import timeit

homemade_s = """\
nd_repeat(Z, np.nanpercentile(Z, 99.9, axis=axes), axes)
"""
homemade_t = timeit.Timer(homemade_s, "from __main__ import nd_repeat,Z,axes,np").timeit(10000)

npbuiltin_s = """\
np.broadcast_to(np.nanpercentile(Z, 99.9, axis=axes, keepdims=True), Z.shape)
"""
npbuiltin_t = timeit.Timer(npbuiltin_s, "from __main__ import Z,axes,np").timeit(10000)

正如所料

>>> np.log(homemade_t/npbuiltin_t)
0.024082885343423521

我的解决方案比 hpaulj 的解决方案慢约 2.5%。

Answer 1

mean 有一个keepdims 参数保留压缩维度：

In [139]: shape=(2,3,4,5)
In [140]: x=np.arange(np.prod(shape)).reshape(shape)
In [141]: m=x.mean(axis=2, keepdims=True)
In [142]: m.shape
Out[142]: (2, 3, 1, 5)

现在很容易沿该维度复制m：

In [144]: m1=np.broadcast_to(m,shape)
In [145]: m1.shape
Out[145]: (2, 3, 4, 5)

repeat 和 tile 也是沿维度复制数组的便捷方式。

broadcast_to 扩展数组而不添加元素，只是通过改变形状和步幅：

In [146]: m1.strides
Out[146]: (120, 40, 0, 8)

repeat 增加数组的大小：

In [148]: m2=np.repeat(m, shape[2], axis=2)
In [149]: m2.shape
Out[149]: (2, 3, 4, 5)
In [150]: m2.strides
Out[150]: (480, 160, 40, 8)

m 可以在没有任何一个的情况下使用，如x-m。 m 在这里与x 一起广播。