 A+
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:
# preallocate 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 loopbased and broadcastingbased 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): # preallocate 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 (3rd 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 (3rd 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