Façon rapide de définir les diagonales d'une matrice (M x N x N)? Einsum/n-dimensionnel fill_diagonal?

Question

J'essaie d'écrire du code rapide et optimisé basé sur des matrices, et j'ai récemment découvert l'einsum comme un outil permettant une accélération significative.

Est-il possible de l'utiliser pour définir efficacement les diagonales d'une matrice multidimensionnelle, ou peut-elle seulement renvoyer des données?

Dans mon problème, j'essaie de définir les diagonales pour un tableau de matrices carrées (forme: M x N x N) en sommant les colonnes dans chaque matrice carrée (N x N).

Ma solution actuelle (slow, boucle) est:

# Build dummy array 
dimx = 2 # Dimension x (likely to be < 100) 
dimy = 3 # Dimension y (likely to be between 2 and 10) 
M = np.random.randint(low=1, high=9, size=[dimx, dimy, dimy]) 

# Blank the diagonals so we can see the intended effect 
np.fill_diagonal(M[0], 0) 
np.fill_diagonal(M[1], 0) 

# Compute diagonals based on summing columns 
diags = np.einsum('ijk->ik', M) 

# Set the diagonal for each matrix 
# THIS IS LOW. CAN IT BE IMPROVED? 
for i in range(len(M)): 
    np.fill_diagonal(M[i], diags[i]) 

# Print result 
M 

peut-il être amélioré du tout s'il vous plaît? Il semble que np.fill_diagonal n'accepte pas les matrices non carrées (d'où la nécessité de forcer ma solution en boucle). Peut-être einsum peut-il aider ici aussi?

Answer 1

Une approche serait de remodeler à 2D, définir les colonnes aux étapes de ncols+1 avec les valeurs diagonales. Le remodelage crée une vue et, en tant que tel, nous permet d'accéder directement à ces positions diagonales. Ainsi, la mise en œuvre serait -

s0,s1,s2 = M.shape 
M.reshape(s0,-1)[:,::s2+1] = diags 

Answer 2

Si vous np.source(np.fill_diagonal) vous verrez que dans le cas 2d, il utilise une solution approche strided '

if a.ndim == 2: 
     step = a.shape[1] + 1 
     end = a.shape[1] * a.shape[1] 
    a.flat[:end:step] = val 

@Divakar's applique ceci à votre cas 3D par "aplatissement" sur 2 dimensions. Vous pouvez additionner les colonnes M.sum(axis=1). Bien que je me rappelle vaguement des horaires qui ont trouvé que einsum était en fait un peu plus rapide. sum est un peu plus conventionnel.

Quelqu'un a demandé une possibilité d'étendre les dimensions en einsum, mais je ne pense pas que cela se produira.