a 2d numpy array element-wise raised to the power of another 2d array-CodePudding

I need to raise a large 2d numpy array element-wise to the power of another 2d array with the same dimensions. I tried to use numba to improve the speed. Is there an efficient implementation using numpy built-in functions instead of numba?

import numpy as np
from time import time


def elewise_power_raw(x, y):
    z = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            z[i,j] = x[i,j]**y[i,j]
    return z
                

import numba
@numba.njit
def elewise_power_numba(x, y):
    n = x.shape[0]
    z = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            z[i,j] = x[i,j]**y[i,j]
    return z
        
def measure_time(n=5000):
    x = np.random.rand(n, n)
    y = np.random.rand(n, n)

    t0 = time()
    elewise_power_raw(x, y)
    print('Raw: ', round(time() - t0, 2), 's' )

    t1 = time()
    elewise_power_numba(x, y)
    print('numba: ', round(time() - t1, 2), 's' )
    
measure_time(5000)
# Raw:  22.31 s
# numba:  1.4 s

CodePudding user response：

You can always vectorize it.

x = np.random.rand(5000, 5000)
y = np.random.rand(5000, 5000)

%timeit x**y
977 ms ± 7.01 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)