Home > Blockchain >  Triple loop into tensor matlab
Triple loop into tensor matlab

Time:09-13

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)
  • Related