What is a vectorized way to create multiple powers of a numpy array?

  • A+
Category:Languages

I have a numpy array:

arr = [[1, 2],        [3, 4]] 

I want to create a new array that contains powers of arr up to a power order:

# arr_new = [arr^0, arr^1, arr^2, arr^3,...arr^order] arr_new = [[1, 1, 1, 2, 1, 4, 1, 8],            [1, 1, 3, 4, 9, 16, 27, 64]] 

My current approach uses for loops:

# pre-allocate an array for powers arr = np.array([[1, 2],[3,4]]) order = 3 rows, cols = arr.shape arr_new = np.zeros((rows, (order+1) * cols))  # iterate over each exponent for i in range(order + 1):     arr_new[:, (i * cols) : (i + 1) * cols] = arr**i print(arr_new) 

Is there a faster (i.e. vectorized) approach to creating powers of an array?


Edit: Benchmarking

Thanks to @hpaulj and @Divakar and @Paul Panzer for the answers. I benchmarked the loop-based and broadcasting-based operations on the following test arrays.

arr = np.array([[1, 2],                 [3,4]]) order = 3 

The loop_based function is:

def loop_based(arr, order):     # pre-allocate an array for powers     rows, cols = arr.shape     arr_new = np.zeros((rows, (order+1) * cols))     # iterate over each exponent     for i in range(order + 1):         arr_new[:, (i * cols) : (i + 1) * cols] = arr**i     return arr_new 

The broadcast_based function using hstack is:

def broadcast_based_hstack(arr, order):     # create a 3D exponent array for a 2D input array to force broadcasting     powers = np.arange(order + 1)[:, None, None]     # generate values (3-rd axis contains array at various powers)     exponentiated = arr ** powers     # reshape and return array     return np.hstack(exponentiated)   # <== using hstack function 

The broadcast_based function using reshape is:

def broadcast_based_reshape(arr, order):     # create a 3D exponent array for a 2D input array to force broadcasting     powers = np.arange(order + 1)[:, None]     # generate values (3-rd axis contains array at various powers)     exponentiated = arr[:, None] ** powers     # reshape and return array     return exponentiated.reshape(arr.shape[0], -1)  # <== using reshape function 

The broadcast_based function using cumulative product cumprod and reshape:

def broadcast_cumprod_reshape(arr, order):     rows, cols = arr.shape     # Create 3D empty array where the middle dimension is     # the array at powers 0 through order     out = np.empty((rows, order + 1, cols), dtype=arr.dtype)     out[:, 0, :] = 1   # 0th power is always 1     a = np.broadcast_to(arr[:, None], (rows, order, cols))     # cumulatively multiply arrays so each multiplication produces the next order     np.cumprod(a, axis=1, out=out[:,1:,:])     return out.reshape(rows, -1) 

On jupyter notebook, I used the timeit command and got these results:

%timeit -n 100000 loop_based(arr, order) 7.41 µs ± 174 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)  %timeit -n 100000 broadcast_based_hstack(arr, order) 10.1 µs ± 137 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)  %timeit -n 100000 broadcast_based_reshape(arr, order) 3.31 µs ± 61.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)  %timeit -n 100000 broadcast_cumprod_reshape(arr, order) 11 µs ± 102 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) 

It seems that the broadcast based approach using reshape is the fastest by far.


Extend arrays to higher dims and let broadcasting do its magic with some help from reshaping -

In [16]: arr = np.array([[1, 2],[3,4]])  In [17]: order = 3  In [18]: (arr[:,None]**np.arange(order+1)[:,None]).reshape(arr.shape[0],-1) Out[18]:  array([[ 1,  1,  1,  2,  1,  4,  1,  8],        [ 1,  1,  3,  4,  9, 16, 27, 64]]) 

Note that arr[:,None] is essentially arr[:,None,:], but we can skip the trailing ellipsis for brevity.

Timings on a bigger dataset -

In [40]: np.random.seed(0)     ...: arr = np.random.randint(0,9,(100,100))     ...: order = 10  # @hpaulj's soln with broadcasting and stacking In [41]: %timeit np.hstack(arr **np.arange(order+1)[:,None,None]) 1000 loops, best of 3: 734 µs per loop  In [42]: %timeit (arr[:,None]**np.arange(order+1)[:,None]).reshape(arr.shape[0],-1) 1000 loops, best of 3: 401 µs per loop 

That reshaping part is practically free and that's where we gain performance here alongwith the broadcasting part of course, as seen in the breakdown below -

In [52]: %timeit (arr[:,None]**np.arange(order+1)[:,None]) 1000 loops, best of 3: 390 µs per loop  In [53]: %timeit (arr[:,None]**np.arange(order+1)[:,None]).reshape(arr.shape[0],-1) 1000 loops, best of 3: 401 µs per loop 

Comment

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen: