Использование numpy.take для более быстрой причудливой индексации

EDIT Я сохранил более сложную проблему, с которой я столкнулся ниже, но мои проблемы с np.take можно резюмировать следующим образом. Скажем, у вас есть массив img формы (planes, rows) и еще один массив lut формы (planes, 256), и вы хотите использовать их для создания нового массива out формы (planes, rows), где out[p,j] = lut[p, img[p, j]], Это может быть достигнуто при фантастическом индексировании следующим образом:

In [4]: %timeit lut[np.arange(planes).reshape(-1, 1), img]
1000 loops, best of 3: 471 us per loop

Но если вместо фантазии индексирования вы используете take и цикл python над planes, события могут значительно ускоряться:

In [6]: %timeit for _ in (lut[j].take(img[j]) for j in xrange(planes)) : pass
10000 loops, best of 3: 59 us per loop

Может ли lut и img быть каким-то образом перестроенным, чтобы вся операция выполнялась без циклов python, но используя numpy.take (или альтернативный метод) вместо обычной причудливой индексации, чтобы сохранить преимущество скорости?

ОРИГИНАЛЬНЫЙ ВОПРОС У меня есть набор поисковых таблиц (LUT), которые я хочу использовать на изображении. Массив, содержащий LUT, имеет форму (planes, 256, n), а изображение имеет форму (planes, rows, cols). Оба имеют dtype = 'uint8', соответствующие оси 256 LUT. Идея состоит в том, чтобы запустить p -й плоскость изображения через каждый из n LUTs из p -й плоскости LUT.

Если мои lut и img следующие:

planes, rows, cols, n = 3, 4000, 4000, 4
lut = np.random.randint(-2**31, 2**31 - 1,
                        size=(planes * 256 * n // 4,)).view('uint8')
lut = lut.reshape(planes, 256, n)
img = np.random.randint(-2**31, 2**31 - 1,
                    size=(planes * rows * cols // 4,)).view('uint8')
img = img.reshape(planes, rows, cols)

Я могу добиться того, что я получаю после использования причудливого индексирования, такого как

out = lut[np.arange(planes).reshape(-1, 1, 1), img]

который дает мне массив формы (planes, rows, cols, n), где out[i, :, :, j] содержит i -я плоскость img, проходящая через j -й LUT i -й плоскости LUT...

Все хорошо, кроме этого:

In [2]: %timeit lut[np.arange(planes).reshape(-1, 1, 1), img]
1 loops, best of 3: 5.65 s per loop

что совершенно неприемлемо, тем более, что у меня есть все следующие не очень приятные альтернативы, использующие np.take, чем работающие быстрее:

Один LUT на одной плоскости быстрее работает примерно на x70:
```
In [2]: %timeit np.take(lut[0, :, 0], img[0])
10 loops, best of 3: 78.5 ms per loop
```

Цикл питона, проходящий через все нужные комбинации, быстрее заканчивается x6:

In [2]: %timeit for _ in (np.take(lut[j, :, k], img[j]) for j in xrange(planes) for k in xrange(n)) : pass
1 loops, best of 3: 947 ms per loop

Даже запуск всех комбинаций плоскостей в LUT и изображении и последующее отбрасывание нежелательных planes**2 - planes быстрее, чем фантастическое индексирование:
```
In [2]: %timeit np.take(lut, img, axis=1)[np.arange(planes), np.arange(planes)]
1 loops, best of 3: 3.79 s per loop
```
И самая быстрая комбинация, которую я смог придумать, имеет цикл python, итератирующий по плоскостям, и быстрее заканчивает x13:
```
In [2]: %timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass
1 loops, best of 3: 434 ms per loop
```

Вопрос, конечно, в том, что нет никакого способа сделать это с помощью np.take без какой-либо петли питона? В идеале любое изменение формы или изменение размера должно происходить на LUT, а не на изображении, но я открыт всем, что вы можете придумать...

Ответ 1

Кулак всего, что я должен сказать, мне очень понравился ваш вопрос. Без перестройки LUT или IMG работало следующее решение:

%timeit a=np.take(lut, img, axis=1)
# 1 loops, best of 3: 1.93s per loop

Но из результата вам нужно запросить диагональ: a [0,0], a [1,1], a [2,2]; чтобы получить то, что вы хотите. Я попытался найти способ сделать это индексирование только для диагональных элементов, но все равно не справился.

Вот несколько способов переупорядочить ваши LUT и IMG: Следующее работает, если индексы в IMG равны 0-255 для 1-й плоскости, 256-511 для 2-й плоскости и 512-767 для третьей плоскости, но это помешает вам использовать 'uint8', что может быть большой проблемой...:

lut2 = lut.reshape(-1,4)
%timeit np.take(lut2,img,axis=0)
# 1 loops, best of 3: 716 ms per loop
# or
%timeit np.take(lut2, img.flatten(), axis=0).reshape(3,4000,4000,4)
# 1 loops, best of 3: 709 ms per loop

в моей машине ваше решение по-прежнему является лучшим вариантом и очень адекватным, поскольку вам просто нужны диагональные оценки, то есть плоскость 1-плоскость1, плоскость 2-плоскость2 и плоскость 3-плоскость3:

%timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass
# 1 loops, best of 3: 677 ms per loop

Надеюсь, это может дать вам некоторое представление о лучшем решении. Было бы неплохо найти дополнительные опции с flatten() и аналогичными методами, такими как np.apply_over_axes() или np.apply_along_axis(), которые кажутся многообещающими.

Я использовал этот код ниже для генерации данных:

import numpy as np
num = 4000
planes, rows, cols, n = 3, num, num, 4
lut = np.random.randint(-2**31, 2**31-1,size=(planes*256*n//4,)).view('uint8')
lut = lut.reshape(planes, 256, n)
img = np.random.randint(-2**31, 2**31-1,size=(planes*rows*cols//4,)).view('uint8')
img = img.reshape(planes, rows, cols)