Search code examples
pythonnumpyvectorization

How to apply a function to each element of a linspace without using a for-loop


Starting with a linspace

t = np.linspace(0,10, 100)

and an array

a = np.array([1,2,44,2, 13,...])

I would like to get an array b, of the same length of the linspace whose elements are the array a raised to the power of the linspace elements (without using a for-loop). So you'd have something like

b[0]=np.array([1^t[0], 2^t[0], 44^t[0],...])
b[1]=np.array([1^t[1], 2^t[1], 44^t[1],...])

etc for the whole linspace.

Since I don't want to use a for-loop, is it possible to use np.apply_along_axis to perform a function on every element of t that gives me the final b I want? I've been struggling to do this, I think because I must just not understand how np.apply_along_axis works entirely.


Solution

  • Use broadcasting:

    If you raise an array a of shape (N,) or (1, N) (a row vector) to an array t of shape (M, 1) (a column vector), numpy automatically broadcasts their shapes and returns an array of shape (M, N), where the i, jth element of that array is a[j]**t[i]. To represent your t array as a column vector, index it as t[:, None], which adds a dimension after the current one.

    a = np.array([1,2,44,2, 13,...])
    t = np.linspace(0,10, 100)
    
    b = a**t[:, None]