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?

Answer 1

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