vectorize nested loops interpolation

I have two nested loops and inside the inner loop I use linear interpolation in one dimension. I would like to vectorize the innermost loop and leave only the outer loop for speed reasons. I attach a minimum working example below, see the comment "How can I vectorize the inner loop over i"? Just to be clear the code works fine, the problem is that for my purposes it is too slow.

Thanks!

% This is a minimum working example

a_min = 1e-10;
a_max = 1;
na    = 200;
a_grid = linspace(a_min,a_max,na)';

Nsim = 10000; % Desired value is larger, around 50000

r = 0.04;
nz = 14; % no. of points 2nd dim of pol_ap
TT = 80; % no. of points 3rd dim of pol_ap
pol_ap    = rand(na,nz,TT);    % fake data
z_ind_sim = randi(nz,Nsim,TT); % fake data
income_sim = rand(Nsim,TT);    % fake data

%% Simulate Nsim agents for TT periods
a_sim  = zeros(Nsim,TT+1);
c_sim  = zeros(Nsim,TT);

% Set initial condition for assets
a_sim(:,1) = a_grid(1); % All agents start their life with a_grid(1)

tic
for j = 1:TT % outer loop over j
    z_c          = z_ind_sim(:,j); % dim: (Nsim,1)
    % =====> How can I vectorize the inner loop over i?  <=====
    for i=1:Nsim
        a_sim(i,j+1) = interp1q(a_grid,pol_ap(:,z_c(i),j), a_sim(i,j));
    end
    % Here it is easy to vectorize the loop over i
    c_sim(:,j) = (1+r)*a_sim(:,j)+income_sim(:,j)-a_sim(:,j+1);
    
end %END j
toc

Elapsed time is 3.441594 seconds.

a_sim(:,TT+1) = []; % we want T periods, not T+1

Answer 1

[Jan]

% Your code:
a_min = 1e-10;
a_max = 1;
na    = 200;
a_grid = linspace(a_min,a_max,na)';
Nsim = 10000; % Desired value is larger, around 50000
r = 0.04;
nz = 14; % no. of points 2nd dim of pol_ap
TT = 80; % no. of points 3rd dim of pol_ap
pol_ap    = rand(na,nz,TT);    % fake data
z_ind_sim = randi(nz,Nsim,TT); % fake data
income_sim = rand(Nsim,TT);    % fake data
%% Simulate Nsim agents for TT periods
a_sim  = zeros(Nsim,TT+1);
c_sim  = zeros(Nsim,TT);
% Set initial condition for assets
a_sim(:,1) = a_grid(1); % All agents start their life with a_grid(1)

tic
for j = 1:TT % outer loop over j
    z_c = z_ind_sim(:,j); % dim: (Nsim,1)
    for i=1:Nsim
        a_sim(i,j+1) = interp1q(a_grid,pol_ap(:,z_c(i),j), a_sim(i,j));
    end
    % Here it is easy to vectorize the loop over i
    c_sim(:,j) = (1+r)*a_sim(:,j)+income_sim(:,j)-a_sim(:,j+1);
    
end %END j
toc

Elapsed time is 2.610400 seconds.

% Interpolate once only to get the indices:
a_sim2  = zeros(Nsim,TT+1);
c_sim2  = zeros(Nsim,TT);
% Set initial condition for assets
a_sim2(:,1) = a_grid(1); % All agents start their life with a_grid(1)

tic
for j = 1:TT % outer loop over j
    z_c = z_ind_sim(:, j); % dim: (Nsim,1)

    index = interp1(a_grid, 1:size(pol_ap, 1), a_sim(:,j));
    F = floor(index);
    S = index - F;
    for i = 1:Nsim
        a_sim2(i, j+1) = pol_ap(F(i),     z_c(i), j) * (1 - S(i)) + ...
                         pol_ap(F(i) + 1, z_c(i), j) * S(i);
    end
    
    % Here it is easy to vectorize the loop over i
    c_sim2(:,j) = (1+r)*a_sim2(:,j)+income_sim(:,j)-a_sim(:,j+1);
    
end %END j
toc

Elapsed time is 0.063530 seconds.

max(abs(a_sim(:) - a_sim2(:)))

ans = 3.4139e-14

max(abs(c_sim(:) - c_sim2(:)))

ans = 3.5305e-14

[Alessandro D]

That's great, thanks a lot! So the key is to pre-compute the indexes, not to eliminate the inner loop entirely. Building on your logic, I created two small functions that pre-compute the indexes and the weights needed for linear interpolation. find_loc_vec1 is basically what you wrote. It returns as outputs the left points and the weights on the left points.

function [jl,omega] = find_loc_vec1(x_grid,xi)

%Find jl s.t. x_grid(jl)<=xi<x_grid(jl+1)
%for jl=1,..,N-1
% omega is the weight on x_grid(jl)
% Similar to find_loc_vec but we use interp1 instead of histc

nx = size(x_grid,1);

index = interp1(x_grid, (1:nx)', xi);
jl = floor(index);  % array of left points
S     = index - jl; % array of weights on jl+1
omega = 1-S;        % array of weight on jl

end %end function "find_loc_vec1"

[Alessandro D]

The second function instead of using interp1 uses histc and is pasted below. I tested them for speed and they are similar. For small data it seems that the version with histc is a bit faster. It is probably better not to use histc however, since Matlab does not recommend it and interp1 performs similarily (at least for this problem).

function [jl,omega] = find_loc_vec(x_grid,xi)

%Find jl s.t. x_grid(jl)<=xi<x_grid(jl+1)
%for jl=1,..,N-1
% omega is the weight on x_grid(jl)

nx = size(x_grid,1);

% For each 'xi', get the left point
[~,jl] = histc(xi,x_grid);
jl = max(jl,1);     % To avoid index=0 when xi < x(1)
jl = min(jl,nx-1);  % To avoid index=nx+1 when xi > x(end).


%Weight on x_grid(j)
omega = (x_grid(jl+1)-xi)./(x_grid(jl+1)-x_grid(jl));
omega = max(min(omega,1),0);

end %end function "find_loc_vec"

[Bruno Luong]

The last time I check the fatest method to find bin location is discretize