I am aware of similar questions, but couldn't find a useful one for my problem
Is there a way to speed up my code by transforming this triple loop into matrix/tensor operations?
% preallocate
X_tensor = zeros(N, N, N);
% loop
for i = 1: N
for k = 1: N
for j = 1:N
X_tensor(i,k,j) = X(i,j) * B(i,k) * Btilde(k,j) / B(i,j);
end
end
end
EDIT I am sorry, I forgot one important information: X, B and Btilde are all NxN matrices
CodePudding user response:
You can vectorize the operation by permuting the matrices and using element-wise multiplication/division.
i
, k
, and j
are dimensions [1,2,3]
. Looking at the first term in the equation, that means we want to take the [1,2,3]
order of X
(where the third dimension is a singleton) and rearrange it to [1,3,2]
. This corresponds to an index of [i,j,k]
. B
in the second term is already in the required [1,2,3]
order.
So the loops can be replaced by:
X_tensor2 = permute(X,[1,3,2]) .* B .* permute(Btilde,[3,1,2]) ./ permute(B,[1,3,2])
Here's a test program to verify correctness:
N = 5;
P = primes(400);
X = reshape(P(1:25),N,N);
B = reshape(P(26:50),N,N);
Btilde = reshape(P(51:75),N,N);
% preallocate
X_tensor = zeros(N, N, N);
X_tensor2 = zeros(N, N, N);
% loop
for i = 1: N
for k = 1: N
for j = 1:N
X_tensor(i,k,j) = X(i,j) * B(i,k) * Btilde(k,j) / B(i,j);
end
end
end
X_tensor2 = permute(X,[1,3,2]) .* B .* permute(Btilde,[3,1,2]) ./ permute(B,[1,3,2]);
isequal(X_tensor, X_tensor2)