diff --git a/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_DiffusionTwoAssets.jl b/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_DiffusionTwoAssets.jl
index 8791794..4926b2a 100644
--- a/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_DiffusionTwoAssets.jl
+++ b/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_DiffusionTwoAssets.jl
@@ -73,7 +73,7 @@ yend = OrderedDict(:v => [(m.ρ / m.γ + (1 - 1 / m.γ) * m.r)^(-m.γ) * (a + y
 y, residual_norm = pdesolve(m, stategrid, yend)
 # 
 # # Important: check marginal value of wealth converges to 1.0
-# # This happens ONLY if a >= 1000.0. Otherwise with 300 it does not work. This is interesting. Maybe it means there should be a better way to have bordering condition at top
+# # This happens ONLY if a >= 1000.0. Otherwise with 300 it does not work. 
 # b = ((m.r + (m.ρ - m.r)/m.γ - (1-m.γ) / (2 * m.γ) * (m.μR - m.r)^2 / (m.γ * m.σR^2)))^(1/(1 - 1/m.γ))
 # pw = (result[:v] * (1-m.γ)).^(1/(1-m.γ)-1) .* result[:va] ./ b
 
diff --git a/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_TwoStates.jl b/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_TwoStates.jl
index 3885445..9ae8f4a 100644
--- a/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_TwoStates.jl
+++ b/examples/ConsumptionProblem/AchdouHanLasryLionsMoll_TwoStates.jl
@@ -17,7 +17,7 @@ mutable struct AchdouHanLasryLionsMoll_TwoStatesModel
     amax::Float64 
 end
 
-function AchdouHanLasryLionsMoll_TwoStatesModel(;yl = 0.5, yh = 1.5, λlh = 0.2, λhl = 0.2, r = 0.03, ρ = 0.04, γ = 2.0, amin = - yl / r, amax = 500.0)
+function AchdouHanLasryLionsMoll_TwoStatesModel(;yl = 0.5, yh = 1.5, λlh = 0.2, λhl = 0.2, r = 0.03, ρ = 0.04, γ = 2.0, amin = - yl / r, amax = 50.0)
     AchdouHanLasryLionsMoll_TwoStatesModel(yl, yh, λlh, λhl, r, ρ, γ, amin, amax)
 end
 
@@ -25,53 +25,53 @@ function (m::AchdouHanLasryLionsMoll_TwoStatesModel)(state::NamedTuple, value::N
     (; yl, yh, λlh, λhl, r, ρ, γ, amin, amax) = m    
     (; a) = state
     (; vl, vla_up, vla_down, vh, vha_up, vha_down) = value
+
+    # low income state
     vla = vla_up
     iter = 0
-    va = vla_up
-    @label start1
-    va = max(va, 1e-8)    
-    c = va^(-1 / γ)
-    μa = yl + r * a - c
-    if (iter == 0) & (μa <= 0)
+    vla = vla_up
+    @label startl
+    vla = max(vla, eps())    
+    cl = vla^(-1 / γ)
+    μla = yl + r * a - cl
+    if (iter == 0) & (μla <= 0)
         iter += 1
-        va = vla_down
-        @goto start1
+        vla = vla_down
+        @goto startl
     end
-    if (a ≈ amin) && (μa <= 0.0)
-        c = yl + r * amin
-        μa = 0.0
-        va = c^(-γ)
+    if (a ≈ amin) && (μla <= 0.0)
+        cl = yl + r * amin
+        μla = 0.0
+        vla = cl^(-γ)
     end
-    vah = va
-    μal = μa
-    vlt = - (c^(1 - γ) / (1 - γ) + μa * va + λlh * (vh - vl) - ρ * vl)
+    vlt = - (cl^(1 - γ) / (1 - γ) + μla * vla + λlh * (vh - vl) - ρ * vl)
    
-
+    # high income state
+    vha = vha_up
     iter = 0
-    va = vha_up
-    @label start2
-    va = max(va, 1e-8)    
-    c = va^(-1 / γ)
-    μa = yh + r * a - c
-    if (iter == 0) & (μa <= 0)
+    vha = vha_up
+    @label starth
+    vha = max(vha, eps())    
+    ch = vha^(-1 / γ)
+    μha = yh + r * a - ch
+    if (iter == 0) & (μha <= 0)
         iter += 1
-        va = vha_down
-        @goto start2
+        vha = vha_down
+        @goto starth
     end
-    if (a ≈ amin) && (μa <= 0.0)
-        c = yh + r * amin
-        μa = 0.0
-        va = c^(-γ)
+    if (a ≈ amin) && (μha <= 0.0)
+        ch = yh + r * amin
+        μha = 0.0
+        vha = ch^(-γ)
     end
-    val = va
-    μah = μa
-    vht = - (c^(1 - γ) / (1 - γ) + μa * va + λhl * (vl - vh) - ρ * vh)
-    return (; vlt, vht), (; vlt, vht, vah, val, μah, μal)
+    vht = - (ch^(1 - γ) / (1 - γ) + μha * vha + λhl * (vl - vh) - ρ * vh)
+    
+    return (; vlt, vht), (; vha, vla, μha, μla)
 end
 
 m = AchdouHanLasryLionsMoll_TwoStatesModel()
 m.amin += 0.001
-stategrid = OrderedDict(:a => m.amin .+ range(0, (m.amax - m.amin)^0.8, length = 5000).^(1/0.8))
+stategrid = OrderedDict(:a => m.amin .+ range(0, (m.amax - m.amin)^(1/2), length = 200).^2)
 yend = OrderedDict(:vl => (m.ρ ./ m.γ .+ (1 .- 1 / m.γ) .* m.r)^(-m.γ) .* (stategrid[:a] .+ m.yl ./ m.r).^(1-m.γ) ./ (1 - m.γ), :vh => (m.ρ ./ m.γ .+ (1 .- m.γ) .* m.r)^(-m.γ)  .* (stategrid[:a] .+ m.yh ./ m.r).^(1-m.γ) ./ (1 - m.γ))
 result = pdesolve(m, stategrid, yend)
 @assert result.residual_norm <= 1e-5
diff --git a/examples/ConsumptionProblem/WangWangYang.jl b/examples/ConsumptionProblem/WangWangYang.jl
index dc410e1..b885483 100644
--- a/examples/ConsumptionProblem/WangWangYang.jl
+++ b/examples/ConsumptionProblem/WangWangYang.jl
@@ -45,9 +45,6 @@ result = pdesolve(m, stategrid, yend, bc = OrderedDict(:pw => (1.0, 1.0)))
 @assert result.residual_norm <= 1e-5
 
 
-
-
-
 # Alternative solution bypassing pdesolve
 # just encode the PDE has a vector equation
 using InfinitesimalGenerators
diff --git a/examples/ConsumptionProblem/try.jl b/examples/ConsumptionProblem/try.jl
deleted file mode 100644
index 0cf80c5..0000000
--- a/examples/ConsumptionProblem/try.jl
+++ /dev/null
@@ -1,59 +0,0 @@
-using EconPDEs
-
-struct SimpleModel
-    # income process parameters
-    y::Float64
-    r::Float64
-
-    # utility parameters
-    ρ::Float64  
-    γ::Float64
-
-    amin::Float64
-    amax::Float64 
-end
-
-function SimpleModel(;y = 1.0, r = 0.03, ρ = 0.05, γ = 2.0, amin = 0.0, amax = 500.0)
-    SimpleModel(y, r, ρ, γ, amin, amax)
-end
-
-function (m::SimpleModel)(state::NamedTuple, value::NamedTuple)
-    (; y, r, ρ, γ, amin, amax) = m    
-    (; a) = state
-    (; v, va_up, va_down, vaa) = value
-    va = va_up
-    iter = 0
-    @label start
-    va_up = max(va_up, eps())    
-    c = va^(-1 / γ)
-    μa = y + r * a - c
-    if (iter == 0) & (μa <= 0)
-        iter += 1
-        va = va_down
-        @goto start
-    end
-    if (a ≈ amin) && (μa <= 0.0)
-        va = (y + r * amin)^(-γ)
-        c = y + r * amin
-        μa = 0.0
-    end
-    vt = - (c^(1 - γ) / (1 - γ) + μa * va - ρ * v)
-    return (; vt), (; c)
-end
-
-m = SimpleModel(amin = 100.0)
-stategrid = OrderedDict(:a =>  range(m.amin, m.amax, length = 100))
-yend = OrderedDict(:v => [log(m.y + max(a, 0.0)) for a in stategrid[:a]])
-result = pdesolve(m, stategrid, yend)
-@assert result.residual_norm <= 1e-5
-
-# finite horizon over 20 years
-yend = OrderedDict(:v => [max(a + y)^(1-m.γ)/(1-m.γ) for y in stategrid[:y], a in stategrid[:a]]) 
-τs = range(0, stop = 100, step = 1)
-result  = pdesolve(m, stategrid, yend, τs)
-@assert maximum(result.residual_norm) <= 1e-5
-
-
-# Check marginal value of wealth converges to 1.0 at infinity
-#b = ((m.r + (m.ρ - m.r)/m.γ))^(1/(1 - 1/m.γ))
-#pw = (result[:v] * (1-m.γ)).^(1/(1-m.γ)-1) .* result[:va] ./ b
diff --git a/test/runtests.jl b/test/runtests.jl
index 32c5c46..f789ef1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -4,7 +4,7 @@ for x in (:CampbellCochrane, :Wachter, :BansalYaron, :GarleanuPanageas, :DiTella
 	@testset "$x" begin include("../examples/AssetPricing/$(x).jl") end
 end
 
-for x in (:AchdouHanLasryLionsMoll_Diffusion, :AchdouHanLasryLionsMoll_DiffusionTwoAssets)
+for x in (:AchdouHanLasryLionsMoll_Diffusion, :AchdouHanLasryLionsMoll_DiffusionTwoAssets, :AchdouHanLasryLionsMoll_TwoStates, :WangWangYang)
 	@testset "$x" begin include("../examples/ConsumptionProblem/$(x).jl") end
 end