How can I speed up this function? #vectorize

I'm struggling to speed up the code inside the version1 function. The only way to vectorize it that I can seem to figure out is version2, but that actually makes the function slower. Help would be much appreciated :-)

n1 = 400;
n2 = 3;
n3 = 7;
n4 = 2;

A = rand(n1, n2, n3);
B = randi(n2, [n4, n1, n3]);
disp(timeit(@() version1(A, B, C, n1, n2, n3, n4)));
disp(timeit(@() version2(A, B, C, n1, n2, n3, n4)));

function C = version1(A, B, n1, n2, n3, n4)
    C = zeros(n4, n1, n3);
    for ii = 1:n3
        for jj = 1:n1
            for kk = 1:n4
                C(kk, jj, ii) = A( ...
                    jj, ...
                    B(kk, jj, ii), ...
                    ii ...
                );
            end
        end
    end
end

function C = version2(A, B, n1, n2, n3, n4)
    C = zeros(n4, n1, n3);
    for ii = 1:n3
        for jj = 1:n1
            C(:, jj, ii) = A( ...
                jj, ...
                B(:, jj, ii), ...
                ii ...
            );
        end
    end
end

Answer 1

[Matt J]

I doubt a vectorized solution would be faster than the loop, but here is one way to vectorize it.

[K,J,I]=ndgrid(1:n4,1:n1,1:n3);  %Recycle this, if possible

Bvals=B(sub2ind(size(B),K,J,I));

C=A(  sub2ind(size(A), J,Bvals,I ));

[Matt J]

For the particular dimensions n1,n2,n3,n4 in your example, there does seem to be a benefit

test()

   4.1479e-04

   1.2279e-04

function test

n1 = 400;
n2 = 3;
n3 = 7;
n4 = 2;


A = rand(n1, n2, n3);
B = randi(n2, [n4, n1, n3]);

disp(timeit(@() version1(A, B, n1, n2, n3, n4)));
disp(timeit(@() version3(A, B, n1, n2, n3, n4)));

end

function C = version1(A, B, n1, n2, n3, n4)
C = zeros(n4, n1, n3);
    for ii = 1:n3
        for jj = 1:n1
            for kk = 1:n4
                C(kk, jj, ii) = A( ...
                    jj, ...
                    B(kk, jj, ii), ...
                    ii ...
                );
            end
        end
    end
end



function C = version3(A, B, n1, n2, n3, n4)


    [K,J,I]=ndgrid(1:n4,1:n1,1:n3);  %Recycle this, if possible
    
    Bvals=B(sub2ind(size(B),K,J,I));
    
    C=A(  sub2ind(size(A), J,Bvals,I ));

end

[Fredrik P]

Thanks! However, I think that there is an error. The reshape line ought to be

C = permute(C, [3 2 1]);

[Matt J]

No, I would do this instead:

function C = version3(A, B, n1, n2, n3, n4)
    [K,J,I]=ndgrid(1:n4,1:n1,1:n3);  %Recycle this, if possible
    
    Bvals=B(sub2ind(size(B),K,J,I));
    
    C=A(  sub2ind(size(A), J,Bvals,I ));
end

It doesn't matter though. Surely you've found by now that the for-loop is fastest.

[Fredrik P]

For me, your latest code is a lot faster. Perhaps it comes down to my actual data not being random. But more likely, the code I initially posted was a bit unfair. There, I preallocated the C matrix, but the preallocation really ought be done inside function version1 and version2, as your version3 does not need any preallocation. Also, there is not really just one A matrix and one C matrix, but rather three of each: A1, A2, and A3 and C1, C2, and C3, which all can use a variable indices=sub2ind(size(A), J,Bvals,I ). Perhaps I was too eager to minimize my minimal working example :-/

Answer 2

[Hassaan]

It seems that the complexity of the indexing in this specific case is not easily amenable to vectorization without encountering shape mismatches or broadcasting issues.it might be more practical to focus on optimizing the original loop-based implementations (version1 and version2). Often, the simplest solution can be the most efficient, especially if the code is already quite optimized and the overhead of additional complexity does not lead to significant performance gains.

If performance is critical, and these functions are a bottleneck in your application, you might consider other strategies, such as:

1. Profiling: Use MATLAB's built-in profiling tools to identify the exact bottlenecks in your code.
2. Parallel Computing: If you have access to MATLAB's Parallel Computing Toolbox, you might gain performance by distributing some of these operations across multiple cores or GPUs.
3. Compiled Languages: For the most intensive computational tasks, rewriting the critical parts in a compiled language like C++ and integrating it with MATLAB might be beneficial.

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

If you find the solution helpful and it resolves your issue, it would be greatly appreciated if you could accept the answer. Also, leaving an upvote and a comment are also wonderful ways to provide feedback.

Professional Interests

• Technical Services and Consulting
• Embedded Systems | Firmware Developement | Simulations
• Electrical and Electronics Engineering

Feel free to contact me.