Optimizing Array Element Shifting in Python / Numpy

Question

Problem:

After running a line profiling of a data analysis code I have written, I have found that around 70% of the total run time is concentrated in calls to two different array manipulation routines. I would eventually like to analyze data in a real-time fashion, so any optimization here would help significantly.

The two functions take the matrix on the left and bring it to the form on the right (and vice versa).

The matrices I am interested in are currently stored as N by N 2d numpy arrays (where N is even).

Code:

I have written the following code to accomplish this:

# 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

I was unaware of numpy's roll function when originally writing this, so I had written the vec_shift functions in its place. They seem to have a ~30% performance increase over using roll on my current system.

Is there any way to further increase the performance of this code?

Answer 1

Let NumPy broadcasting help you out for a vectorized solution!

# 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)