numpyで複数の集合から複数の部分集合を選んで割り当てよう

Question

各要素が4次元である2組の3つの要素として解釈できる形状（2、3、4）の基本配列Xが与えられたら、この配列Xから次のようにして 2つのセットのそれぞれから、長さ3のバイナリ配列で定義された2つのサブセットを選択したい場合、他のサブセットは0に設定されます。したがって、サンプリングプロセスは形状の配列（2、2、3）で定義されます。このサンプリングの結果は、形状（2、2、3、4）でなければなりません。

私が必要とするコードはここにありますが、がnumpyインデックスを使用してより効率的に書き換えられるかどうか疑問に思っています。

import numpy as np 
np.random.seed(3) 

sets = np.random.randint(0, 10, [2, 3, 4]) 
subset_masks = np.random.randint(0, 2, [2, 2, 3]) 

print('Base set\n', sets, '\n') 
print('Subset masks\n', subset_masks, '\n') 

result = np.empty([2, 2, 3, 4]) 
for set_index in range(sets.shape[0]): 
    for subset_index, subset in enumerate(subset_masks[set_index]): 
     print('----') 
     picked_subset = subset.reshape(3, 1) * sets[set_index] 
     result[set_index][subset_index] = picked_subset 
     print('Picking subset ', subset, 'from set #', set_index) 
     print(picked_subset, '\n') 

出力

Base set 
[[[8 9 3 8] 
    [8 0 5 3] 
    [9 9 5 7]] 

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

Subset masks 
[[[0 0 1] 
    [1 0 0]] 

[[1 0 1] 
    [0 1 1]]] 

---- 
Picking subset [0 0 1] from set # 0 
[[0 0 0 0] 
[0 0 0 0] 
[9 9 5 7]] 

---- 
Picking subset [1 0 0] from set # 0 
[[8 9 3 8] 
[0 0 0 0] 
[0 0 0 0]] 

---- 
Picking subset [1 0 1] from set # 1 
[[6 0 4 7] 
[0 0 0 0] 
[2 1 3 5]] 

---- 
Picking subset [0 1 1] from set # 1 
[[0 0 0 0] 
[8 1 6 2] 
[2 1 3 5]] 

Answer 1

は最後に沿って第2の軸としてsetsためsubset_masksための新たな軸を追加することによって4Dにそれらの各々を拡張。これらの新しい軸を追加するにはNone/np.newaxisを使用できます。次に、そのような要素毎の乗算を実行するために活用しNumPy broadcasting - 両方のオプションが動作し

np.einsum('ijk,ilj->iljk',sets,subset_masks)