python - Sort multidimensional NumPy array by norm over a dimension -
i working multidimensional numpy array a "vector" of 2x2 matrices. want sort a such 2x2 matrices sorted row-norms.
import numpy np = np.array([[[3, 4], [1, 2]], [[5, 6], [7, 8]]]) sortidxs = np.argsort(np.linalg.norm(a, axis=-1)) = np.array([a[_][sortidxs[_]] _ in range(a.shape[0])]) # , final output should be: print(a) [[[1 2] [3 4]] [[5 6] [7 8]]] the above code snippet looking (not quite, @ edit below). looking way avoid loop
a = np.array([a[_][sortidxs[_]] _ in range(a.shape[0])]) -- edit --
the example above missed crucial part of of question.a have more "empty" dimensions i.e.
a = np.array([[[3, 4], [1, 2]], [[5, 6], [7, 8]]]) = a.reshape((2,1,2,2)) a looks like:
in [257]: out[257]: array([[[[3, 4], [1, 2]]], [[[5, 6], [7, 8]]]]) and after sorting should be
in [259]: out[259]: array([[[[1, 2], [3, 4]]], [[[5, 6], [7, 8]]]]) a have following dimension (1,2,2,2) or more such "empty" dimensions in beginning. want sort work in these cases too.
you can use advanced-indexing -
a[np.arange(a.shape[0])[:,none], sortidxs] sample run -
in [144]: = np.random.randint(0,9,(2,3,4)) in [145]: out[145]: array([[[1, 1, 5, 5], [1, 1, 7, 5], [6, 1, 2, 8]], [[7, 2, 5, 4], [3, 7, 3, 7], [8, 4, 4, 6]]]) in [146]: sortidxs = np.argsort(np.linalg.norm(a, axis=-1)) in [147]: np.array([a[_][sortidxs[_]] _ in range(a.shape[0])]) out[147]: array([[[1, 1, 5, 5], [1, 1, 7, 5], [6, 1, 2, 8]], [[7, 2, 5, 4], [3, 7, 3, 7], [8, 4, 4, 6]]]) in [149]: a[np.arange(a.shape[0])[:,none], sortidxs] out[149]: array([[[1, 1, 5, 5], [1, 1, 7, 5], [6, 1, 2, 8]], [[7, 2, 5, 4], [3, 7, 3, 7], [8, 4, 4, 6]]]) further performance boost
we can optimize on computing sortidxs np.einsum -
sortidxs = np.einsum('ijk,ijk->ij',a,a).argsort() let's time , verify idea -
in [94]: = np.random.randint(0,9,(20,30,40)) in [95]: %timeit np.argsort(np.linalg.norm(a, axis=-1)) 10000 loops, best of 3: 63.5 µs per loop in [96]: %timeit np.einsum('ijk,ijk->ij',a,a).argsort() 10000 loops, best of 3: 19.7 µs per loop in [97]: = np.random.randint(0,9,(200,300,400)) in [98]: %timeit np.argsort(np.linalg.norm(a, axis=-1)) 10 loops, best of 3: 88.6 ms per loop in [99]: %timeit np.einsum('ijk,ijk->ij',a,a).argsort() 10 loops, best of 3: 22.6 ms per loop array higher dimensions
for additional case of a being 4d array, need use more arrays indexing.
1] first axis : use np.arange(a.shape[0]) 2 new axes @ end.
2] second axis : use np.arange(a.shape[0]) 1 new axis @ end.
3] third axis : use sortidxs indexing this.
thus, have :
m,n,r,s = a.shape out = a[np.arange(m)[:,none,none],np.arange(n)[:,none], sortidxs] arrays singleton dim (dim length=1)
as special case, let's second axis of input array singleton one, use 0 axis , simplify things, -
a[np.arange(m)[:,none,none],0, sortidxs] sample run -
in [58]: = np.array([[[3, 4], ...: [1, 2]], ...: ...: [[5, 6], ...: [7, 8]]]) ...: ...: = a.reshape((2,1,2,2)) ...: in [59]: sortidxs = np.argsort(np.linalg.norm(a, axis=-1)) in [60]: a[np.arange(a.shape[0])[:,none,none],0, sortidxs] out[60]: array([[[[1, 2], [3, 4]]], [[[5, 6], [7, 8]]]]) another sample run array generic shape of (2,3,4) make things clear -
in [70]: = np.random.randint(0,9,(2,1,3,4)) in [71]: out[71]: array([[[[6, 4, 8, 6], [4, 0, 1, 0], [5, 3, 2, 5]]], [[[3, 6, 0, 4], [6, 2, 5, 2], [0, 8, 0, 8]]]]) in [72]: sortidxs = np.argsort(np.linalg.norm(a, axis=-1)) in [73]: sortidxs out[73]: array([[[1, 2, 0]], [[0, 1, 2]]]) in [74]: a[np.arange(a.shape[0])[:,none,none],0, sortidxs] out[74]: array([[[[4, 0, 1, 0], [5, 3, 2, 5], [6, 4, 8, 6]]], [[[3, 6, 0, 4], [6, 2, 5, 2], [0, 8, 0, 8]]]])
Comments
Post a Comment