How can I assign a value to the diagonals of a 4-D matrix using linear indexing in MATLAB?

Question

I have a 4-D matrix A of size NxNxPxQ. How can I easily change the diagonal values to 1 for each NxN 2-D submatrix in a vectorized way

Answer 1

You can actually do this very simply by directly computing the linear indices for every diagonal element, then setting them to 1:

[N,N,P,Q] = size(A);
diagIndex = cumsum([1:(N+1):N^2; N^2.*ones(P*Q-1,N)]);
A(diagIndex) = 1;

The above example finds the N diagonal indices for the first N-by-N matrix (1:(N+1):N^2). Each subsequent N-by-N matrix (P*Q-1 of them) is offset by N^2 elements from the last, so a matrix of size PQ-1-by-N containing only the value N^2 is appended to the linear indices for the diagonal of the first matrix. When a cumulative sum is performed over each column using the function CUMSUM, the resulting matrix contains the linear indices for all diagonal elements of the 4-D matrix.