numpy.choose

numpy.choose(a, choices, out=None, mode='raise')

Функция choose() создает массив на основе индексного массива из набора массивов.

Параметры:
a - массив NumPy, массивоподобный объект или число.
Должен содержать только целые числа в интервале [0, n - 1], где n - это количество массивов. В случае если mode не равно 'raise' может содержать любые целые числа. Если указано число, то из choices будет выбран массив с индексом равным этому числу.
choices - список или кортеж массивов NumPy или массивоподобных объектов.
Набор массивов, причем, все массивы и индексный массив a должны быть транслируемы в одну и ту же форму. Данный параметр может принимать массив NumPy, но в этом случае он все равно рассматривается как последовательность, а длина его первой оси (choices.shape[0]) будет определять длину этой последовательности.
out - массив NumPy (необязательный параметр).
Позволяет сразу напрямую поместить результат в указанный массив, при условии, что он имеет подходящую форму и тип данных.
mode - {'raise', 'wrap', 'clip'} (необязательный параметр).

Определяет метод обработки индексов, которые выходят за пределы формы массивов в choices.

  • 'raise' - вызывать исключение (по умолчанию);
  • 'wrap' - обогнуть вокруг оси, т.е. циклически сместиться по ней;
  • 'clip' - обрезает до диапазона индексов массивов в choices, причем отрицательные индексы обрезаются до 0.
Возвращает:
ndarray - массив NumPy
объединение массивов из множества choices в соответствии с массивом индексов в a.
Смотрите так же:
take, take_along_axis, compress

Замечание

Данная функция так же реализована в виде метода базового класса ndarray.choose() с аналогичной сигнатурой (ndarray.choose(choices, out=None, mode='raise')) и принципом работы.



Примеры

Данная функция работает довольно просто, допустим, у нас есть три массива из которых можно выбирать элементы и есть массив индексов который указывает какие именно элементы выбирать. Единственное что сбивает с толку, так это то, что само число в массиве индексов указывает не на элемент который нужно выбрать, а на массив из которого выбирается элемент, а вот индекс самого числа в индексном массиве определяет какой именно элемент будет выбран:

>>> import numpy as np
>>> 
>>> a = np.r_[:4]
>>> b = np.r_[-4:0]
>>> c = np.r_[100:500:100]
>>> a, b, c
(array([0, 1, 2, 3]), array([-4, -3, -2, -1]), array([100, 200, 300, 400]))
>>> 
>>> np.choose([0, 0, 0, 0], [a, b, c])
array([0, 1, 2, 3])

Глядя на результат, становится понятно, что из массивов b и c, мы вообще ничего не выбирали, потому что в массиве индексов указаны одни нули. Но ведь цифры то в результате разные?.. Дело в том что каждый из этих нулей имеет свой индекс, который и является критерием того какой именно элемент будет вытащен, т.е.индексы в a и choices совпадают, но числа в индексном массиве указывают из какого массива будет выбран элемент.

Если указать одни единицы, то мы получим все элементы из массива b:

>>> np.choose([1, 1, 1, 1], [a, b, c])
array([-4, -3, -2, -1])

А если одни двойки, то мы получим все элементы из массива c:

>>> np.choose([2, 2, 2, 2], [a, b, c])
array([100, 200, 300, 400])

Если же мы хотим вытащить все элементы кратные трем, то у нас остается всего один вариант индексоного массива:

>>> np.choose([0, 1, 2, 0], [a, b, c])
array([  0,  -3, 300,   3])

Чтож, с тем как работает функция мы разобрались, теперь рассмотрим как работает параметр mode. По умолчанию он установлен в значение 'raise', что приводит к логичному появлению ошибки если вдруг один из индексов в a выходит за пределы choices. Но такое поведение нужно не всегда, например, чрезмерно большие индексы могут быть урезаны до длины choices:

>>> np.choose([0, 1, 1, 2], [a, b, c])    #  все индексы в пределах choices
array([  0,  -3,  -2, 400])
>>> 
>>> np.choose([0, 1, 1, 4], [a, b, c], mode = 'clip')    #  4 обрезается до 2
array([  0,  -3,  -2, 400])

Если установить значение wrap то индекс будет циклически сдвигаться вдоль choices:

>>> np.choose([0, 1, 1, 7], [a, b, c], mode = 'wrap')    #  7 сдвигается до 1
array([ 0, -3, -2, -1])

Если параметр axis не указан, то извлекаются элементы из массива сжатого до одной оси:

>>> np.take(a, [0, -1])
array([ 0, 71])
>>> 
>>> np.take(a[1], [0, -1])
array([24, 47])

Если массив индексов содержит элементы чьи значения выходят за пределы исходного массива то вызывается исключение, но такое поведение желательно не всегда. Например эти значения могут просто обрезаться до максимальных размеров исходного массива:

>>> a = np.arange(7)
>>> a
array([0, 1, 2, 3, 4, 5, 6])
>>> 
>>> ind = [0, 1, 2, 10]
>>> 
>>> np.take(a, ind, mode = 'clip')
array([0, 1, 2, 6])

В примере выше, мы должны были увидеть сообщение об ошибке, так как последний элемент в ind явно превышает дину индексируемого массива. Но благодаря режиму 'clip' 10 было урезано до 6.

Режим 'wrap' в случае превышения размеров массива позволяет вести отсчет циклически смещаясь по оси:

>>> np.take(a, ind, mode = 'wrap')
array([0, 1, 2, 3])

Параметр out позволяет выгружать результат напрямую в указанный массив, минуя операцию присваивания, что позволяет несколько ускорить выполнение программы, особенно в циклах:

>>> out_arr = np.empty(4, dtype = np.int)
>>> out_arr
array([0, 1, 1, 2])
>>> 
>>> np.choose([0, 1, 1, 2], [a, b, c], out = out_arr)
array([  0,  -3,  -2, 400])
>>> 
>>> out_arr
array([  0,  -3,  -2, 400])

Обратите внимание на то, что форма и тип данных как возвращаемого массива, так и массива в который помещается результат должны полностью совпадать.

Массивы в a и choices вовсе не обязаны быть одномерным, единственное требование, все они должны быть транслируемы к единой форме. Самый простой случай, когда в choices находятся не массивы, а скалярные значения:

>>> choices = [-1, 1]
>>> a = [[0, 1, 0], [1, 1, 1], [0, 1, 0]]
>>> 
>>> np.choose(a, choices)
array([[-1,  1, -1],
       [ 1,  1,  1],
       [-1,  1, -1]])

Абсолютно все массивы могут иметь разную, но транслируему форму:

>>> a = np.array([2, 0, 1]).reshape(3, 1, 1)
>>> a
array([[[2]],

       [[0]],

       [[1]]])
>>> 
>>> c1 = 0
>>> c2 = np.full(3, 7).reshape(1, 3, 1)
>>> c2
array([[[7],
        [7],
        [7]]])
>>> 
>>> c3 = np.full(5, 3).reshape(1, 1, 5)
>>> c3
array([[[3, 3, 3, 3, 3]]])
>>> 
>>> 
>>> np.choose(a, [c1, c2, c3])    #  вернет массив с формой (3, 3, 5)
array([[[3, 3, 3, 3, 3],
        [3, 3, 3, 3, 3],
        [3, 3, 3, 3, 3]],

       [[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

       [[7, 7, 7, 7, 7],
        [7, 7, 7, 7, 7],
        [7, 7, 7, 7, 7]]])