(M x N x N) 행렬의 대각선을 설정하는 빠른 방법? Einsum/n 차원 fill_diagonal?

Question

필자는 행렬을 기반으로 한 빠르고 최적화 된 코드를 작성하려고 노력 중이며 최근에는 einsum이 속도 향상을위한 도구로 사용되었습니다.

다차원 배열의 대각선을 효율적으로 설정하는 데 사용할 수 있습니까? 아니면 데이터 만 반환 할 수 있습니까?

내 문제는 각 사각형 (N x N) 행렬의 열을 합산하여 사각형 행렬 (모양 : M x N x N)의 배열에 대한 대각선을 설정하려고합니다.

나의 현재 (느린 루프 기반) 솔루션입니다 :

# 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 

이 전혀주십시오 개선 할 수 있습니까? np.fill_diagonal이 비 정사 매트릭스를 수용하지 않는 것 같습니다 (그러므로 내 루프 기반 솔루션을 강요합니다). 아마 어둠이 여기에서도 도움이 될 수 있을까요?

Answer 1

한 가지 방법은 2D으로 다시 형성하고, 대각선 값을 사용하여 ncols+1 단계로 열을 설정하는 것입니다. 모양을 변경하면보기가 만들어 지므로 대각선 위치에 직접 액세스 할 수 있습니다. 따라서, 구현 될 것이다 - 당신이 np.source(np.fill_diagonal)을 할 경우

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

Answer 2

당신은 그 차원 케이스에 의해 당신의 3D 케이스에

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

@Divakar's 솔루션이 적용되는 '스트라이드'접근 방식을 사용하여 볼 수 있습니다 2 차원에서 '평면화'.

열을 M.sum(axis=1)과 합할 수 있습니다. 나는 약간 타이밍을 회상했지만 실제로는 einsum이 조금 더 빠르다는 것을 발견했다. sum은 좀 더 일반적입니다.

누군가가 einsum의 크기를 확장 할 수 있는지 묻는 질문이 있지만 그럴 것이라고 생각하지 않습니다.