Unified model types.

Project.toml

@@ -30,6 +30,7 @@ RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 ResumableFunctions = "c5292f4c-5179-55e1-98c5-05642aab7184"
 SIMD = "fdea26ae-647d-5447-a871-4b548cad5224"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Term = "22787eb5-b846-44ae-b979-8e399b8463ab"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
 
@@ -46,3 +47,4 @@ Dolang = "≥3.3.0"
 Format = "1.3.7"
 LabelledArrays = "≥1.16.0"
 StaticArrays = "≥1.9.8"
+Term = "2.0.6"

examples/ymodels/consumption_savings_iid.jl

@@ -18,6 +18,7 @@ model = let
     σ = 0.1
     ρ = 0.0
     r = 1.02
+
     w = 1.0
     c = 0.9*w
     cbar = c
@@ -40,14 +41,14 @@ model = let
 
 
     states = CartesianSpace(;
-        w=(0.5, 20.0)
+        w=(0.5, 5.0)
     )
 
     controls = CartesianSpace(;
         c=(0.0,Inf),
     )
 
-    Σ = @SMatrix [0.0001;;]
+    Σ = @SMatrix [0.01;;]
     exogenous = MvNormal( (:y,), Σ)
 
 
@@ -76,27 +77,21 @@ end
 function arbitrage(mod::typeof(model), s::NamedTuple, x::NamedTuple, S::NamedTuple, X::NamedTuple)
     p = mod.calibration
     eq = p.β*( X.c/x.c )^(-p.γ)*(p.r) - 1 # - x.λ
-    # @warn "The euler equation is satisfied only if c<w. If c=w, it can be strictly positive."
-    # eq2 = x.λ ⟂ s.y-x.c
     (eq,)
 end
 
 function complementarities(mod::typeof(model), s::NamedTuple, x::NamedTuple, Fv::SVector)
-    eq = Fv[1] ⫫ s.w-x.c
-    # eq = Fv[1] ⟂ s.y-x.c
+    eq = Fv[1] ⫫ x.c
+    eq = -eq ⫫ s.w-x.c
+    eq = -eq
     (eq,)
 end
 
-# function initial_guess(model, m::SLArray, s::SLArray, p)
-#     # c = min( 1.0 + 0.01*(s.y - 1.0), s.y)
-#     c = 0.8*s.y
-#     λ = 0.01 # max( 1.0 + 0.01*(s.y - 1.0), 0.01*(s.y-1))
-#     return SLVector(;c, λ)
-# end
-
 function initial_guess(mod::typeof(model), s::NamedTuple)
     p = mod.calibration
-    c = s.y*0.9
+    # c = p.c
+    c = s.w*0.95
+    # c = min(s.w, 1.0 + (s.w-1.0)*0.02)
     return (;c)
 end
 

examples/ymodels/consumption_savings_iid.yaml

@@ -13,7 +13,7 @@ equations:
         w[t] = exp(y[t]) + (w[t-1]-c[t-1])*r
 
     arbitrage: |
-        β*( c[t+1]/c[t] )^(-γ)*r - 1  | 0<=c[t]<=w[t]
+        β*( c[t+1]/c[t] )^(-γ)*r - 1 | 0<=c[t]<=w[t]
 
 
 calibration:
@@ -36,7 +36,7 @@ calibration:
     rew: 0.0
 
 domain:
-    w: [0.5, 20.0]
+    w: [0.5, 5.0]
 
 exogenous:
     y: !UNormal

src/Dolo.jl

@@ -10,6 +10,7 @@ module Dolo
     using Printf
     using Format 
     using Crayons.Box
+    using Term
 
     import LinearAlgebra: cross, norm, ×
 

src/dolo_model.jl

@@ -35,12 +35,26 @@ get_controls(model::YModel) = variables(model.controls)
 # discretize(cc::CartesianSpace; n=10) = CGrid( tuple(( (cc.min[i],cc.max[i], n) for i=1:length(cc.min))...) )
 
 
+import Term: Panel, tprint, tprintln
+
 function Base.show(io::IO, m::YModel) 
-    println("Model")
-    println("* name: ", name(m))
-    println("* states: ", join(get_states(m), ", "))
-    println("* controls: ", join(get_controls(m), ", "))
-    println("* exogenous: ", m.exogenous)
+    # println(Panel("this is {red}RED{/red}"; fit=true))
+
+    exovars = [e for e in Dolo.variables(m.exogenous)]
+    exotype = typeof(m.exogenous).name.name
+    hcontrols = join(get_controls(m), ", ")
+    hstates = join([
+        (e in exovars ? "{red}$e{/red}" : e)
+        for e in get_states(m)
+    ], ", ")
+
+    txt = """
+    Model: {blue}$(name(m)){/blue}
+    * states: $hstates
+    * controls: $hcontrols
+    * exogenous: $exotype({red}$(join(exovars,",")){/red})"""
+
+    tprintln(Panel(txt; fit=true))
 end
 
 function Base.show(io::IO, m::ADModel) 

src/grids.jl

@@ -151,9 +151,9 @@ end
 @inline function getindex(s::CGrid{d}, z::Complex{Int64}) where d
     # Tf = eltype(g)
     # SVector{d,Tf}(g.grids[1][i]..., g.grids[2][j]...)
-    i,j = from_linear(s, z.re)
-    val = s[i,j]
-    Dolo.QP((i,j), val)
+    ind = from_linear(s, z.re)
+    val = s[ind...]
+    Dolo.QP(ind, val)
 end
 
 iterate(s::SGrid) = (s.points[1], 2)

src/model_extensions.jl

@@ -169,10 +169,10 @@ function transition(model::YModel{<:VAR1}, s::NamedTuple, x::NamedTuple)
 
     m = get_exo(model, s)
 
-    M = rand(model.exogenous, m)
-    S_ = transition(model, s, x, M)
+    E = rand(model.exogenous, m)
+    S_ = transition(model, s, x, E)
 
-    S = merge(M, S_)
+    S = merge(E, S_)
 
     return S