Vectorize nested loops for performance

I have to solve a dynamic programming problem with a finite horizon and I am trying to vectorize as much as possible for speed. I attach here a MWE so that everybody can run it

I have run the code with the Matlab profiler and indetified two bottlenecks, marked with the comment line % THIS IS SLOW ACCORDING TO PROFILER

The first bottleneck is the call to the function ReturnFn that builds the n_a*n_a matrix RetMat. I have done this in a vectorized way (inside the function RetMat, I do not use loops)

The second bottleneck is the maximization of RetMat along the first dimension.

I'd be very grateful for any comment/suggestion!

P.S. I have an if condition inside the loops to check for bugs, I removed it now that I am confident about the code but the speed improvement is marginal.

clear,clc,close all
% STATE VARIABLES
% V(a,h,z,j)
% a: asset holdings
% h: human capital (female)
% z: labor productivity shocks (eta_m, eta_f are shocks, theta is permanent)
% j: age (from 1 to N_j)
% CHOICE VARIABLES
% d:  Female labor supply (only extensive margin: either 0 or 1)
% a': Next-period assets
% h' is implied by (d,a) and consumption is implied by the budget
% constraint
% DYNAMIC PROGRAMMING PROBLEM
% V(a,h,z,j) = max_{d,a'} F(d,a',a,h,z,j)+beta*s_j*E[V(a',h',z',j+1)|z]
% subject to
% h'=G(d,h), law of motion for human capital
verbose = 1;
%% Define grids and grid sizes
N_j = 80;
n_a = 51;
n_h = 11;
n_z = 50;
n_d = 2;
a_grid = linspace(0,450,n_a)';
h_grid = linspace(0,0.72,n_h)';
z_grid = linspace(0.9,1.1,n_z)';
d_grid = [0,1]';
z_grid = repmat(z_grid,[1,3]);
pi_z = rand(n_z,n_z);
pi_z = pi_z./sum(pi_z,2);
aprime_val = a_grid;  %(a',1)
a_val      = a_grid'; %(1,a)
%% Set parameters that do not depend on age
beta = 0.98;
r    = 0.04;
w_m  = 1;
w_f  = 0.75;
crra = 2;
nu   = 0.12;
Jr   = 45;
xi_1  = 0.05312;
xi_2  = -0.00188;
del_h = 0.074;
h_l   = 0;
p.eff_j     = ones(N_j,1);
p.pchild_j  = ones(N_j,1);
p.pen_j     = ones(N_j,1);
p.nchild_j  = ones(N_j,1);
p.s_j       = ones(N_j,1);
p.age_j     = (1:1:N_j)';
% Initialize output arrays
V  = zeros(n_a,n_h,n_z,N_j);
Policy = zeros(2,n_a,n_h,n_z,N_j);
tic
%% Solve problem in the last period
% Set age-dependent parameters
eff_j     = p.eff_j(N_j);
pchild_j  = p.pchild_j(N_j);
pen_j     = p.pen_j(N_j);
nchild_j  = p.nchild_j(N_j);
age_j     = p.age_j(N_j);
s_j       = p.s_j(N_j);
% V(a,h,z,N_j) = max_{d,a'}
V_d = zeros(n_a,n_h,n_z,n_d);
Pol_aprime_d = zeros(n_a,n_h,n_z,n_d);
for d_c = 1:n_d
    d_val = d_grid(d_c);
    for z_c = 1:n_z
        eta_m_val = z_grid(z_c,1);
        eta_f_val = z_grid(z_c,2);
        theta_val = z_grid(z_c,3);
        for h_c = 1:n_h
            h_val = h_grid(h_c);
            % RetMat is (a',a)
            RetMat = ReturnFn(d_val,aprime_val,a_val,h_val,eta_m_val,eta_f_val,theta_val,...
                w_m,eff_j,w_f,pchild_j,pen_j,r,nchild_j,crra,nu,age_j,Jr);
            [max_val,max_ind] = max(RetMat,[],1);
            V_d(:,h_c,z_c,d_c) = max_val;
            Pol_aprime_d(:,h_c,z_c,d_c) = max_ind;
        end %end h
    end %end z
end %end d
[V(:,:,:,N_j),d_max] = max(V_d,[],4);
Policy(1,:,:,:,N_j) = d_max; % Optimal d
for z_c=1:n_z
    for h_c=1:n_h
        for a_c = 1:n_a
            d_star = d_max(a_c,h_c,z_c);
            Policy(2,a_c,h_c,z_c,N_j) = Pol_aprime_d(a_c,h_c,z_c,d_star); % Optimal a'
        end
    end
end
%% Backward iteration over age
for j = N_j-1:-1:1
    if verbose==1; fprintf('Age %d out of %d \n',j,N_j); end
    V_next = V(:,:,:,j+1); %V(a',h',z')
    % Set age-dependent parameters
    eff_j     = p.eff_j(j);
    pchild_j  = p.pchild_j(j);
    pen_j     = p.pen_j(j);
    nchild_j  = p.nchild_j(j);
    age_j     = p.age_j(j);
    s_j       = p.s_j(j);
    for z_c = 1:n_z
        eta_m_val = z_grid(z_c,1);
        eta_f_val = z_grid(z_c,2);
        theta_val = z_grid(z_c,3);
        % Compute EV(a',h'), given z
        % EV = zeros(n_a,n_h);
        z_prob = pi_z(z_c,:)';
        % for zprime_c = 1:n_z
        %     EV = EV+V_next(:,:,zprime_c)*z_prob(zprime_c);
        % end %end z'
        EV = V_next.*shiftdim(z_prob,-2); %V(a',h',z')*Prob(1,1,z')
        EV = sum(EV,3); %EV(a',h')
        for d_c=1:n_d
            d_val = d_grid(d_c);
            for h_c = 1:n_h
                h_val = h_grid(h_c);
                hprime_val = f_HC_accum(d_val,h_val,age_j,xi_1,xi_2,del_h,h_l);
                % Ret_mat is (a',a)
                % THIS IS SLOW ACCORDING TO PROFILER
                Ret_mat = ReturnFn(d_val,aprime_val,a_val,h_val,eta_m_val,eta_f_val,theta_val,...
                    w_m,eff_j,w_f,pchild_j,pen_j,r,nchild_j,crra,nu,age_j,Jr);
                %[ind_l,weight_l] = interp_toolkit(hprime_val,h_grid);
                [ind_l,weight_l] = find_loc(h_grid,hprime_val);
                if ind_l>length(h_grid)-1
                    error('ind_l out of bounds')
                end
                EV_interp = EV(:,ind_l)*weight_l+EV(:,ind_l+1)*(1-weight_l);
                RHS_mat = Ret_mat+beta*s_j*EV_interp;
                % THIS IS SLOW ACCORDING TO PROFILER
                [max_val,max_ind] = max(RHS_mat,[],1);
                % max_val and max_ind are (1,a)
                V_d(:,h_c,z_c,d_c) = max_val;  %best V given d
                Pol_aprime_d(:,h_c,z_c,d_c) = max_ind; %best a' given d
            end %end h
        end %end d
    end %end z
    [V(:,:,:,j),d_max] = max(V_d,[],4);
    Policy(1,:,:,:,j) = d_max; % Optimal d
    for z_c=1:n_z
        for h_c=1:n_h
            for a_c = 1:n_a
                d_star = d_max(a_c,h_c,z_c);
                Policy(2,a_c,h_c,z_c,j) = Pol_aprime_d(a_c,h_c,z_c,d_star); % Optimal a'
            end
        end
    end
end %end j

Age 79 out of 80 
Age 78 out of 80 
Age 77 out of 80 
Age 76 out of 80 
Age 75 out of 80 
Age 74 out of 80 
Age 73 out of 80 
Age 72 out of 80 
Age 71 out of 80 
Age 70 out of 80 
Age 69 out of 80 
Age 68 out of 80 
Age 67 out of 80 
Age 66 out of 80 
Age 65 out of 80 
Age 64 out of 80 
Age 63 out of 80 
Age 62 out of 80 
Age 61 out of 80 
Age 60 out of 80 
Age 59 out of 80 
Age 58 out of 80 
Age 57 out of 80 
Age 56 out of 80 
Age 55 out of 80 
Age 54 out of 80 
Age 53 out of 80 
Age 52 out of 80 
Age 51 out of 80 
Age 50 out of 80 
Age 49 out of 80 
Age 48 out of 80 
Age 47 out of 80 
Age 46 out of 80 
Age 45 out of 80 
Age 44 out of 80 
Age 43 out of 80 
Age 42 out of 80 
Age 41 out of 80 
Age 40 out of 80 
Age 39 out of 80 
Age 38 out of 80 
Age 37 out of 80 
Age 36 out of 80 
Age 35 out of 80 
Age 34 out of 80 
Age 33 out of 80 
Age 32 out of 80 
Age 31 out of 80 
Age 30 out of 80 
Age 29 out of 80 
Age 28 out of 80 
Age 27 out of 80 
Age 26 out of 80 
Age 25 out of 80 
Age 24 out of 80 
Age 23 out of 80 
Age 22 out of 80 
Age 21 out of 80 
Age 20 out of 80 
Age 19 out of 80 
Age 18 out of 80 
Age 17 out of 80 
Age 16 out of 80 
Age 15 out of 80 
Age 14 out of 80 
Age 13 out of 80 
Age 12 out of 80 
Age 11 out of 80 
Age 10 out of 80 
Age 9 out of 80 
Age 8 out of 80 
Age 7 out of 80 
Age 6 out of 80 
Age 5 out of 80 
Age 4 out of 80 
Age 3 out of 80 
Age 2 out of 80 
Age 1 out of 80 

toc

Elapsed time is 6.267489 seconds.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function F = ReturnFn(l_f,aprime,a,h_f,eta_m,eta_f,theta,w_m,eff_j,w_f,...
    pchild_j,pen_j,r,nchild_j,crra,nu,agej,Jr)
% Calculate earnings (incl. child care costs) of men and women
y_m = w_m*eff_j*theta*eta_m;
y_f = w_f*l_f*(exp(h_f)*theta*eta_f - pchild_j);
% l_f can be either 0 or 1
% calculate available resources
cash = (1+r)*a + pen_j*(agej>=Jr) + (y_m + y_f)*(agej<Jr);
cons = cash-aprime;
%pos    = cons>0;
F = (cons/(sqrt(2+nchild_j))).^(1-crra)/(1-crra) - nu*l_f;
F(cons<=0) = -inf;
end %end function "f_ReturnFn"
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [h_f_prime] = f_HC_accum(l_f,h_f,age_j,xi_1,xi_2,del_h,h_l)
% l_f: d variable that affects h'
% h_f: current-period value of h
h_f_prime = h_f + (xi_1 + xi_2*age_j)*l_f - del_h*(1-l_f);
h_f_prime = max(h_f_prime, h_l);
%h_f_prime = h_f;
%h_f_prime = max(h_f_prime, h_l);
end %end function f_HC_accum
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [jl,omega] = find_loc(x_grid,xi)
%-------------------------------------------------------------------------%
% DESCRIPTION
% 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) so that
% omega*x_grid(jl)+(1-omega)*x_grid(jl+1)=xi
% INPUTS
% x_grid must be a strictly increasing column vector (nx,1)
% xi must be a scalar
% OUTPUTS
% jl: Left point (scalar)
% omega: weight on the left point (scalar)
% NOTES
% See find_loc_vec.m for a vectorized version.
%-------------------------------------------------------------------------%
nx = size(x_grid,1);
jl = max(min(locate(x_grid,xi),nx-1),1);
%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"
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function jl = locate(xx,x)
%function jl = locate(xx,x)
%
% x is between xx(jl) and xx(jl+1)
%
% jl = 0 and jl = n means x is out of range
%
% xx is assumed to be monotone increasing
n = length(xx);
if x<xx(1)
    jl = 0;
elseif x>xx(n)
    jl = n;
else
    jl = 1;
    ju = n;
    while (ju-jl>1)
        jm = floor((ju+jl)/2);
        if x>=xx(jm)
            jl = jm;
        else
            ju=jm;
        end
    end
end
end %end function locate

[dpb]

In

Ret_mat = ReturnFn(d_val,aprime_val,a_val,h_val,eta_m_val,eta_f_val,theta_val,...
                    w_m,eff_j,w_f,pchild_j,pen_j,r,nchild_j,crra,nu,age_j,Jr);

there may be no internal loops but you're calling the function (N_j-1)*n_d*n_h times with single values for

eta_m_val = z_grid(z_c,1);
eta_f_val = z_grid(z_c,2);
theta_val = z_grid(z_c,3);
...
d_val = d_grid(d_c);
...
h_val = h_grid(h_c);

It's complex enough I've not tried to read the function code in detail, but where you could possibly save some time is vectorizing over those loops. It may (or may not) be possible, but there's where to look.

Or, precompute and save in an ND array, maybe...

[Alessandro Maria Marco]

I can vectorize the loop over h_c but the problem is the interpolation. Now ind_l and weight_l are scalars, since hprime_val is a scalar. If I vectorize h_c, the index and the weight for interpolation become vectors of n_h elements. How do I index EV using vector indexes in

EV_interp = EV(:,ind_l)*weight_l+EV(:,ind_l+1)*(1-weight_l)

?

Can ind_l be a vector instead of a scalar in the expression above?

Note that now EV_interp has size (n_a,1), but if I vectorize h_c, it will have size (n_a,n_h)

[dpb]

It may not be possible (or may become a so much more complex structure in order to hold the result that the memory and other overhead are more costly than "dead ahead" solution) as noted earlier.

It would take far more time to try to understand the bowels of the computation than I have time to devote to be able to do more than general comments, sorry, but as the final earlier comment alluded to, perhaps what is now a 2D array can be turned into a 3D array by plane for one of the looping variables; theoretically it could be a 5D array over the other two, but then the issue of memory likely becomes the problem instead.

Or, as said, unless the values are variable every time the code is run, the alternative approach could be to precompute externally and then read the precomputed result...the time to do the precomputation would be the same but if something else is what changes from one execution to another, you can make up for it there. But, unless some of it is invariant from one run to another, then likely it's going to be "you simply have to compute what you have to compute".

[Alessandro Maria Marco]

Thanks for your suggestions. I managed to vectorized the code (see answer) but the speedup varies a lot with the grid size. If n_a=50 (which is the MWE I posted here), then the vectorized code is faster than the loops. But if n_a=201 (what I have in the real code), then the vectorized code is slower! This is surprising for me (I have more than enough RAM on my PC).

n_a = 50

loops: 2.19 seconds

vec: 1.357690 seconds

n_a = 201;

loops: 14.807141 seconds

vec: 25.714558 seconds

This is the relevant vectorized part. Is there anything that I can improve?

[ind_l,weight_l] = find_loc_vec(h_grid,hprime_val); %(h,1)

Unrecognized function or variable 'h_grid'.

%EV_interp = zeros(n_a,n_h); %(a',h)
%for h_c = 1:n_h
%EV_interp(:,h_c) = EV(:,ind_l(h_c))*weight_l(h_c)+EV(:,ind_l(h_c)+1)*(1-weight_l(h_c));
%end
EV_interp = EV(:,ind_l(:)).*weight_l'+EV(:,ind_l(:)+1).*(1-weight_l');
% entireRHS is (a',a,h) since (a',a,h) + (a',h)
entireRHS = Ret_mat+beta*s_j*permute(EV_interp,[1,3,2]);
[max_val,max_ind] = max(entireRHS,[],1);
% max_val and max_ind are (1,a)
V_d(:,:,z_c,d_c) = max_val;  %best V given d
Pol_aprime_d(:,:,z_c,d_c) = max_ind; %best a' given d

[dpb]

I don't know that I'm terribly surprised; for loop code optimization has advanced by leaps and bounds since early days of MATLAB so while vectorized code can often be quicker if can get to a linear addressing mode in internal code, if the vecorized code is complex there may be little optimization that can be done and the combination of memory and code complexity can be counterproductive to the "deadahead" solution (as you've just demonstrated).

You do need to run the process monitor and ensure you really are not page faulting with the larger sizes; it may be you have a lot of installed memory but MATLAB isn't using it all...

Answer 1

[Alessandro Maria Marco]

I managed to vectorize the code but the vectorized code is faster than the loops only if the number of points on the grid is small. Here is the code for the MWE

clear,clc,close all
% check 13.0085
% STATE VARIABLES
% V(a,h,z,j)
% a: asset holdings
% h: human capital (female)
% z: labor productivity shocks (eta_m, eta_f are shocks, theta is permanent)
% j: age (from 1 to N_j)
% CHOICE VARIABLES
% d:  Female labor supply (only extensive margin: either 0 or 1)
% a': Next-period assets
% h' is implied by (d,a) and consumption is implied by the budget
% constraint
% DYNAMIC PROGRAMMING PROBLEM
% V(a,h,z,j) = max_{d,a'} F(d,a',a,h,z,j)+beta*s_j*E[V(a',h',z',j+1)|z]
% subject to
% h'=G(d,h), law of motion for human capital
verbose = 1;
%% Define grids and grid sizes
N_j = 80;
n_a = 201;
n_h = 11;
n_z = 50;
n_d = 2;
a_grid = linspace(0,450,n_a)';
h_grid = linspace(0,0.72,n_h)';
z_grid = linspace(0.9,1.1,n_z)';
d_grid = [0,1]';
z_grid = repmat(z_grid,[1,3]);
pi_z = rand(n_z,n_z);
pi_z = pi_z./sum(pi_z,2);
aprime_val = a_grid;  %(a',1)
a_val      = a_grid'; %(1,a)
h_grid3    = shiftdim(h_grid,-2);
%% Set parameters that do not depend on age
beta = 0.98;
r    = 0.04;
w_m  = 1;
w_f  = 0.75;
crra = 2;
nu   = 0.12;
Jr   = 45;
xi_1  = 0.05312;
xi_2  = -0.00188;
del_h = 0.074;
h_l   = 0;
p.eff_j     = ones(N_j,1);
p.pchild_j  = ones(N_j,1);
p.pen_j     = ones(N_j,1);
p.nchild_j  = ones(N_j,1);
p.s_j       = ones(N_j,1);
p.age_j     = (1:1:N_j)';
% Initialize output arrays
V  = zeros(n_a,n_h,n_z,N_j);
Policy = zeros(2,n_a,n_h,n_z,N_j);
tic
%% Solve problem in the last period
% Set age-dependent parameters
eff_j     = p.eff_j(N_j);
pchild_j  = p.pchild_j(N_j);
pen_j     = p.pen_j(N_j);
nchild_j  = p.nchild_j(N_j);
age_j     = p.age_j(N_j);
s_j       = p.s_j(N_j);
% V(a,h,z,N_j) = max_{d,a'}
V_d = zeros(n_a,n_h,n_z,n_d);
Pol_aprime_d = zeros(n_a,n_h,n_z,n_d);
for d_c = 1:n_d
    d_val = d_grid(d_c);
    for z_c = 1:n_z
        eta_m_val = z_grid(z_c,1);
        eta_f_val = z_grid(z_c,2);
        theta_val = z_grid(z_c,3);
        for h_c = 1:n_h
            h_val = h_grid(h_c);
            % RetMat is (a',a)
            RetMat = ReturnFn_ale(d_val,aprime_val,a_val,h_val,eta_m_val,eta_f_val,theta_val,...
                w_m,eff_j,w_f,pchild_j,pen_j,r,nchild_j,crra,nu,age_j,Jr);
            [max_val,max_ind] = max(RetMat,[],1);
            V_d(:,h_c,z_c,d_c) = max_val;
            Pol_aprime_d(:,h_c,z_c,d_c) = max_ind;
        end %end h
    end %end z
end %end d
[V(:,:,:,N_j),d_max] = max(V_d,[],4);
Policy(1,:,:,:,N_j) = d_max; % Optimal d
for z_c=1:n_z
    for h_c=1:n_h
        for a_c = 1:n_a
            d_star = d_max(a_c,h_c,z_c);
            Policy(2,a_c,h_c,z_c,N_j) = Pol_aprime_d(a_c,h_c,z_c,d_star); % Optimal a'
        end
    end
end
%% Backward iteration over age
for j = N_j-1:-1:1
    if verbose==1; fprintf('Age %d out of %d \n',j,N_j); end
    V_next = V(:,:,:,j+1); %V(a',h',z')
    % Set age-dependent parameters
    eff_j     = p.eff_j(j);
    pchild_j  = p.pchild_j(j);
    pen_j     = p.pen_j(j);
    nchild_j  = p.nchild_j(j);
    age_j     = p.age_j(j);
    s_j       = p.s_j(j);
    for z_c = 1:n_z
        eta_m_val = z_grid(z_c,1);
        eta_f_val = z_grid(z_c,2);
        theta_val = z_grid(z_c,3);
        % Compute EV(a',h'), given z
        % EV = zeros(n_a,n_h);
        z_prob = pi_z(z_c,:)';
        % for zprime_c = 1:n_z
        %     EV = EV+V_next(:,:,zprime_c)*z_prob(zprime_c);
        % end %end z'
        EV = V_next.*shiftdim(z_prob,-2); %V(a',h',z')*Prob(1,1,z')
        EV = sum(EV,3); %EV(a',h')
        for d_c=1:n_d
            d_val = d_grid(d_c);
            % Ret_mat is (a',a)
            Ret_mat = ReturnFn(d_val,aprime_val,a_val,h_grid3,eta_m_val,eta_f_val,theta_val,...
                w_m,eff_j,w_f,pchild_j,pen_j,r,nchild_j,crra,nu,age_j,Jr);
            h_val      = h_grid;
            hprime_val = f_HC_accum(d_val,h_val,age_j,xi_1,xi_2,del_h,h_l); %(h,1)
            %[ind_l,weight_l] = interp_toolkit(hprime_val,h_grid);
            % ind_l and weight_l are arrays with size (n_h,1)
            [ind_l,weight_l] = find_loc_vec(h_grid,hprime_val); %(h,1)
            %EV_interp = zeros(n_a,n_h); %(a',h)
            %for h_c = 1:n_h
            %EV_interp(:,h_c) = EV(:,ind_l(h_c))*weight_l(h_c)+EV(:,ind_l(h_c)+1)*(1-weight_l(h_c));
            %end
            EV_interp = EV(:,ind_l(:)).*weight_l'+EV(:,ind_l(:)+1).*(1-weight_l');
            % entireRHS is (a',a,h) since (a',a,h) + (a',h)
            entireRHS = Ret_mat+beta*s_j*permute(EV_interp,[1,3,2]);
            [max_val,max_ind] = max(entireRHS,[],1);
            % max_val and max_ind are (1,a)
            V_d(:,:,z_c,d_c) = max_val;  %best V given d
            Pol_aprime_d(:,:,z_c,d_c) = max_ind; %best a' given d
        end %end d
    end %end z
    [V(:,:,:,j),d_max] = max(V_d,[],4);
    Policy(1,:,:,:,j) = d_max; % Optimal d
    for z_c=1:n_z
        for h_c=1:n_h
            for a_c = 1:n_a
                d_star = d_max(a_c,h_c,z_c);
                Policy(2,a_c,h_c,z_c,j) = Pol_aprime_d(a_c,h_c,z_c,d_star); % Optimal a'
            end
        end
    end
end %end j
toc
mean(Policy,'all')
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function F = ReturnFn(l_f,aprime,a,h_f,eta_m,eta_f,theta,w_m,eff_j,w_f,...
    pchild_j,pen_j,r,nchild_j,crra,nu,agej,Jr)
% Calculate earnings (incl. child care costs) of men and women
y_m = w_m*eff_j*theta*eta_m;
y_f = w_f*l_f*(exp(h_f)*theta*eta_f - pchild_j);
% l_f can be either 0 or 1
% calculate available resources
cash = (1+r)*a + pen_j*(agej>=Jr) + (y_m + y_f)*(agej<Jr);
cons = cash-aprime;
F = (cons/(sqrt(2+nchild_j))).^(1-crra)/(1-crra) - nu*l_f;
F(cons<=0) = -inf;
%F = -inf(size(cons)); %(a',a)
%pos = cons>0;         %(a',a)
%F(pos) = (cons(pos)/(sqrt(2+nchild_j))).^(1-crra)/(1-crra) - nu*l_f;%(a',a)
end %end function "f_ReturnFn"
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [h_f_prime] = f_HC_accum(l_f,h_f,age_j,xi_1,xi_2,del_h,h_l)
% l_f: d variable that affects h'
% h_f: current-period value of h
h_f_prime = h_f + (xi_1 + xi_2*age_j)*l_f - del_h*(1-l_f);
h_f_prime = max(h_f_prime, h_l);
%h_f_prime = h_f;
%h_f_prime = max(h_f_prime, h_l);
end %end function f_HC_accum
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [jl,omega] = find_loc_vec(x_grid,xi)
%-------------------------------------------------------------------------%
% DESCRIPTION
% 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) so that
% omega*x_grid(jl)+(1-omega)*x_grid(jl+1)=xi
% INPUTS
% x_grid must be a strictly increasing column vector (nx,1)
% xi can be a N-dim array with dim (s1,s2,..)
% OUTPUTS
% jl: Left point, same size as xi
% omega: weight on the left point, same size as xi
% NOTES
% Matlab recommneds to replace "histc" with "discretize".
%-------------------------------------------------------------------------%
nx = size(x_grid,1);
% For each 'xi', get the position of the 'x' element bounding it on the left [p x 1]
[~,jl] = histc(xi,x_grid); %#ok<HISTC>
jl(xi<=x_grid(1)) = 1;
jl(xi>=x_grid(nx)) = nx-1;
%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"