Efficient way of serializing and mapping numpy array

user_4236833 0
Viewed 134

Assume that we have a 2d-array (such as a depth image). I'm looking for an efficient way to serialize this array and apply a function.

  1. array is a heightxwidth 2d-array.
  2. Serializing array --> list of [x,y,array[y,x]] for all combinations of x in range(0,width) and y in range(0,height).
  3. Applying a function func: func(x,y,array[y,x])

Example code:

import numpy as np
import itertools
width,height= 640,480
xstep,ystep= 1,1
array= np.array([np.sqrt(x*x+y*y) for y,x in itertools.product(range(0,height,ystep), range(0,width,xstep))]).reshape(height,width)
func= lambda x,y,z: [np.sqrt(x),np.sqrt(y),np.sqrt(x*x+y*y+z*z)]

I tried several approaches:

%%time
points1= [func(x,y,array[y,x]) for y,x in itertools.product(range(0,height,ystep), range(0,width,xstep))]
#--> Wall time: 1.15 s

%%time
points2= [func(x,y,array[y,x]) for y,x in itertools.product(np.arange(0,height,ystep), np.arange(0,width,xstep))]
#--> Wall time: 2.09 s

%%time
points3= [func(x,y,array[y,x]) for y,x in np.array(np.meshgrid(range(0,height,ystep),range(0,width,xstep))).T.reshape(-1,2)]
#--> Wall time: 2.16 s

points1==points2, points1==points3
#--> (True, True)

Is there a faster solution than these?

numpy arrays python
1 Answers

The problem is that you are computing square roots for individual elements, without taking advantage of vectorization. You can replace func by array operations:

def compute(array):
    indices = np.indices(array.shape)
    merged = np.empty((*array.shape, 3))
    merged[:,:,0] = np.sqrt(indices[1])
    merged[:,:,1] = np.sqrt(indices[0])
    merged[:,:,2] = np.sqrt(np.square(indices[1]) + np.square(indices[0]) + np.square(array))
    return merged.reshape((-1, 3)).tolist()

merged is a 3-dimensional array, containing values of func(x, y, array[y, x]) in merged[y, x, :]. Also it is convenient to convert this array to list by reshaping it to x*y by 3 array.

In fact, func is already vectorized. You just need to pass arrays into it instead of individual elements. However it can take advantage of using np.square:

func = lambda x,y,z: [np.sqrt(x),np.sqrt(y),np.sqrt(np.square(x)+np.square(y)+np.square(z))]

After these modifications creating the list runs about 15 times faster than before:

def compute(array):
    indices = np.indices(array.shape)
    merged = np.dstack(func(indices[1], indices[0], array))
    return merged.reshape((-1, 3)).tolist()

list1 = [func(x,y,array[y,x]) for y,x in itertools.product(range(0,height,ystep), range(0,width,xstep))]
list2 = [func(x,y,array[y,x]) for y, x in np.ndindex(array.shape)]
list3 = [func(x,y,value) for (y, x), value in np.ndenumerate(array)]
list4 = compute(array)
assert(list1 == list2 and list1 == list3 and list1 == list4)

timeit.timeit(lambda: [func(x,y,array[y,x]) for y, x in itertools.product(range(0,height,ystep), range(0,width,xstep))], number=1)
# 1.2148581651999848

timeit.timeit(lambda: [func(x,y,array[y,x]) for y, x in np.ndindex(array.shape)], number=1)
# 1.3134885265999856

timeit.timeit(lambda: [func(x,y,value) for (y, x), value in np.ndenumerate(array)], number=1)
# 1.2504884549999815

timeit.timeit(lambda: compute(array), number=1)
# 0.0784944145899999
user_5021440
Related