Numpy linalg.norm with ufunc.reduceat functionality

105 Views Asked by At

Solution: @QuangHoang's first comment namely np.linalg.norm(arr,axis=1).

I would like to apply Numpy's linalg.norm function column wise to sub-arrays of a 3D array by using ranges (or indices?), similar in functionality to what ufunc.reduceat does.

Given the following array:

import numpy as np

In  []: arr = np.array([[0,1,2,3], [2,2,3,4], [3,2,5,6],
                        [1,7,1,9], [1,4,8,6], [2,3,5,8],
                        [2,5,7,3], [2,3,4,6], [2,5,3,2]]).reshape(3,3,4)


Out []: array([[[0, 1, 2, 3],
                [2, 2, 3, 4],
                [3, 2, 5, 6]],

               [[1, 7, 1, 9],
                [1, 4, 8, 6],
                [2, 3, 5, 8]],

               [[2, 5, 7, 3],
                [2, 3, 4, 6],
                [2, 5, 3, 2]]])

I would like to apply linalg.norm column wise to the three sub-arrays separately i.e. for the first column it would be linalg.norm([0, 2, 3]), linalg.norm([1, 1, 2]) and linalg.norm([2, 2, 2]), for the second linalg.norm([1, 2, 2]), linalg.norm([7, 4, 3]) and linalg.norm([5, 3, 5]) etc. resulting in a 2D vector with shape (3,4) containing the results of the linalg.norm calls.

Doing this with a 2D array is straightforward by specifying the axis:

import numpy.linalg as npla

In  []: npla.norm(np.array([[0,1,2,3], [2,2,3,4], [3,2,5,6]]), axis=0)
Out []: array([3.60555128, 3.        , 6.164414  , 7.81024968])

But I don't understand how to do that for each sub-array separately. I believe that reduceat with a ufunc like add allows to set indices and ranges. Would something similar be possible here but with linalg.norm?

Edit 1:

I followed @hpaulj's advice to look at the code used for add.reduce. Getting a better understanding of the method I was able to search more precisely and I found np.apply_along_axis which is exactly what I was looking for:

In  []: np.apply_along_axis(npla.norm, 1, arr)
Out []: array([[ 3.60555128,  3.        ,  6.164414  ,  7.81024968],
               [ 2.44948974,  8.60232527,  9.48683298, 13.45362405],
               [ 3.46410162,  7.68114575,  8.60232527,  7.        ]])

However, this method is very slow. Is there a way to use linalg.nrom in a vectorized manner instead?

Edit 2:

@QuangHoang's first comment is actually the correct answer I was looking for. I misunderstood the method which is why I misunderstood their comment. Specifying the axis in the linalg.norm call is what is required here:

np.linalg.norm(arr,axis=1)
0

There are 0 best solutions below