WIP

Project.toml

@@ -3,6 +3,7 @@ uuid = "9d24351c-2990-5e1b-a277-04c4b809c898"
 version = "0.5.0-alpha"
 
 [deps]
+AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -14,6 +15,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Dolang = "e5c7262c-e9d2-5620-ad8e-1af14eb8a8e3"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Formatting = "59287772-0a20-5a39-b81b-1366585eb4c0"
@@ -42,3 +44,5 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StringDistances = "88034a9c-02f8-509d-84a9-84ec65e18404"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"

src/dolo_model.jl

@@ -15,6 +15,13 @@ name(::YModel{C,A,B,D,N}) where C where A where B where D where N = N
 
 bounds(model::YModel, s) = model.controls
 
+
+# TODO: check somewhere that the type of
+#  states/controls/exogenous
+# is the same as calibration
+
+eltype(model::YModel) = eltype(model.calibration)
+
 function recalibrate(model::YModel; kwargs...)
     calib = merge(model.calibration, kwargs)
     YModel(model.states, model.controls, model.exogenous, calib, model.source)

src/grids.jl

@@ -6,8 +6,8 @@ import Base: eltype, iterate, size
 eltype(cg::AGrid{d}) where d = SVector{d, Float64}
 ndims(cg::AGrid{d}) where d = d
 
-struct CGrid{d} <: AGrid{d}
-    ranges::NTuple{d, Tuple{Float64, Float64, Int64}}
+struct CGrid{d,Tf} <: AGrid{d}
+    ranges::NTuple{d, Tuple{Tf, Tf, Int64}}
 end
 
 size(cg::CGrid{d}) where d = tuple((e[3] for e in cg.ranges)... )
@@ -39,8 +39,8 @@ getindex(g::CGrid{1}, ::Colon) = [SVector(i) for i in range(g.ranges[1]...)]
 
 
 
-struct SGrid{n,d} <: ASGrid{d}
-    points::SVector{n,SVector{d, Float64}}
+struct SGrid{n,d,T} <: ASGrid{d}
+    points::SVector{n,SVector{d, T}}
 end
 
 function SGrid(Q::Matrix)
@@ -134,7 +134,7 @@ end
 iterate(s::SGrid) = (s.points[1], 2)
 iterate(s::SGrid, state) = state<=length(s) ? (s.points[state], state+1) : nothing
 
-eltype(cg::CGrid{d}) where d = SVector{d, Float64}
+eltype(cg::CGrid{d,Tf}) where d where Tf = SVector{d, Tf}
 
 function iterate(cg::CGrid{d}) where d
 

src/processes.jl

@@ -302,9 +302,9 @@ end
 
 discretize(var::VAR1, d::Dict) = length(d)>=1 ? discretize(var, d[:n]) : discretize(var)
 
-struct MarkovChain{names, d, d2, k}
-    P::SMatrix{d,d,Float64,d2}
-    Q::SVector{d, SVector{k, Float64}}
+struct MarkovChain{names, d, d2, k, Tf}
+    P::SMatrix{d,d,Tf,d2}
+    Q::SVector{d, SVector{k, Tf}}
 end
 
 variables(mc::MarkovChain{names}) where names = names
@@ -313,14 +313,15 @@ ndims(mc::MarkovChain{names}) where names = length(names)
 # MarkovChain(names, P, Q) = MarkovChain{names, typeof(P), typeof(Q)}(P,Q)
 
 function MarkovChain(names, P::Matrix, Q::Matrix) 
+    Tf = eltype(P)
     d = size(P,1)
-    sm = SMatrix{d,d,Float64,d*d}(P)
+    sm = SMatrix{d,d,Tf,d*d}(P)
     sv = SVector( tuple( ( SVector(Q[i,:]...) for i in 1:size(Q,1))...)  )
-    MarkovChain{names,d,d^2,length(sv[1])}(sm,sv)
+    MarkovChain{names,d,d^2,length(sv[1]),Tf}(sm,sv)
 end
 
-MarkovChain(names, P::SMatrix, Q::SVector{d,SVector{k,Float64}}) where d where k = MarkovChain{names, size(P,1), length(P), length(Q[1])}(P, Q) # TODO: specify type arguments
-function MarkovChain(P::SMatrix, Q::SVector{d,SVector{k,Float64}}) where d where k
+MarkovChain(names, P::SMatrix, Q::SVector{d,SVector{k,Tf}}) where d where k where Tf = MarkovChain{names, size(P,1), length(P), length(Q[1]),Tf}(P, Q) # TODO: specify type arguments
+function MarkovChain(P::SMatrix, Q::SVector{d,SVector{k,Tf}}) where d where k where Tf
     names = tuple((Symbol(string("e", i)) for i=1:k)...)
     MarkovChain(names, P, Q)
 end

src/space.jl

@@ -1,31 +1,31 @@
 abstract type Space{d} end
 
-struct CartesianSpace{d,dims}
+struct CartesianSpace{d,dims,Tf}
     # names::NTuple{d, Symbol}
-    min::NTuple{d, Float64}
-    max::NTuple{d, Float64}
+    min::NTuple{d, Tf}
+    max::NTuple{d, Tf}
 end
 
 const CSpace = CartesianSpace
 
-CartesianSpace(a::Tuple{Float64}, b::Tuple{Float64}) = CartesianSpace{length(a), (:x,)}(a,b)
-CartesianSpace(a::Tuple{Float64, Float64}, b::Tuple{Float64, Float64}) = CartesianSpace{length(a), (:x_1, :x_2)}(a,b)
+CartesianSpace(a::Tuple{Tf}, b::Tuple{Tf}) where Tf = CartesianSpace{length(a), (:x,), Tf}(a,b)
+CartesianSpace(a::Tuple{Tf, Tf}, b::Tuple{Tf, Tf}) where Tf = CartesianSpace{length(a), (:x_1, :x_2), Tf}(a,b)
 
-function CartesianSpace(kwargs::Pair{Symbol, Tuple{Float64, Float64}}...)
+function CartesianSpace(kwargs::Pair{Symbol, Tuple{Tf, Tf}}...) where Tf
 
     names = tuple( (e[1] for e in kwargs)... )
     a = tuple( (v[2][1] for v in values(kwargs))... )
     b = tuple( (v[2][2] for v in values(kwargs))... )
     d = length(names)
 
-    return CartesianSpace{d, names}(a,b)
+    return CartesianSpace{d, names, Tf}(a,b)
 
 end
 
 CartesianSpace(;kwargs...) = CartesianSpace(kwargs...)
 
 
-getindex(cs::CartesianSpace{d}, ind::SVector{d, Float64}) where d = ind
+getindex(cs::CartesianSpace{d}, ind::SVector{d, Tf}) where d where Tf= ind
 
 import Base: in
 
@@ -46,15 +46,15 @@ ndims(dom::CartesianSpace{d, dims}) where d where dims = d
 variables(dom::CartesianSpace{d,t}) where d where t = t
 dims(dom::CartesianSpace) = variables(dom)
 
-struct GridSpace{N,d,dims}
-    points::SVector{N,SVector{d,Float64}}
+struct GridSpace{N,d,dims,Tf}
+    points::SVector{N,SVector{d,Tf}}
 end
 
 const GSpace = GridSpace
 
-GridSpace(v::SVector{N, SVector{d, Float64}}) where d where N = GridSpace{length(v), d, (:i_,)}(SVector(v...))
-GridSpace(v::Vector{SVector{d, Float64}}) where d = GridSpace{length(v), d, (:i_,)}(SVector(v...))
-GridSpace(names, v::SVector{k, SVector{d, Float64}}) where k where d = GridSpace{length(v), d, names}(v)
+GridSpace(v::SVector{N, SVector{d, Tf}}) where d where N where Tf = GridSpace{length(v), d, (:i_,), Tf}(SVector(v...))
+GridSpace(v::Vector{SVector{d, Tf}}) where d where Tf = GridSpace{length(v), d, (:i_,), Tf}(SVector(v...))
+GridSpace(names, v::SVector{k, SVector{d, Tf}}) where k where d where Tf = GridSpace{length(v), d, names, Tf}(v)
 
 getindex(gs::GridSpace, i::Int64) = gs.points[i]
 
@@ -111,8 +111,8 @@ function dropnames(namedtuple::NamedTuple, names::Tuple{Vararg{Symbol}})
 end
 
 
-QP(space::CartesianSpace{d}; values...) where d = let
-    s_ =zero(SVector{d,Float64})*NaN
+QP(space::CartesianSpace{d, Tf}; values...) where d where Tf = let
+    s_ =zero(SVector{d,Tf})*NaN
     s0 = QP(s_,s_)
     QP(space, s0; values...)
 end
@@ -136,10 +136,10 @@ QP(space::GridSpace{d}; i_=1) where d = QP(i_,space[i_])
 
 
 
-function QP(space::ProductSpace{<:GridSpace{d1},<:CartesianSpace{d2}}; values...) where d1 where d2
+function QP(space::ProductSpace{<:GridSpace{d1},<:CartesianSpace{d2, Tf}}; values...) where d1 where d2 where Tf
     i0 = get(values, :i_, 1)
     m_ = space.spaces[1][i0]
-    s_ =zero(SVector{d2,Float64})*NaN
+    s_ =zero(SVector{d2,Tf})*NaN
     s0 = QP((i0,s_),SVector(m_...,s_...))
     QP(space, s0; values...)
 end