加快矩阵中每个x，y点的角度计算

Question

I have a 3-d Numpy array flow as follows:

flow = np.random.uniform(low=-1.0, high=1.0, size=(720,1280,2))
# Suppose flow[0] are x-coordinates. flow[1] are y-coordinates.

Need to calculate the angle for each x,y point. Here is how I have implemented it:

def calcAngle(a):
    assert(len(a) == 2)
    (x, y) = a
    # angle_deg = 0
    angle_deg = np.angle(x + y * 1j, deg=True)
    return angle_deg

fangle = np.apply_along_axis(calcAngle, axis=2, arr=flow) 
# The above statement takes 14.0389318466 to execute

The calculation of angle at each point takes 14.0389318466 seconds to execute on my Macbook Pro.

Is there a way I could speed this up, probably by using some matrix operation, rather than processing each pixel one at a time.

Answer 1

numpy.angle supports vectorized operation. So, just feed in the first and second column slices to it for the final output, like so -

fangle = np.angle(flow[...,0] + flow[...,1] * 1j, deg=True)

Verification -

In [9]: flow = np.random.uniform(low=-1.0, high=1.0, size=(720,1280,2))

In [17]: out1 = np.apply_along_axis(calcAngle, axis=2, arr=flow)

In [18]: out2 = np.angle(flow[...,0] + flow[...,1] * 1j, deg=True)

In [19]: np.allclose(out1, out2)
Out[19]: True

Runtime test -

In [10]: %timeit np.apply_along_axis(calcAngle, axis=2, arr=flow)
1 loop, best of 3: 8.27 s per loop

In [11]: %timeit np.angle(flow[...,0] + flow[...,1] * 1j, deg=True)
10 loops, best of 3: 47.6 ms per loop

In [12]: 8270/47.6
Out[12]: 173.73949579831933

173x+ speedup!

Answer 2

You can use numpy.arctan2() to get the angle in radians, and then convert to degrees with numpy.rad2deg() :

fangle = np.rad2deg(np.arctan2(flow[:,:,1], flow[:,:,0]))

On my computer, this is a little faster than Divakar's version:

In [17]: %timeit np.angle(flow[...,0] + flow[...,1] * 1j, deg=True)
10 loops, best of 3: 44.5 ms per loop

In [18]: %timeit np.rad2deg(np.arctan2(flow[:,:,1], flow[:,:,0]))
10 loops, best of 3: 35.4 ms per loop

A more efficient way to use np.angle() is to create a complex view of flow . If flow is an array of type np.float64 with shape (m, n, 2) , then flow.view(np.complex128)[:,:,0] will be an array of type np.complex128 with shape (m, n) :

fangle = np.angle(flow.view(np.complex128)[:,:,0], deg=True)

This appears to be a smidge faster than using arctan2 followed by rad2deg (but the difference is not far above the measurement noise of timeit ):

In [47]: %timeit np.angle(flow.view(np.complex128)[:,:,0], deg=True)
10 loops, best of 3: 35 ms per loop

Note that this might not work if flow was creating as the tranpose of some other array, or as a slice of another array using steps bigger than 1.