Efficiently reshape numpy array

Question

I am working with NumPy arrays.

I have a 2N length vector D and want to reshape part of it into an N x N array C.

Right now this code does what I want, but is a bottleneck for larger N:

```

import numpy as np
M = 1000
t = np.arange(M)
D = np.sin(t)    # initial vector is a sin() function
N = M / 2
C = np.zeros((N,N))
for a in xrange(N):
    for b in xrange(N):
        C[a,b] = D[N + a - b]

```

Once C is made I go ahead and do some matrix arithmetic on it, etc.

This nested loop is pretty slow, but since this operation is essentially a change in indexing I figured that I could use NumPy's builtin reshape (numpy.reshape) to speed this part up.

Unfortunately, I cannot seem to figure out a good way of transforming these indices.

Any help speeding this part up?

Divakar · Accepted Answer · 2016-06-14 19:45:21Z

10

You can use NumPy broadcasting to remove those nested loops -

C = D[N + np.arange(N)[:,None] - np.arange(N)]

One can also use np.take to replace the indexing, like so -

C = np.take(D,N + np.arange(N)[:,None] - np.arange(N))

A closer look reveals the pattern to be close to toeplitz and hankel matrices. So, using those, we would have two more approaches to solve it, though with comparable speedups as with broadcasting. The implementations would look something like these -

from scipy.linalg import toeplitz
from scipy.linalg import hankel

C = toeplitz(D[N:],np.hstack((D[0],D[N-1:0:-1])))
C = hankel(D[1:N+1],D[N:])[:,::-1]

Runtime test

In [230]: M = 1000
     ...: t = np.arange(M)
     ...: D = np.sin(t)    # initial vector is a sin() function
     ...: N = M / 2
     ...: 

In [231]: def org_app(D,N):
     ...:     C = np.zeros((N,N))
     ...:     for a in xrange(N):
     ...:         for b in xrange(N):
     ...:             C[a,b] = D[N + a - b]
     ...:     return C
     ...: 

In [232]: %timeit org_app(D,N)
     ...: %timeit D[N + np.arange(N)[:,None] - np.arange(N)]
     ...: %timeit np.take(D,N + np.arange(N)[:,None] - np.arange(N))
     ...: %timeit toeplitz(D[N:],np.hstack((D[0],D[N-1:0:-1])))
     ...: %timeit hankel(D[1:N+1],D[N:])[:,::-1]
     ...: 
10 loops, best of 3: 83 ms per loop
100 loops, best of 3: 2.82 ms per loop
100 loops, best of 3: 2.84 ms per loop
100 loops, best of 3: 2.95 ms per loop
100 loops, best of 3: 2.93 ms per loop

edited Jun 14, 2016 at 19:45

answered Jun 14, 2016 at 19:06

Divakar

222k19 gold badges273 silver badges374 bronze badges

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

2 Comments

jjgoings Over a year ago

Awesome! This provided a considerable speed-up over the looping.

Divakar Over a year ago

@jjgoings Yup, confirmed that!

Collectives™ on Stack Overflow

Efficiently reshape numpy array

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related