I rewrote (translated) Ben Moll’s Matlab code for the Huggett model (partial equilibrium) in continuous time in Julia. When I change the value of the shocks, the concavity condition during the iteration fails. The Matlab code converges with changing the shock values, but the Julia code does not. Any suggestions as to why this is happening and how to correct it?

When I change the shock values from z = [0.1, 0.2] to z = [0.1, 0.3], it results in complex values and the iteration fails. The Julia code is below and the original Matlab code is in this link. I did not code the KF part since the problem is with the iteration.

```
using LinearAlgebra, SparseArrays
s = 1.2;
r = 0.035;
rho = 0.05;
# shock values
z = [0.1, 0.2];
nz = size(z)[1];
# co-state initial values
la = [1.5, 1.1];
# grid
na = 500;
amin = -0.02;
amax = 3.0;
a = collect(range(amin,amax,na));
da = (amax-amin)/(na-1);
# making the grid for all the states
aa = repeat(a,1,nz);
zz = repeat(z',na,1);
# iteration parameters
maxit = 100;
crit = 10^(-6);
# HJB parameters
Delta = 1000;
global dVf = zeros(na,nz);
global dVb = zeros(na,nz);
global c = zeros(na,nz);
# sparse matrix for the co-state variables
global Aswitch = [-sparse(I,na,na).*la[1] sparse(I,na,na).*la[1];
sparse(I,na,na).*la[2] -sparse(I,na,na).*la[2]];
# initial guess for the value functions
v0 = zeros(na,nz);
for i in 1:nz
v0[:,i] = (z[i].+r.*a).^(1-s)./((1-s)*rho);
end
# start the loop
global v = v0;
for n in 1:maxit
V = v;
# FD
dVf[1:na-1,:] = (V[2:na,:]-V[1:na-1,:])/da;
dVf[na,:] = (z.+r.*amax).^(-s);
# BD
dVb[2:na,:] = (V[2:na,:]-V[1:na-1,:])/da;
dVb[1,:] = (z.+r.*amin).^(-s);
# concavity check
global I_concave = dVb .> dVf;
# consumption and savings with FD
cf = dVf.^(-1/s);
global sf = zz.+r.*aa.-cf;
# consumption and savings with BD
cb = dVb.^(-1/s);
global sb = zz.+r.*aa.-cb;
# consumption and savings at steady state
c0 = zz.+r.*aa;
dV0 = c0.^(-s);
# upwind scheme
If = sf .> 0;
Ib = sb .< 0;
I0 = (1).-If.-Ib;
Ib[na,:] .= 1.0;
If[na,:] .= 0.0;
# approximate derivate based on the drift,
# the last term is important for stability
dV_upwind = dVf.*If.+dVb.*Ib.+dV0.*I0;
global c = dV_upwind.^(-1/s);
global u = c.^(1-s)/(1-s);
# construct the A matrix using sparse routines
# first construct the pieces
global lower = -min.(sb,0)/da;
global middle = (min.(sb,0).-max.(sf,0))/da;
global upper = max.(sf,0)/da;
global A_list = [];
for j in 1:nz
push!(A_list,spdiagm(-1 => lower[2:na,j], 0 => middle[:,j], 1 => upper[1:na-1,j]))
end
global A = [A_list[1] spzeros(na,na);spzeros(na,na) A_list[2]] .+ Aswitch;
if findmax(abs.(sum(A,dims=2)))[1] > 10^(-12)
println("Improper Transition Matrix")
break
end
#set up the now linear system and solve it!
global B = (1/Delta+rho).*sparse(I,nz*na,nz*na).-A;
# stack the shock states on top of each other (automated)
global u_stacked = u[:,1];
global V_stacked = V[:,1];
for j in 2:nz
u_stacked = cat(u_stacked, u[:,j], dims=1);
V_stacked = cat(V_stacked, V[:,j], dims=1);
end
global b = u_stacked.+V_stacked./Delta;
global V_sol = B\b;
global V = [V_sol[1:na] V_sol[na+1:nz*na]];
global V_change = V .- v
global v = V;
dist = findmax(abs.(V_change))[1]
if dist<crit
println("Value function converged!")
println("Number of iterations: ",n,".")
println("Error = ", dist,".")
break
end
end
```