CODEWITHSUNDEEP
pythonperformancenumpymultidimensional-arrayvectorization
threnna
threnna

Reputation: 61

How to turn a 1D radial profile into a 2D array in python

I have a list that models a phenomenon that is a function of radius. I want to convert this to a 2D array. I wrote some code that does exactly what I want, but since it uses nested for loops, it is quite slow.

l = len(profile1D)/2
critDim = int((l**2 /2.)**(1/2.))
profile2D = np.empty([critDim, critDim])
for x in xrange(0, critDim):
    for y in xrange(0,critDim):
        r = ((x**2 + y**2)**(1/2.))
        profile2D[x,y] = profile1D[int(l+r)]

Is there a more efficient way to do the same thing by avoiding these loops?

Upvotes: 2

Views: 825

Answers (1)

Divakar
Divakar

Reputation: 221614

Here's a vectorized approach using broadcasting -

a = np.arange(critDim)**2
r2D = np.sqrt(a[:,None] + a)
out = profile1D[(l+r2D).astype(int)]

If there are many repeated indices generated by l+r2D, we can use np.take for some further performance boost, like so -

out = np.take(profile1D,(l+r2D).astype(int))

Runtime test

Function definitions -

def org_app(profile1D,l,critDim):
    profile2D = np.empty([critDim, critDim])
    for x in xrange(0, critDim):
        for y in xrange(0,critDim):
            r = ((x**2 + y**2)**(1/2.))
            profile2D[x,y] = profile1D[int(l+r)]
    return profile2D

def vect_app1(profile1D,l,critDim):
    a = np.arange(critDim)**2
    r2D = np.sqrt(a[:,None] + a)
    out = profile1D[(l+r2D).astype(int)]
    return out

def vect_app2(profile1D,l,critDim):
    a = np.arange(critDim)**2
    r2D = np.sqrt(a[:,None] + a)
    out = np.take(profile1D,(l+r2D).astype(int))
    return out

Timings and verification -

In [25]: # Setup input array and params
    ...: profile1D = np.random.randint(0,9,(1000))
    ...: l = len(profile1D)/2
    ...: critDim = int((l**2 /2.)**(1/2.))
    ...: 

In [26]: np.allclose(org_app(profile1D,l,critDim),vect_app1(profile1D,l,critDim))
Out[26]: True

In [27]: np.allclose(org_app(profile1D,l,critDim),vect_app2(profile1D,l,critDim))
Out[27]: True

In [28]: %timeit org_app(profile1D,l,critDim)
10 loops, best of 3: 154 ms per loop

In [29]: %timeit vect_app1(profile1D,l,critDim)
1000 loops, best of 3: 1.69 ms per loop

In [30]: %timeit vect_app2(profile1D,l,critDim)
1000 loops, best of 3: 1.68 ms per loop

In [31]: # Setup input array and params
    ...: profile1D = np.random.randint(0,9,(5000))
    ...: l = len(profile1D)/2
    ...: critDim = int((l**2 /2.)**(1/2.))
    ...: 

In [32]: %timeit org_app(profile1D,l,critDim)
1 loops, best of 3: 3.76 s per loop

In [33]: %timeit vect_app1(profile1D,l,critDim)
10 loops, best of 3: 59.8 ms per loop

In [34]: %timeit vect_app2(profile1D,l,critDim)
10 loops, best of 3: 59.5 ms per loop

Upvotes: 2

Related Questions