Trier un tableau numpy par un autre tableau, le long d'un axe particulier
Similaire à cette réponse, j'ai une paire de 3D tableaux numpy, a et b, et je veux trier les entrées de b par les valeurs de a. Contrairement à cette réponse , Je ne veux trier que le long d'un axe des tableaux.
Ma lecture naïve de la documentation numpy.argsort():
Returns
-------
index_array : ndarray, int
Array of indices that sort `a` along the specified axis.
In other words, ``a[index_array]`` yields a sorted `a`.
M'a amené à croire que je pouvais faire mon tri avec le code suivant:
import numpy
a = numpy.zeros((3, 3, 3))
a += numpy.array((1, 3, 2)).reshape((3, 1, 1))
print "a"
print a
"""
[[[ 1. 1. 1.]
[ 1. 1. 1.]
[ 1. 1. 1.]]
[[ 3. 3. 3.]
[ 3. 3. 3.]
[ 3. 3. 3.]]
[[ 2. 2. 2.]
[ 2. 2. 2.]
[ 2. 2. 2.]]]
"""
b = numpy.arange(3*3*3).reshape((3, 3, 3))
print "b"
print b
"""
[[[ 0 1 2]
[ 3 4 5]
[ 6 7 8]]
[[ 9 10 11]
[12 13 14]
[15 16 17]]
[[18 19 20]
[21 22 23]
[24 25 26]]]
"""
print "a, sorted"
print numpy.sort(a, axis=0)
"""
[[[ 1. 1. 1.]
[ 1. 1. 1.]
[ 1. 1. 1.]]
[[ 2. 2. 2.]
[ 2. 2. 2.]
[ 2. 2. 2.]]
[[ 3. 3. 3.]
[ 3. 3. 3.]
[ 3. 3. 3.]]]
"""
##This isnt' working how I'd like
sort_indices = numpy.argsort(a, axis=0)
c = b[sort_indices]
"""
Desired output:
[[[ 0 1 2]
[ 3 4 5]
[ 6 7 8]]
[[18 19 20]
[21 22 23]
[24 25 26]]
[[ 9 10 11]
[12 13 14]
[15 16 17]]]
"""
print "Desired shape of b[sort_indices]: (3, 3, 3)."
print "Actual shape of b[sort_indices]:"
print c.shape
"""
(3, 3, 3, 3, 3)
"""
Quelle est la bonne façon de faire?
1 réponses
Vous devez toujours fournir des indices pour les deux autres dimensions pour que cela fonctionne correctement.
>>> a = numpy.zeros((3, 3, 3))
>>> a += numpy.array((1, 3, 2)).reshape((3, 1, 1))
>>> b = numpy.arange(3*3*3).reshape((3, 3, 3))
>>> sort_indices = numpy.argsort(a, axis=0)
>>> static_indices = numpy.indices((3, 3, 3))
>>> b[sort_indices, static_indices[1], static_indices[2]]
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8]],
[[18, 19, 20],
[21, 22, 23],
[24, 25, 26]],
[[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17]]])
numpy.indices calcule les indices de chaque axe du tableau lorsqu'il est "aplati" à travers les deux autres axes (ou n-1 axes où n = nombre total d'axes). En d'autres termes, ceci (excuses pour le long post):
>>> static_indices
array([[[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]],
[[1, 1, 1],
[1, 1, 1],
[1, 1, 1]],
[[2, 2, 2],
[2, 2, 2],
[2, 2, 2]]],
[[[0, 0, 0],
[1, 1, 1],
[2, 2, 2]],
[[0, 0, 0],
[1, 1, 1],
[2, 2, 2]],
[[0, 0, 0],
[1, 1, 1],
[2, 2, 2]]],
[[[0, 1, 2],
[0, 1, 2],
[0, 1, 2]],
[[0, 1, 2],
[0, 1, 2],
[0, 1, 2]],
[[0, 1, 2],
[0, 1, 2],
[0, 1, 2]]]])
Ce sont les indices d'identité pour chaque axe; lorsqu'ils sont utilisés pour indexer b, ils recréent B.
>>> b[static_indices[0], static_indices[1], static_indices[2]]
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8]],
[[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17]],
[[18, 19, 20],
[21, 22, 23],
[24, 25, 26]]])
, Comme une alternative à numpy.indices, vous pouvez utiliser numpy.ogrid, comme le suggère unutbu. Puisque l'objet généré par ogrid est plus petit, je vais créer les trois axes, juste pour des raisons de cohérence, mais notez le commentaire d'unutbu pour un moyen de le faire en générant seulement deux.
>>> static_indices = numpy.ogrid[0:a.shape[0], 0:a.shape[1], 0:a.shape[2]]
>>> a[sort_indices, static_indices[1], static_indices[2]]
array([[[ 1., 1., 1.],
[ 1., 1., 1.],
[ 1., 1., 1.]],
[[ 2., 2., 2.],
[ 2., 2., 2.],
[ 2., 2., 2.]],
[[ 3., 3., 3.],
[ 3., 3., 3.],
[ 3., 3., 3.]]])