How can I apply a function over last two dimensions? E.g. I generated an array below of (2,3,3) dimensions, the resulting array should have the same dimenstions where the function is apply to a[0,:,:] and a[1,:,:]. I understand I can go with a for loop, but might there be an in-built function specially for these type of operations?
a = np.arange(18).reshape(2,3,3)
c = []
for i in range(2):
c.append(np.linalg.pinv(a[i,:,:]))
result = np.array(c)
Here I used np.linalg.pinv but assume a generic function f: (n,k)-> (n,k), e.g f = lambda x: x**3 +61
>Solution :
Assuming a 3D array, you can directly map your function on the array, this will loop over the first dimension and apply the function on the remaining dimensions:
out = np.array(list(map(np.linalg.pinv, a)))
NB. this is not vectorized.
Output:
array([[[-5.55555556e-01, -1.66666667e-01, 2.22222222e-01],
[-5.55555556e-02, 1.98795266e-16, 5.55555556e-02],
[ 4.44444444e-01, 1.66666667e-01, -1.11111111e-01]],
[[-1.30555556e+00, -1.66666667e-01, 9.72222222e-01],
[-5.55555556e-02, 1.31569730e-15, 5.55555556e-02],
[ 1.19444444e+00, 1.66666667e-01, -8.61111111e-01]]])