我有一个 4 维数组,它是 3 维数组的时间序列。我想沿时间轴打乱三维数组中的每个点。这是我使用嵌套for
循环编写的代码。这可以通过花哨的 numpy 索引来完成吗?速度是一个因素。谢谢。
import numpy as np
timepoints = 2
x = 4
y = 4
z = 3
vol_1 = np.zeros((x, y, z))
vol_2 = np.ones((x, y, z))
timeseries = np.array((vol_1, vol_2))
timeseries.shape # (2, 4, 4, 3)
# One voxel over time.
timeseries[:, 0, 0, 0]
for xx in range(x):
for yy in range(y):
for zz in range(z):
np.random.shuffle(timeseries[:, xx, yy, zz])
我们可以沿第一个轴生成所有随机的索引,然后简单地使用advanced-indexing
来获得随机版本。现在,为了获得这些所有随机的索引,我们可以生成一个与输入数组形状相同的随机数组,并沿第一个轴获取 argsort 索引。这在之前已经探索过,如here
.
因此,我们将有一个这样的矢量化实现 -
m,n,r,p = a.shape # a is the input array
idx = np.random.rand(*a.shape).argsort(0)
out = a[idx, np.arange(n)[:,None,None], np.arange(r)[:,None], np.arange(p)]
只是为了向读者解释问题到底是什么,这里有一个示例运行 -
1) 输入 4D 数组:
In [711]: a
Out[711]:
array([[[[60, 22, 34],
[29, 18, 79]],
[[11, 69, 41],
[75, 30, 30]]],
[[[63, 61, 42],
[70, 56, 57]],
[[70, 98, 71],
[29, 93, 96]]]])
2) 使用沿第一轴索引的拟议方法生成的随机索引:
In [712]: idx
Out[712]:
array([[[[1, 0, 1],
[0, 1, 1]],
[[0, 0, 1],
[1, 0, 1]]],
[[[0, 1, 0],
[1, 0, 0]],
[[1, 1, 0],
[0, 1, 0]]]])
3)最后索引到输入数组进行随机输出:
In [713]: out
Out[713]:
array([[[[63, 22, 42],
[29, 56, 57]],
[[11, 69, 71],
[29, 30, 96]]],
[[[60, 61, 34],
[70, 18, 79]],
[[70, 98, 41],
[75, 93, 30]]]])
仔细观察,我们将看到63
a[0,0,0,0]
和60
由于idx
值分别在idx
的相应位置1
和0
,因此交换了a[1,0,0,0]
。接下来,22
和61
留在自己的位置,因为idx
值是0
和1
等等。
运行时测试
In [726]: timeseries = np.random.rand(10,10,10,10)
In [727]: %timeit org_app(timeseries)
100 loops, best of 3: 5.24 ms per loop
In [728]: %timeit proposed_app(timeseries)
1000 loops, best of 3: 289 µs per loop
In [729]: timeseries = np.random.rand(50,50,50,50)
In [730]: %timeit org_app(timeseries)
1 loop, best of 3: 720 ms per loop
In [731]: %timeit proposed_app(timeseries)
1 loop, best of 3: 426 ms per loop
在大尺寸下,创建随机数组的成本被证明是所提出的方法的瓶颈,但仍然显示出比原始循环版本良好的加速。
我添加这个作为答案,因为它不适合评论,因为它只是在@Divakar的出色答案之上的一个小补充:
def divakar(a):
m,n,r,p = a.shape # a is the input array
idx = np.random.rand(*a.shape).argsort(0)
return a[idx, np.arange(n)[:,None,None], np.arange(r)[:,None], np.arange(p)]
a = np.random.rand(50,50,50,50)
%timeit divakar(a)
560 ms ± 2.62 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
我观察到通过多次使用重塑而不是广播来加速一些加速,例如:
def norok2(a):
shape = a.shape
idx = np.random.rand(*a.shape).argsort(0).reshape(shape[0], -1)
return a.reshape(shape[0], -1)[idx, np.arange(shape[1] * shape[2] * shape[3])].reshape(shape)
a = np.random.rand(50,50,50,50)
%timeit norok2(a)
495 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
与OP的提案相比:
def jakub(a):
t, x, y, z = a.shape
for xx in range(x):
for yy in range(y):
for zz in range(z):
np.random.shuffle(a[:, xx, yy, zz])
%timeit jakub(a)
2 s ± 30.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
顺便说一下,我提出的修改更容易扩展到n维数组和任意洗牌轴,例如:
import numpy as np
import functools
def shuffle_axis(arr, axis=0):
arr = np.swapaxes(arr, 0, axis)
shape = arr.shape
i = np.random.rand(*shape).argsort(0).reshape(shape[0], -1)
return arr.reshape(shape[0], -1)[i, np.arange(functools.reduce(lambda x, y: x * y, shape[1:]))].reshape(shape).swapaxes(axis, 0)
具有类似的速度:
a = np.random.rand(50,50,50,50)
%timeit shuffle_axis(a)
499 ms ± 2.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
重新审视编辑
。而且时间并不比随机化所有东西更糟糕:
a = np.random.rand(50,50,50,50)
%timeit np.random.shuffle(a.ravel())
310 ms ± 1.84 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
这应该是此问题的任何解决方案的性能的某种下限(但它不能解决 OP 问题)。