Speed up vectorized code automatic expansion logical indexing

I have to compute the maximum of a 2-dimensional array repeatedly in a nested loop and I am trying to speed up the code, also using automatic expansion.

Basically I have an array F(l,a',a,z) and for each a and z (these are the loop variables) I find the maximum over l and a' and store it in some arrays. Consider that to compute the maximand I call a function. I already use automatic expansion in some places (see comments in the MWE) and I use logical indexing inside the function.

I attach a minimum working example so that anyone can run the problem. I tested it with Matlab online in this forum and it runs. Any help is greatly appreciated!

clear
clc
close all

na = 300;
nz = 66;
nl = 50;

% Generate fake data
l_grid = linspace(0.001,0.999,nl)' ;
a_grid = linspace(1e-6,250,na)';
z_grid = linspace(0.02,25,nz)';

% pi_z is a transition matrix
pi_z = rand(nz,nz);
pi_z = pi_z./sum(pi_z,2);

V_next = rand(na,nz);

% Initialize output of the loop
V         = zeros(na,nz);
pol_ap    = zeros(na,nz);
pol_l     = zeros(na,nz);

% Choice variables: (l,a')
l_choice      = l_grid;  %(nl,1)
aprime_choice = a_grid'; %(1,na)

%% Numerical computation starts here
tic

% Expected value EV(a') given z,theta and z' integrated out
EV = V_next*pi_z'; % V(a',z')*pi(z,z')T ==> EV(a',z)

% Loop over shock z today
for z_c = 1:nz
    EVz = EV(:,z_c)';
    % Loop over assets today
    for a_c = 1:na
        a_val   = a_grid(a_c);  % assets today, scalar
        z_val   = z_grid(z_c);  % shock today, scalar
        RHS_mat = myfun(l_choice,aprime_choice,a_val,z_val)+EVz;
        [V(a_c,z_c),opt_ind]   = max(RHS_mat(:));
        [l_opt_ind,ap_opt_ind] = ind2sub([nl,na],opt_ind);
        aprime_opt      = a_grid(ap_opt_ind);
        l_opt           = l_grid(l_opt_ind);
        pol_ap(a_c,z_c) = aprime_opt;
        pol_l(a_c,z_c)  = l_opt;
    end % end loop over assets a
end % end loop over z

toc

Elapsed time is 6.147432 seconds.

function F = myfun(l,aprime,a_val,z_val)

aux     = 3.5;
cons    = a_val+l*z_val - aprime; % Use automatic expansion, dim: (nl,na)
indx    = cons<=0; % NOT feasible! 
F       = log(cons)-l.^aux/aux; % Use automatic expansion
F(indx) = -inf; % Logical indexing


end

[dpb]

Have you profiled the code to determine where the time is being used?

Peephole optimization on the basis of an assumption of what is the culprit is, more often than not, futile...

I've not figured out just what this is supposed to be doing in a brief overview, but that makes me think it could probably all be vectorized directly by using MATLAB builtin propensity to operate by column -- IOW, it doesn't look to me like a point-by-point nested loop is needed at all...

[Alessandro Maria Marco]

I have profiled the code, the bottleneck seems to be inside the call to myfun, especially the following lines:

F       = log(cons)-l.^aux/aux; % Use automatic expansion
F(indx) = -inf; % Logical indexing

[Alessandro Maria Marco]

In other words, the slow part is the computation of RHS_mat, whereas the max and the call of ind2sub (I need this to convert the linear index to matrix subscript) is fast

[dpb]

Well, the slow part is almost certainly the calculation of the log() and the exponential; they're intrinsic math lib calls...you can confirm that by rearranging the code to see them individually. There was a recent case where rewriting x.^3 as x.*x.*x achieved something like a 10-20X speedup. aux not being integer-valued here, I'd suppose the additional sqrt() would be a killer trying to go that route.

I'd try

F=-inf(size(cons));
ix=cons>0;                            % feasible
F(ix)=log(cons(ix))-l.^aux/aux;

and/or variations on a theme of the above to avoid creating the complex variables your present does by not removing the bad values first.

[Alessandro Maria Marco]

@dpb Thanks for your suggestion, I tried it but it doesn't work since the array sizes are not compatible. Specifically, the following line of your code generates a mistake:

F(ix)=log(cons(ix))-l.^aux/aux;

Before applying the logical mask, cons has dimension (nl,na) and l has dimension (nl,1). The addition is OK because Matlab automatically expand the l term to make it compatible with the first. However, if you do cons(ix) you are changing the number of elements, so that the sizes become not compatible.

I think a better route would be to eliminate the nested loops over a_c and z_c, but I found it challenging.

[dpb]

OK, overlooked that in the indexing expression, yes.

Use

F(ix)=log(cons(ix));
F=F-l.^aux/aux;

instead.

I've some other commitments have to get ready for; if get a chance I'll see if I can find enough time to figure out what it is this is actually doing.

First guess that might help would be to look at meshgrid examples if what your doing is trying to evaluate a function over a 2D grid as it looks like might be...for debugging/testing, I'd also suggest cutting your array sizes down to something that fits on one screen in the command line window -- it's MUCH easier to figure out what's going on when the sizes are like 5x3 or somesuch instead of "cast of thousands" and the logic is no different with size.

Answer 1

[dpb]

function F = myfun(l,aprime,a_val,z_val)
  aux     = 3.5;
  cons    = a_val+l*z_val - aprime;
  F=-inf(size(cons));
  indx=cons>0;        % feasible! 
  F(indx)=log(cons(indx));
  F=F-l.^aux/aux;
end

Using the above modfication of your function and a slightly smaller array size just so it wouldn't take too long I got the following results on local machine --

>> amm      % original
Elapsed time is 2.460195 seconds.
>> amm      % above modification
Elapsed time is 1.465042 seconds.
>>

That's about a 40% reduction which is not quite a factor of 2X, but getting close....

I didn't put the original back to test, but when I bumped the array sizes back to match yours, it was

>> amm
Elapsed time is 3.748369 seconds.
>> 

Using the 60% factor, that would scale to about 6.3 sec for the original that is in ballpark of the above...

I did look at the code some more; it doesn't appear to be particularly amenable to further vectorizing.

I would suspect although I did not investigate further that the bigger speed up could be gained by screening the input vectors for the ranges that can't be feasible first and only iterating over those instead of the full, dead-ahead, just try 'em all approach.

[Alessandro Maria Marco]

Thanks a lot for your answer! I tested your code on my PC with the original grid size and I got a good speed-up: from 3.73 seconds to 2.03 seconds. The reason is that now we take the log only of positive numbers, so we avoid doing unnecessary computations that will be discarded anyways.

I reproduce the experiment below

clear
clc
close all

na = 300;
nz = 66;
nl = 50;

% Generate fake data
l_grid = linspace(0.001,0.999,nl)' ;
a_grid = linspace(1e-6,250,na)';
z_grid = linspace(0.02,25,nz)';


% pi_z is a transition matrix
pi_z = rand(nz,nz);
pi_z = pi_z./sum(pi_z,2);

V_next = rand(na,nz);

% Initialize output of the loop
V         = zeros(na,nz);
pol_ap    = zeros(na,nz);
pol_l     = zeros(na,nz);

% Choice variables: (l,a')
l      = l_grid;  %(nl,1)
aprime = a_grid'; %(1,na)

%% Numerical computation starts here
disp('Method 1')

Method 1

tic

% Expected value EV(a') given z,theta and z' integrated out
EV = V_next*pi_z'; % V(a',z')*pi(z,z')T ==> EV(a',z)

% Loop over shock z today
for z_c = 1:nz
    EVz = EV(:,z_c)';
    % Loop over assets today
    for a_c = 1:na
        a_val   = a_grid(a_c);  % assets today, scalar
        z_val   = z_grid(z_c);  % shock today, scalar
        RHS_mat = myfun(l,aprime,a_val,z_val)+EVz;
        [V(a_c,z_c),opt_ind]   = max(RHS_mat(:));
        [l_opt_ind,ap_opt_ind] = ind2sub([nl,na],opt_ind);
        aprime_opt      = a_grid(ap_opt_ind);
        l_opt           = l_grid(l_opt_ind);
        pol_ap(a_c,z_c) = aprime_opt;
        pol_l(a_c,z_c)  = l_opt;
    end % end loop over assets a
end % end loop over z

t1=toc

t1 = 4.1172

%% Method 2
disp('Method 2, courtesy of dpb on Matlab answers')

Method 2, courtesy of dpb on Matlab answers

% Initialize output of the loop
V2         = zeros(na,nz);
pol_ap2    = zeros(na,nz);
pol_l2     = zeros(na,nz);

tic

% Expected value EV(a') given z,theta and z' integrated out
EV = V_next*pi_z'; % V(a',z')*pi(z,z')T ==> EV(a',z)

% Loop over shock z today
for z_c = 1:nz
    EVz = EV(:,z_c)';
    % Loop over assets today
    for a_c = 1:na
        a_val   = a_grid(a_c);  % assets today, scalar
        z_val   = z_grid(z_c);  % shock today, scalar
        RHS_mat = myfun2(l,aprime,a_val,z_val)+EVz;
        [V2(a_c,z_c),opt_ind]   = max(RHS_mat(:));
        [l_opt_ind,ap_opt_ind] = ind2sub([nl,na],opt_ind);
        aprime_opt      = a_grid(ap_opt_ind);
        l_opt           = l_grid(l_opt_ind);
        pol_ap2(a_c,z_c) = aprime_opt;
        pol_l2(a_c,z_c)  = l_opt;
    end % end loop over assets a
end % end loop over z

t2=toc

t2 = 2.7294

err_l = max(abs(pol_l(:)-pol_l2(:)));
err_ap = max(abs(pol_ap(:)-pol_ap2(:)));
err_V = max(abs(V(:)-V2(:)));
disp([err_l,err_ap,err_V])

     0     0     0

fprintf('Speed up gain: %f \n',t2/t1)

Speed up gain: 0.662923 

function F = myfun(l,aprime,a_val,z_val)

aux     = 3.5;
cons    = a_val+l.*z_val - aprime; % Use automatic expansion, dim: (nl,na)
indx    = cons<=0; % NOT feasibility 
F       = log(cons)-l.^aux/aux; % Use automatic expansion
F(indx) = -inf;

end

function F = myfun2(l,aprime,a_val,z_val)
  aux     = 3.5;
  cons    = a_val+l*z_val - aprime;
  F       = -inf(size(cons));   % (nl,na)
  indx    = cons>0; % feasible! % (nl,na)
  F(indx) = log(cons(indx));    % (nl,na)
  F       = F-l.^aux/aux;  % Array (nl,na) + array (nl,1)==> array (nl,na)
end