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.

$Matrix Forms$

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?

Divakar · Accepted Answer · 2016-07-05 22:47:18Z

5

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)

edited Jul 5, 2016 at 22:47

answered Jul 5, 2016 at 22:41

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

DMilden Over a year ago

Thanks! I didn't realize you can perform operations on otherwise incompatible matrix shapes in such an elegant manner like this. This helps considerably. :D

Collectives™ on Stack Overflow

Optimizing Array Element Shifting in Python / Numpy

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related