How to vectorize this code?

Could you let me know how to vectorize the below for-loop ?

pigrid   = zeros(knum,epnum,sepnum,znum,sznum);

    for ki = 1:knum
        for epi = 1:epnum
           for sepi = 1:sepnum
               for zi = 1:znum
                   for szi = 1:sznum

           pigrid(ki,epi,sepi,zi,szi) = epgrid(epi,sepi)*zgrid(zi,szi)*kgrid(ki)^alpha-f;  

                    end
                end
            end
        end
    end

[Adam Danz]

Please provide values of knum,epnum,sepnum,znum,sznum. If the current nested loop strategy works, why do you need to vectorize it? Assuming epgrid, zgrid, and kgrid and matricies, this should be really fast.

[Stephen23]

matwork's "Answer" moved here:

Thanks for your response.

knum=100,epnum=10,sepnum=5,znum=10,sznum=5

Yes. The code works and doesn't take much time. Just I want to see if vectorization helps save time.

[Adam Danz]

Here's what I got so far. It's a simplification but still relies on 3 of your 5 loops.

for ki = 1:knum
    for epi = 1:epnum
        for sepi = 1:sepnum
            pigrid(ki,epi,sepi,:,:) = epgrid(epi,sepi)*zgrid(1:znum,1:sznum)*kgrid(ki)^alpha-f;
        end
    end
end

if zgrid(1:znum,1:sznum) == zgrid, then a further simplification would be

for ki = 1:knum
    for epi = 1:epnum
        for sepi = 1:sepnum
            pigrid(ki,epi,sepi,:,:) = epgrid(epi,sepi)*zgrid*kgrid(ki)^alpha-f;
        end
    end
end

Further vectorization would require a few repmat() and reshape() function calls to store the results in your current 5D variable and I'm not sure you're going to save time by doing that. Nevertheless, I'd love to see a complete vectorization of this from someone who has more time to fiddle with it.

[icdi]

Thanks! Your code works and is simpler than mine. This helps me.

[OCDER]

Moved from Answers to here.

"Thanks for your response. znum is not 10 but 7. Sorry about my typo."

In that case, my other answer wouldn't have worked... hm... The only way I can see this loop go faster is to move the

kgrid(ki)^alpha

outside the nested loops to prevent recalculating this value many times.

[OCDER]

By the way, just realized your ID name is "matwork" with a mathworks logo. Perhaps try changing your ID and logo so that it does not seem to impersonate MathWorks employees.

https://www.mathworks.com/matlabcentral/termsofuse.html Read Content 2.ii

[icdi]

Oh I see. Thanks for your comment.

[OCDER]

No problem. And make sure to accept @Adam's answer, once he copies it over to the Answer section :)

[icdi]

Yes yes. Thanks again for your advice!

Answer 1

[Adam Danz]

Creating a fully vectorized version of these nested loops seems either not possible or not practical. Here's a summary of some improvements listed in the comment section under the question.

This simplification still relies on 3 of your 5 loops.

for ki = 1:knum
    for epi = 1:epnum
        for sepi = 1:sepnum
            pigrid(ki,epi,sepi,:,:) = epgrid(epi,sepi)*zgrid(1:znum,1:sznum)*kgrid(ki)^alpha-f;
        end
    end
end

if zgrid(1:znum,1:sznum) == zgrid, then a further simplification would be

for ki = 1:knum
    for epi = 1:epnum
        for sepi = 1:sepnum
            pigrid(ki,epi,sepi,:,:) = epgrid(epi,sepi)*zgrid*kgrid(ki)^alpha-f;
        end
    end
end

@OCDER also recommended moving kgrid(ki)^alpha outside of the loop which would eliminate ~10000 exponential calculations.

Assuming kgrid(1:knum)==kgrid,

kgalpha = kgrid.^alpha;              % or kgrid(1:knum).^alpha
for ki = 1:knum
    for epi = 1:epnum
        for sepi = 1:sepnum
            pigrid(ki,epi,sepi,:,:) = epgrid(epi,sepi)*zgrid*kgalpha(ki)-f;
        end
    end
end

Lastly, you could remove the " -f" to after the loop like this:

kgalpha = kgrid.^alpha;         %or kgrid(1:knum).^alpha
for ki = 1:knum
    for epi = 1:epnum
        for sepi = 1:sepnum
            pigrid(ki,epi,sepi,:,:) = epgrid(epi,sepi)*zgrid*kgalpha(ki);
        end
    end
end
pigrid = pigrid - f;

[Adam Danz]

All of these versions have been tested with fake data where epgrid, zgrid, and kgrid were matrices with random numbers and the loop variables used the values you shared above. Additionally, you should test your current nested loop code with the simplified version you choose to ensure you've got the same values. To do that, you could use

isequal(pigrid, pigrid2).

[icdi]

Oh Thanks! I learned a lot from your comments. Also thanks for letting me know the way to compare the versions.