Sélection et assignation de plusieurs sous-ensembles à partir de plusieurs ensembles en numpy

Question

Étant donné le tableau de base X de forme (2, 3, 4) qui peut être interprété comme deux ensembles de 3 éléments chacun, où chaque élément est en 4 dimensions, je veux échantillonner à partir de ce tableau X de la manière suivante. De chacun des 2 ensembles je veux choisir 2 sous-ensembles chacun défini par le tableau binaire de longueur 3, d'autres sous-ensembles seraient mis à 0. Donc le processus d'échantillonnage est défini par le tableau de forme (2, 2, 3). Le résultat de cet échantillonnage devrait avoir la forme (2, 2, 3, 4).

Voici le code qui fait ce dont j'ai besoin, mais je me demande si le pourrait être réécrit plus efficacement en utilisant l'indexation numérique.

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') 

sortie

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

étendent chacune d'entre elles à 4D en ajoutant de nouveaux axes pour subset_masks le long de la dernière et de sets que le deuxième axe. Pour ajouter ces nouveaux axes, nous pouvons utiliser None/np.newaxis. Ensuite, l'effet de levier NumPy broadcasting pour effectuer la multiplication élément par élément, comme si -

subset_masks[...,None]*sets[:,None] 

Juste pour les coups de pied probablement, nous pouvons aussi utiliser np.einsum -

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