Копирование между массивами путем индексации

Question

Я хотел бы, чтобы скопировать элементы массива (a) в результирующий массив b в соответствии с некоторыми правилами, которые определяют, какой элемент из a идет туда, где в b. Я привел пример ниже.

Возможно ли (посредством некоторого умного индексации?) Выполнить последний шаг (b[x,mask] = a[mask]) за один шаг для всех x или это может быть достигнуто только в цикле, как показано (расширенным способом) ниже? x в моем реальном примере довольно большой (~ 100), поэтому я думаю, что это может извлечь выгоду из удаления цикла.

a = np.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) 

m1 = a > 6 
m2 = ~m1 & ((a < 8) & (a > 4)) 
m3 = ~m1 & ~m2 & (a > 1) 
m4 = ~m1 & ~m2 & ~m3 & (a < 10) 


b = np.zeros((4, 3, 3)) 
b[0, m1] = a[m1] 
b[1, m2] = a[m2] 
b[2, m3] = a[m3] 
b[3, m4] = a[m4] 

print '0\n', b[0] 
print '1\n', b[1] 
print '2\n', b[2] 
print '3\n', b[3] 

Выход:

0 
[[ 0. 0. 0.] 
[ 0. 0. 0.] 
[ 7. 8. 9.]] 
1 
[[ 0. 0. 0.] 
[ 0. 5. 6.] 
[ 0. 0. 0.]] 
2 
[[ 0. 2. 3.] 
[ 4. 0. 0.] 
[ 0. 0. 0.]] 
3 
[[ 1. 0. 0.] 
[ 0. 0. 0.] 
[ 0. 0. 0.]] 

Answer 1

Я не знаю, если это улучшает его или нет, но это можно сделать копию в один шаг, если мы строим распущенной индексы:

Start со списком масок:

In [1934]: m=[m1,m2,m3,m4] 

Сделать список индексов, присоединившись к отдельному списку, скорректирован для работы с плоская или распущенная 1d версия a:

In [1935]: ida=np.concatenate([np.ravel_multi_index(np.nonzero(z),a.shape) for z in m]) 
In [1936]: ida 
Out[1936]: array([6, 7, 8, 4, 5, 1, 2, 3, 0], dtype=int32) 
In [1937]: a.flat[ida] 
Out[1937]: array([7, 8, 9, 5, 6, 2, 3, 4, 1]) 

Построить же для 3d b, принимая во внимание различного 1-го значения измерений:

In [1938]: idb = [np.ravel_multi_index(((i,)+np.nonzero(z)),bshape) for i,z in enumerate(m)] 
In [1939]: idb 
Out[1939]: 
[array([6, 7, 8], dtype=int32), 
array([13, 14], dtype=int32), 
array([19, 20, 21], dtype=int32), 
array([27], dtype=int32)] 
In [1940]: idb=np.concatenate(idb) 

Теперь сделайте копию:

In [1941]: b.flat[id1]=a.flat[ida] 
In [1942]: b 
Out[1942]: 
array([[[ 0., 0., 0.], 
     [ 0., 0., 0.], 
     [ 7., 8., 9.]], 

     [[ 0., 0., 0.], 
     [ 0., 5., 6.], 
     ... 
     [ 0., 0., 0.]]]) 

Таким образом, копия выполняется одним вызовом, но для построения требуемых индексов мне пришлось использовать два контекста списка.