How to rewrite this high dimensional matrix calculation

Here is the code:

X = randn(A,B,C);
Z = zeros(A,B,A,B,C);
for a=1:A
    for b=1:B
        Z(a,b,:,:,:) = X - X(a,b,:);
        Z(a,b,a,b,:) = X(a,b,:);
    end
end

X is a given three dimensional matrix with dimension A*B*C, i want to obtain the 5 dimension matrix Z. I am not familiar with the high dimension matrix manipulation so i just write with for loop, but it is really time-consuming. Is there any way to accelerate the calculation of Z?

Answer 1

[Bruno Luong]

Fully vectorized

Y = reshape(X, [], 1, C);
Z = reshape(Y, 1, [], C)-Y;
Z = reshape(Z, [], C);
Z(1:A*B+1:end,:) = Y;
Z = reshape(Z,[A B A B C]);

[Hancheng Zhu]

Wow, this code is right and the running time is much faster than the original code. I have a question, i actually want to obtain S, which is calculated by Z and their relationship is in this code

S = squeeze(prod(Z,[3,4]))

I haven't really understood your code, but i think after you have obtained

Z(1:A*B+1:end,:) = Y;

Maybe you don't have to reshape Z (your last line), you can also obtain S. If it is possible, could you please write this code to me? I want to short the running time of my code. Thanks for your help, bro.

[Bruno Luong]

Reshape is not at all CPU consuming, it is just create another array container on the same underlined data in memory. So don't worry about the last reshape statement.

But summing in last dimensions is less efficient than first dimensions for the reason that linear memory access and caching by CPU. You might want to reorganize the dimensions of your X and Z arrays with that in mind. It also depends if summing is a bootleneck or not in your case. You have to profile the code and see.

[Hancheng Zhu]

Thanks for your answer, Dr. Luong. I only have last question. I have a very similar time-consuming code as follows

X = randn(A,B,C);
Z = zeros(A,A,B,C);
for a=1:A
    Z(a,:,:,:) = X - X(a,:,:);
    Z(a,a,:,:) = X(a,:,:);
end
S = squeeze(prod(Z,2));

Could you help me write in a fully vectorized way? Thanks a lot.

[Bruno Luong]

Y = reshape(X, A, 1, []);
Z = reshape(X, 1, A, [])-Y;
Z = reshape(Z, [], B*C);
Z(1:A+1:end,:) = Y;
Z = reshape(Z,[A A B C]);
S = reshape(prod(Z,2), [A B C]);

I recommend to keep your original code; at least as comment while you don't fully master the vectorized version, which is not obviously clear.

[Hancheng Zhu]

This really helps me a lot, thanks for your answer.

Answer 2

[Zinea]

Hi @Hancheng Zhu,

To accelerate the calculation of the 5-dimensional matrix Z, you can make the following changes:

1. Preallocate the matrix Z with zeros.

Z = zeros(A, B, A, B, C);

2. Use reshape to manipulate the dimensions of X. Also, element-wise operations is used to compute the difference for each element in X without explicit loops.

X_expanded = reshape(X, [1, 1, A, B, C]);
Z = X_expanded - reshape(X, [A, B, 1, 1, C]);

You may refer to the following documention on vectorization for more information:

https://www.mathworks.com/help/matlab/matlab_prog/vectorization.html

3. A loop is still used to assign values where the indices of Z match (a,b) in the 3rd and 4th dimensions, but this is a relatively small operation compared to the entire matrix computation.

for a = 1:A
    for b = 1:B
        Z(a, b, a, b, :) = X(a, b, :);
    end
end