Optimize my for loop with vectors

Hi, I read that vectorizing your matlab program will speed it up so I am trying to vectorize the following for loop:

mysum = 0;
  for kx = -1:dk:1
      kyMax = sqrt(1^2-kx^2);
      for ky = 0:dk:kyMax
        myeps = kx^2 + ky^2;
        mysum = mysum + myeps;
       end
  end

I'm not sure how to do it since {kx,ky} is not a rectangular array but rather one who's ky dimension changes depending on the kx.

Answer 1

[Andrei Bobrov]

kx = -1:dk:1;
[xx,yy] = ndgrid(0:dk:1,kx);
t = xx <= sqrt(1-kx.^2);
m = [xx(:),yy(:)];
a = m(t(:),:).^2;
out = sum(a(:));

or

kx = -1:dk:1;
B = arrayfun(@(x)0:dk:x,sqrt(1-kx.^2),'un',0);
n = cellfun(@numel,B);
out = sum(repelem(kx,n).^2 + [B{:}].^2);

[olabaz]

This is what I was looking for. Thanks.

[Chad Greene]

Can you explain these steps Andrei? When I try it with any value of dk I get a dimension mismatch at the t= line.

Answer 2

[Chad Greene]

Hi olabaz,

It looks like eps is getting overwritten each time through the loop, so you'll just end up with eps = 1. As I understand it you want eps to be a 2D matrix for all the values of kx and ky?

A minor note as an aside, eps is actually a different thing in Matlab. The way you're overwriting it works fine here, but it's probably best to avoid using that variable name.

[olabaz]

I did not know about eps being used, I will change the name. Also, I left part of the code out that does a calculation with eps so I will change the code. But ideally, I would like to have two corresponding vectors kx,ky that satisfy that relationship. So I can just do eps(1) = kx(1)^2 + ky(1)^2, eps(2) = kx(2)^2 + ky(2)^2, etc. And then I could easily do eps = sum(eps).

Answer 3

[Chad Greene]

I really appreciate your interesting in vectorizing. It will take you far in life. First things first, 1^2 will always be 1, so we can remove a calculation by rewriting

kyMax = sqrt(1^2-kx^2);

as

kyMax = sqrt(1-kx^2);

Second, it can be helpful to get in the habit of rewriting this:

 for kx = -1:dk:1
    ...
 end

as

kx = -1:dk:1; 
for k = 1:length(kx) 
  ...
end

and then inside the loop, when you want to reference the kth value of kx, use kx(k). The difference may seem subtle, abstract, or even like an unnecessary extra step. But it means we can pull

kyMax = sqrt(1-kx^2);

out of the loop. Just calculate it once for all the values of kx. That gives us this:

mysum = 0;

kx = -1:dk:1; 
kyMax = sqrt(1-kx.^2);

for k = 1:length(kx) 

     for ky = 0:dk:kyMax(k)
       myeps = kx(k)^2 + ky^2;
       mysum = mysum + myeps;
     end

end

Note the .^ instead of just ^ which lets us square all the values at once. We can then do the same thing with the next for statement. That becomes this:

kx = -1:dk:1; 
kyMax = sqrt(1-kx.^2);

for k = 1:length(kx) 

     ky = 0:dk:kyMax(k); 
     myeps = kx(k)^2 + ky.^2; 

     for kk = 1:length(ky) 
       mysum = mysum + myeps(kk);
     end

end

and now there's really no need for that inner loop because we can just use sum. Taking out the inner loop it looks like this:

mysum = 0;
kx = -1:dk:1; 
kyMax = sqrt(1-kx.^2);

for k = 1:length(kx) 

     ky = 0:dk:kyMax(k); 
     myeps = kx(k)^2 + ky.^2; 
     mysum = mysum + sum(myeps); 

end

Already we've made some great progress because we no longer have a loop within a loop. Loops within loops have a way of slowing things down tremendously because everything within the inner loop needs to be done not just N times, but N*M times. Now I have to think about how to get rid of that other loop.

[olabaz]

Oh, that's interesting. I would've never thought about pulling the kyMax out of the loop.