Использование объектов диапазона Python для индексации массивов

Я видел это один или два раза раньше, но я не могу найти никаких официальных документов по нему: Использование Python range объекты как индексы в NumPy.

import numpy as np
a = np.arange(9).reshape(3,3)
a[range(3), range(2,-1,-1)]
# array([2, 4, 6])

Давайте запустим ошибку индекса, чтобы подтвердить, что диапазоны не находятся в официальном диапазоне (каламбур) законных методов индексации:

a['x']

# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

Теперь небольшое расхождение между NumPy и его документами не совсем неслыханно и не обязательно указывает на то, что функция не предназначена (см., Например, здесь).

Итак, кто-нибудь знает, почему это работает вообще? И если это предполагаемая особенность, какова точная семантика / для чего она нужна? И есть ли какие-либо обобщения НД?

2 ответа

Решение

Просто чтобы обернуть это (спасибо @WarrenWeckesser в комментариях): Это поведение на самом деле задокументировано. Нужно только понять, что range объекты являются последовательностями питона в строгом смысле.

Так что это всего лишь случай необычной индексации. Имейте в виду, однако, что это очень медленно:

>>> a = np.arange(100000)
>>> timeit(lambda: a[range(100000)], number=1000)
12.969507368048653
>>> timeit(lambda: a[list(range(100000))], number=1000)
7.990526253008284
>>> timeit(lambda: a[np.arange(100000)], number=1000)
0.22483703796751797

Не правильный ответ, но слишком длинный для комментариев.

На самом деле, он работает с любым индексируемым объектом:

import numpy as np

class MyIndex:
    def __init__(self, n):
        self.n = n
    def __getitem__(self, i):
        if i < 0 or i >= self.n:
            raise IndexError
        return i
    def __len__(self):
        return self.n

a = np.array([1, 2, 3])
print(a[MyIndex(2)])
# [1 2]

Я думаю, что соответствующие строки в коде NumPy находятся ниже этого комментария в core/src/multiarray/mapping.c:

/*
 * Some other type of short sequence - assume we should unpack it like a
 * tuple, and then decide whether that was actually necessary.
 */

Но я не совсем уверен. По какой-то причине это зависает, если вы удалите if i < 0 or i >= self.n: raise IndexError хотя есть __len__, так что в какой-то момент он, кажется, перебирает данный объект до IndexError Поднялся.

Другие вопросы по тегам