优化Python / Numpy中的数组元素移位

Question

问题：

在运行我编写的数据分析代码的行分析后，我发现大约70％的总运行时间集中在对两个不同的数组操作例程的调用上。我最终希望以实时方式分析数据，因此这里的任何优化都会有很大帮助。

这两个函数采用左边的矩阵并将其带到右边的表格（反之亦然）。

我感兴趣的矩阵目前存储为N个N 2d numpy数组（其中N是偶数）。

代码：

我已经编写了以下代码来完成此任务：

# Shifts elements of a vector to the left by the given amount.
def Vec_shift_L(vec, shift=0):
    s = vec.size
    out = np.zeros(s, dtype=complex)
    out[:s-shift] = vec[shift:]
    out[s-shift:] = vec[:shift]
    return out

# Shifts elements of a vector to the right by the given amount.
def Vec_shift_R(vec,shift=0):
    s=vec.size
    out=np.zeros(s, dtype=complex)
    out[:shift] = vec[s-shift:]
    out[shift:] = vec[:s-shift]
    return out

# Shifts a matrix from the left form (above) to the right form.
def OP_Shift(Trace):
    s = Trace.shape
    Out = np.zeros(s, dtype=complex)

    for i in np.arange(s[0]):
        Out[i,:] = Vec_shift_L(Trace[i,:], (i+s[0]/2) % s[0])

    for i in np.arange(s[0]):
        Out[i,:] = np.flipud(Out[i,:])

    return Out

# Shifts a matrix from the right form (above) to the left form.
def iOP_Shift(Trace):
    s = Trace.shape
    Out = np.zeros(s, dtype=complex)
    for i in np.arange(s[0]):
        Out[i,:] = np.flipud(Trace[i,:])

    for i in np.arange(s[0]):
        Out[i,:] = Vec_shift_R(Out[i,:], (i+s[0]/2) % s[0])

    return Out

我最初写这篇文章的时候并没有意识到numpy的滚动功能，所以我写了vec_shift函数。与使用当前系统上的roll相比，它们的性能似乎提高了约30％。

有没有办法进一步提高此代码的性能？

Answer 1

让NumPy broadcasting帮助您提供矢量化解决方案！

# Store shape of input array
s = Trace.shape

# Store arrays corresponding to row and column indices
I = np.arange(s[0])
J = np.arange(s[1]-1,-1,-1)

# Store all iterating values in "(i+s[0]/2) % s[0]" as an array
shifts = (I + s[0]/2)%s[0]

# Calculate all 2D linear indices corresponding to 2D transformed array
linear_idx = (((shifts[:,None] + J)%s[1]) + I[:,None]*s[1])

# Finally index into input array with indices for final output
out = np.take(Trace,linear_idx)