From 1f1bf1d4bd5173d56c1effaa26109e9482b38297 Mon Sep 17 00:00:00 2001 From: Jason Jensen Date: Wed, 17 Jul 2024 17:18:56 -0400 Subject: [PATCH 1/2] MTK adjustments --- .github/workflows/main.yml | 2 +- Project.toml | 11 +- ext/StateSpaceEconMTKExt.jl | 268 ++++++++++++++++++++++++++++++++++ src/MTKExt.jl | 235 +++++++++++++++++++++++++++++ src/StateSpaceEcon.jl | 1 + src/stackedtime/solverdata.jl | 14 +- test/mtkext.jl | 99 +++++++++++++ test/runtests.jl | 2 + 8 files changed, 623 insertions(+), 9 deletions(-) create mode 100644 ext/StateSpaceEconMTKExt.jl create mode 100644 src/MTKExt.jl create mode 100644 test/mtkext.jl diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index a50e3f2..8f8094b 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -20,7 +20,7 @@ jobs: strategy: matrix: - julia-version: ["1.7", "1.8", "1.9"] + julia-version: ["1.10"] julia-arch: [x64] os: [ubuntu-latest, windows-latest, macOS-latest] diff --git a/Project.toml b/Project.toml index a33c76b..d8b2c2d 100644 --- a/Project.toml +++ b/Project.toml @@ -18,6 +18,12 @@ Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" TimeSeriesEcon = "8b6756d2-c55c-11ea-2998-5f67ea17da60" TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" +[weakdeps] +ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" + +[extensions] +StateSpaceEconMTKExt = "ModelingToolkit" + [compat] DelimitedFiles = "1.7" DiffResults = "1.0" @@ -36,8 +42,11 @@ julia = "1.7" [extras] DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" +ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" +NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" +SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Test", "ModelingToolkit", "SymbolicIndexingInterface", "NonlinearSolve"] diff --git a/ext/StateSpaceEconMTKExt.jl b/ext/StateSpaceEconMTKExt.jl new file mode 100644 index 0000000..8a536a3 --- /dev/null +++ b/ext/StateSpaceEconMTKExt.jl @@ -0,0 +1,268 @@ +module StateSpaceEconMTKExt + +using ModelBaseEcon, StateSpaceEcon, ModelingToolkit +using ModelingToolkit.SciMLBase: solve, AbstractSolution +using StateSpaceEcon.StackedTimeSolver: StackedTimeSolverData +using StateSpaceEcon.SteadyStateSolver: SolverData as SteadyStateSolverData +using TimeSeriesEcon: rangeof, MVTSeries +import StateSpaceEcon.MTKExt: + stacked_time_system, + compute_residuals_stacked_time, + _create_system, + rename_variables, + get_var_names, + solve_steady_state!, + steady_state_system, + compute_residuals_steady_state, + sol2data + +# Note: Unfortunately, Documenter.jl seems not to include docstrings from package extensions. +# (At least it's not obvious how to do so.) +# As a result, the docstrings for the following methods live in `../src/MTKExt.jl`. + +############################## +# Stacked time system +############################## + +function stacked_time_system(m::Model, exog_data::AbstractArray{Float64,2}; fctype = fcgiven, plan::Plan=Plan(m, 1+m.maxlag:size(exog_data, 1)-m.maxlead)) + + data = zeroarray(m, plan) + data[plan.exogenous .== true] = exog_data[plan.exogenous .== true] + data[1:m.maxlag, :] = exog_data[1:m.maxlag, :] # Initial conditions + if fctype === fcgiven + data[end-m.maxlead+1:end, :] = exog_data[end-m.maxlead+1:end, :] # Final conditions + end + + # TODO: Should other variants be accepted, or is `:default` sufficient? + sd = StackedTimeSolverData(m, plan, fctype, :default) + + return stacked_time_system(sd, data) + +end + +function stacked_time_system(sd::StackedTimeSolverData, data::AbstractArray{Float64,2}) + + f = let sd = sd, data = data + (u, p) -> compute_residuals_stacked_time(u, sd, data) + end + + # Since we are just creating a MTK system out of the `NonlinearProblem`, + # the `u0` we specify here is not actually used for solving the system. + u0 = zeros(count(sd.solve_mask)) + prob = NonlinearProblem(f, u0) + s = _create_system(prob, sd) + s.defaults[:data_in] = data # hack to pass along the data to be used when constructing an MVTSeries + + return s + +end + +function compute_residuals_stacked_time(u, sd::StackedTimeSolverData, data::AbstractArray{Float64,2}) + + # We need to use `zeros` instead of `similar` because `assign_exog_data!` + # doesn't assign final conditions for the shocks. + point = zeros(eltype(u), size(data)) + point[sd.solve_mask] = u + # Note: `assign_exog_data!` assigns both exogenous data (including initial conditions) *and* final conditions. + StateSpaceEcon.StackedTimeSolver.assign_exog_data!(point, data, sd) + + # Emulate `StackedTimeSolver.stackedtime_R!`, computing the residual + # for each equation at each simulation period. + resid = map(sd.TT[1:length(sd.II)]) do tt + # `tt` grabs all the data from `point` for the current simulation period, lags, and leads (e.g., `tt = [1, 2, 3, 4, 5, 6]`). + # The last value of `tt` contains the indices for just the leads (e.g., `tt = [100, 101, 102]`) + # (but we don't use it in this loop). + p = view(point, tt, :) + med = sd.evaldata + map(med.alleqns, med.allinds) do eqn, inds + eqn.eval_resid(p[inds]) + end + end + + return vcat(resid...) + +end + +function _create_system(prob::NonlinearProblem, sd::StackedTimeSolverData) + + # Convert `prob` into a ModelingToolkit system. + @named old_sys = modelingtoolkitize(prob) + + # `modelingtoolkitize` creates a system with automatically generated variable names + # `x₁`, `x₂`, etc. + # Rename the variables to make it easier to query the solution to problems + # created from this system. + sys = rename_variables(old_sys, sd) + + # We need `conservative = true` to prevent errors in `structural_simplify`. + s = structural_simplify(complete(sys); conservative = true) + + return s + +end + +function rename_variables(old_sys::NonlinearSystem, sd::StackedTimeSolverData) + + # reshape the solve_mask + solve_mask = reshape(sd.solve_mask, :, maximum(last, sd.evaldata.var_to_idx)) + + # Find all column indices of columns that have at least one `true`. + var_cols = vec(mapslices(any, solve_mask; dims = 1)) |> findall + + # Convert from column indices to variable/shock names. + var_names = map(var_cols) do i + findfirst(p -> last(p) == i, sd.evaldata.var_to_idx) + end + + var_lengths = Dict{Symbol,Int64}() + col_count = 1 + for col in var_cols + var_lengths[var_names[col_count]] = count(solve_mask[:, col]) + col_count += 1 + end + + # For each model variable, create a ModelingToolkit variable vector + # with the correct number of time steps. + # `vars` will end up being, e.g., `[pinf[1:97], rate[1:97], ygap[1:97]]`. + vars = map(var_names) do var_name + NE = var_lengths[var_name] + ModelingToolkit.@variables $(var_name)[1:NE] + end |> x -> reduce(vcat, x) + + # Collect all the individual variables to enable substituting the variable names + # that were automatically generated by `modelingtoolkitize`. + # `vars_enumerated` will end up being, e.g., `[pinf[1], pinf[2], ..., ygap[97]]`. + vars_enumerated = reduce(vcat, map(collect, vars)) + old_vars = unknowns(old_sys) + subs = Dict(old_vars .=> vars_enumerated) + + # Take all the equations from `old_sys`, rename the variables, + # and create a new system with the new variable names. + eqs = substitute.(ModelingToolkit.equations(old_sys), Ref(subs)) + @named sys = NonlinearSystem(eqs, vars_enumerated, []) + + return sys + +end + +function get_var_names(sd) + + # Reshape the solve mask to have a column for each variable/shock. + solve_mask = reshape(sd.solve_mask, :, maximum(last, sd.evaldata.var_to_idx)) + # Find all column indices of columns that have at least one `true`. + var_cols = vec(mapslices(any, solve_mask; dims = 1)) |> findall + # Convert from column indices to variable/shock names. + var_names = map(var_cols) do i + findfirst(p -> last(p) == i, sd.evaldata.var_to_idx) + end + + return var_names + +end + +############################## +# Steady state system +############################## + +function solve_steady_state!(m::Model, sys::NonlinearSystem; u0 = zeros(length(unknowns(sys))), solver = nothing, solve_kwargs...) + + # Creating a `NonlinearProblem` from a `NonlinearFunction` seems to be faster + # than creating a `NonlinearProblem` directly from a `NonlinearSystem`. + nf = NonlinearFunction(sys) + prob = NonlinearProblem(nf, u0) + sol = isnothing(solver) ? solve(prob; solve_kwargs...) : solve(prob, solver; solve_kwargs...) + + # Copy the steady state solution to the model. + m.sstate.values[.!m.sstate.mask] = sol.u + # Mark the steady state as solved. + m.sstate.mask .= true + + return m + +end + +steady_state_system(m::Model) = steady_state_system(SteadyStateSolverData(m)) + +function steady_state_system(sd::SteadyStateSolverData) + + f = let sd = sd + (u, p) -> compute_residuals_steady_state(u, sd) + end + + # Since we are just creating a MTK system out of the `NonlinearProblem`, + # the `u0` we specify here is not actually used for solving the system. + u0 = zeros(count(sd.solve_var)) + prob = NonlinearProblem(f, u0) + @named sys = modelingtoolkitize(prob) + s = complete(sys) + + return s + +end + +# `ModelBaseEcon.__to_dyn_pt` doesn't work with automatic differentiation, +# so copy it here to allow passing in an apropriately typed `buffer`. +# TODO: Should this method be added directly to ModelBaseEcon.jl? +function __to_dyn_pt!(buffer, pt, s) + # This function applies the transformation from steady + # state equation unknowns to dynamic equation unknowns + for (i, jt) in enumerate(s.JT) + if length(jt.ssinds) == 1 + pti = pt[jt.ssinds[1]] + else + pti = pt[jt.ssinds[1]] + ModelBaseEcon.__lag(jt, s) * pt[jt.ssinds[2]] + end + buffer[i] += pti + end + return buffer +end + +function compute_residuals_steady_state(u, sd) + + # `point` needs to contain the levels and slopes for all variables and shocks, + # but these values are all zero except for those included in `u`. + point = zeros(eltype(u), length(sd.solve_var)) + point[sd.solve_var] = u + + # Emulate `SteadyStateSolver.global_SS_R!`, computing the residual for each equation. + resid = map(sd.alleqns) do eqn + if hasproperty(eqn.eval_resid, :s) + # `eqn.eval_resid` closes over `s::SSEqData`. + # There's no function to create just `s`, so just grab it from the closure. + s = eqn.eval_resid.s + buffer = zeros(eltype(u), length(s.JT)) + s.eqn.eval_resid(__to_dyn_pt!(buffer, point[eqn.vinds], s)) + else + # `eqn.eval_resid` is an `EquationEvaluator` (not a closure) + # that computes the residual for an equation defined by a steady state constraint. + eqn.eval_resid(point[eqn.vinds]) + end + end + + return resid + +end + +function sol2data(sol::AbstractSolution, s::ModelingToolkit.AbstractSystem, m::Model, p::Plan) + rng = rangeof(p) + + # get endogenous points + _result = ModelingToolkit.SymbolicIndexingInterface.getu(sol, + reduce(vcat, map(collect, map(var -> hasproperty(s, var.name) ? getproperty(s, var.name) : [], m.varshks))) + )(sol) # is this a fair assumption, that we will get back the variables? + + # get mask for endogenous points + mask = p.exogenous .== false + mask[begin:begin+m.maxlag-1, :] .= false # initial conditions + mask[end-m.maxlead+1:end, :] .= false # final conditions + + # construct results from exogenous data + result = MVTSeries(rng, m.varshks, s.defaults[:data_in]) + + # update endogenous points + result[mask] = _result + + return result +end + +end diff --git a/src/MTKExt.jl b/src/MTKExt.jl new file mode 100644 index 0000000..b856b5a --- /dev/null +++ b/src/MTKExt.jl @@ -0,0 +1,235 @@ +""" + MTKExt + +A module that is part of StateSpaceEcon.jl. +Declares functions for creating ModelingToolkit systems out of `Model`s. +By default, these functions have no methods; +the relevant methods will be added when ModelingToolkit.jl is loaded. +""" +module MTKExt + +const MTK_NEEDED_STRING = """ +!!! note + This method requires ModelingToolkit.jl to be loaded + and only works with Julia 1.9.0 or later. +""" + +""" + stacked_time_system(m::Model, exog_data::Matrix; fctype = fcgiven) + +Convert a `Model` into a `ModelingToolkit.NonlinearSystem` +that incorporates the stacked time algorithm. + +$MTK_NEEDED_STRING + +# Inputs +- `m::Model`: Model to convert. +- `exog_data::AbstractArray{Float64,2}`: Data matrix of size `(NT, nvars + nshks)`, + where `NT` is the number of simulation periods to simulate (plus lags and leads), + and `nvars` and `nshks` are the number of variables and shocks in the model. + This data is used to specify the exogenous data, initial conditions, and + (if applicable) the final conditions. + +# Options +- `fctype = fcgiven`: The class of final conditions to use in the simulation. + The default is [`fcgiven`](@ref StateSpaceEcon.StackedTimeSolver.fcgiven). + +!!! note + If `fctype` is [`fclevel`](@ref StateSpaceEcon.StackedTimeSolver.fclevel) + or [`fcslope`](@ref StateSpaceEcon.StackedTimeSolver.fcslope), + `m` will need its steady state solved prior to calling this function. + See [`sssolve!`](@ref StateSpaceEcon.SteadyStateSolver.sssolve!) or [`solve_steady_state!`](@ref). + +# Example +This function is used to bring a `Model` into +the ModelingToolkit/SciML ecosystem. +Here is an example of calling this function +and then converting the returned system into a `ModelingToolkit.NonlinearProblem` +to be solved with one of the solvers from [NonlinearSolve.jl](https://github.com/SciML/NonlinearSolve.jl): +```julia +using ModelBaseEcon, StateSpaceEcon, ModelingToolkit, NonlinearSolve +@using_example E3 +m = E3.newmodel() +exog_data = rand(102, 6) # Replace with desired data. +s = create_system(m, exog_data) +nf = NonlinearFunction(s) +u0 = zeros(length(unknowns(s))) +prob = NonlinearProblem(nf, u0) +solver = NewtonRaphson() # Replace with desired solver. +sol = solve(prob, solver) +``` + +See the [ModelingToolkit.jl docs](https://docs.sciml.ai/ModelingToolkit/stable/tutorials/ode_modeling/) +to see what one can do with the solution object `sol`. + +--- + + stacked_time_system(sd::StackedTimeSolverData, data::Matrix) + +Alternative call signature when `sd` is already available (e.g., created manually). +""" +function stacked_time_system end +export stacked_time_system + +""" + compute_residuals_stacked_time(u, sd::StackedTimeSolverData, data::Matrix) + +Compute the residuals of the stacked time system in `sd` +given variable values `u` and exogenous data `data`. +This function is closed over in [`stacked_time_system`](@ref) +to create a function that can be passed to `ModelingToolkit.NonlinearProblem`. + +$MTK_NEEDED_STRING + +!!! warning + Internal function not part of the public interface. +""" +function compute_residuals_stacked_time end + +""" + _create_system(prob::ModelingToolkit.NonlinearProblem, sd::StackedTimeSolverData) + +Create a `ModelingToolkit.NonlinearSystem` from the given problem. +The solver data `sd` should be the same as used for creating `prob`. + +$MTK_NEEDED_STRING + +!!! warning + Internal function not part of the public interface. +""" +function _create_system end + +""" + rename_variables(old_sys::ModelingToolkit.NonlinearSystem, sd::StackedTimeSolverData) + +Create a new `ModelingToolkit.NonlinearSystem` by replacing the variable names in `old_sys` +with variable names from the solver data `sd`. +The solver data `sd` should be the same as used for creating `old_sys`. + +$MTK_NEEDED_STRING + +!!! warning + Internal function not part of the public interface. +""" +function rename_variables end + +""" + get_var_names(sd::StackedTimeSolverData) + +Return the names of endogenous variables and/or shocks that are solved for. +A variable/shock name is returned if the corresponding variable/shock +is included in `sd.solve_mask` for at least one simulation period. + +$MTK_NEEDED_STRING + +!!! warning + Internal function not part of the public interface. +""" +function get_var_names end + +""" + solve_steady_state!(m::Model, sys::ModelingToolkit.NonlinearSystem; u0 = zeros(...), solver = nothing, solve_kwargs...) + +Solve the steady state system `sys` (created with [`steady_state_system`](@ref)) +and store the result in `m`. +The model `m` should be the same model used to create `sys`. + +!!! note + This function is a replacement for [`sssolve!`](@ref StateSpaceEcon.SteadyStateSolver.sssolve!) + and uses the ModelingToolkit ecosystem for solving. + +$MTK_NEEDED_STRING + +# Inputs +- `m::Model`: Model whose steady state should be solved. +- `sys::ModelingToolkit.NonlinearSystem`: Steady state system for `m`. + +# Options +- `u0 = zeros(length(unknowns(sys)))`: Initial guess of steady state variables. + `length(u0)` should be twice the number of variables in the model. + The ordering of the elements of `u0` should be + `[var1_level, var1_slope, var2_level, ..., varN_slope]`. +- `solver = nothing`: Solver to use. + `nothing` means use the default solver determined by `solve`. +- Additional options are passed as keyword arguments to `solve`. +""" +function solve_steady_state! end +export solve_steady_state! + +""" + steady_state_system(m::Model) + +Convert the steady state model associated with model `m` +into a `ModelingToolkit.NonlinearSystem`. + +$MTK_NEEDED_STRING + +# Example +This function is used to allow solving a `Model`'s steady state +using the ModelingToolkit/SciML ecosystem. +Here is an example of calling this function +and using [`solve_steady_state!`](@ref) on the result +to solve the model's steady state: +```julia +using ModelBaseEcon, StateSpaceEcon, ModelingToolkit, NonlinearSolve +@using_example E3 +m = E3.newmodel() +sys = steady_state_system(m) +solver = NewtonRaphson() # Replace with desired solver. +solve_steady_state!(m, sys; solver) +``` + +--- + + steady_state_system(sd::SteadyStateSolver.SolverData) + +Alternative call signature when `sd` is already available (e.g., created manually). +""" +function steady_state_system end +export steady_state_system + +""" + compute_residuals_steady_state(u, sd::SteadyStateSolver.SolverData) + +Compute the residuals of the steady state system in `sd` +given variable levels and slopes `u`. +This function is closed over in [`steady_state_system`](@ref) +to create a function that can be passed to `ModelingToolkit.NonlinearProblem`. + +$MTK_NEEDED_STRING + +!!! warning + Internal function not part of the public interface. +""" +function compute_residuals_steady_state end + +""" + sol2data(sol::AbstractSolution, s::ModelingToolkit.AbstractSystem, m::Model, p::Plan) + +Construct an MVTSeries of the full solution from the MTK solution. + +$MTK_NEEDED_STRING + +# Example +```julia +m = E3.newmodel() +IDX = size(plan_data, 1) +plan = Plan(m1, 1+m.maxlag:IDX-m.maxlead) +s = stacked_time_system(m, plan_data; fctype, plan) +nf = NonlinearFunction(s) +prob = NonlinearProblem(nf, zeros(length(unknowns(s)))) +sol = solve(prob, NewtonRaphson()) +result_mvts = sol2data(sol, s, m, plan) +``` +""" +function sol2data end +export sol2data + +end + +using .MTKExt + +export stacked_time_system +export steady_state_system +export solve_steady_state! +export sol2data diff --git a/src/StateSpaceEcon.jl b/src/StateSpaceEcon.jl index df7361d..8b5aed1 100644 --- a/src/StateSpaceEcon.jl +++ b/src/StateSpaceEcon.jl @@ -39,5 +39,6 @@ function getsolvermodule(solvername::Symbol) end include("simulate.jl") +include("MTKExt.jl") end # module diff --git a/src/stackedtime/solverdata.jl b/src/stackedtime/solverdata.jl index 6f51993..bb53207 100644 --- a/src/stackedtime/solverdata.jl +++ b/src/stackedtime/solverdata.jl @@ -88,7 +88,7 @@ function assign_fc! end return Dict{Int,Vector{Float64}}() end -@inline function assign_fc!(x::AbstractVector{Float64}, exog::AbstractVector{Float64}, vind::Int, sd::StackedTimeSolverData, ::FCNone) +@inline function assign_fc!(x::AbstractVector{<:Real}, exog::AbstractVector{Float64}, vind::Int, sd::StackedTimeSolverData, ::FCNone) # nothing to see here return x end @@ -100,7 +100,7 @@ end return Dict{Int,Vector{Float64}}() end -@inline function assign_fc!(x::AbstractVector{Float64}, exog::AbstractVector{Float64}, ::Int, sd::StackedTimeSolverData, ::FCGiven) +@inline function assign_fc!(x::AbstractVector{<:Real}, exog::AbstractVector{Float64}, ::Int, sd::StackedTimeSolverData, ::FCGiven) # values given exogenously if x !== exog x[sd.TT[end]] .= exog[sd.TT[end]] @@ -115,7 +115,7 @@ end return Dict{Int,Vector{Float64}}() end -@inline function assign_fc!(x::AbstractVector{Float64}, ::AbstractVector{Float64}, vind::Int, sd::StackedTimeSolverData, ::FCMatchSSLevel) +@inline function assign_fc!(x::AbstractVector{<:Real}, ::AbstractVector{Float64}, vind::Int, sd::StackedTimeSolverData, ::FCMatchSSLevel) # values come from the steady state x[sd.TT[end]] .= sd.SS[sd.TT[end], vind] return x @@ -144,7 +144,7 @@ end return Dict{Int,Vector{Float64}}(0 => fill(-1.0, length(sd.TT[end]))) end -@inline function assign_fc!(x::AbstractVector{Float64}, ::AbstractVector{Float64}, vind::Int, sd::StackedTimeSolverData, ::FCMatchSSRate) +@inline function assign_fc!(x::AbstractVector{<:Real}, ::AbstractVector{Float64}, vind::Int, sd::StackedTimeSolverData, ::FCMatchSSRate) # compute values using the slope from the steady state for t in sd.TT[end] x[t] = x[t-1] + sd.SS[t, vind] - sd.SS[t-1, vind] @@ -181,7 +181,7 @@ end return Dict{Int,Vector{Float64}}(-1 => foo, 0 => -1.0 .- foo) end -@inline function assign_fc!(x::AbstractVector{Float64}, ::AbstractVector{Float64}, ::Int, sd::StackedTimeSolverData, ::FCConstRate) +@inline function assign_fc!(x::AbstractVector{<:Real}, ::AbstractVector{Float64}, ::Int, sd::StackedTimeSolverData, ::FCConstRate) # compute values using the slope at the end of the simulation for t in sd.TT[end] x[t] = 2x[t-1] - x[t-2] @@ -430,7 +430,7 @@ exogenous data from `exog`. Also call [`assign_final_condition!`](@ref). !!! warning Internal function not part of the public interface. """ -function assign_exog_data!(x::AbstractArray{Float64,2}, exog::AbstractArray{Float64,2}, sd::StackedTimeSolverData) +function assign_exog_data!(x::AbstractArray{<:Real,2}, exog::AbstractArray{Float64,2}, sd::StackedTimeSolverData) # @assert size(x,1) == size(exog,1) == sd.NT x[sd.exog_mask] = exog[sd.exog_mask] assign_final_condition!(x, exog, sd) @@ -446,7 +446,7 @@ are stored in the the solver data `sd`. `exog` is used for [`fcgiven`](@ref). !!! warning Internal function not part of the public interface. """ -function assign_final_condition!(x::AbstractArray{Float64,2}, exog::AbstractArray{Float64,2}, sd::StackedTimeSolverData) +function assign_final_condition!(x::AbstractArray{<:Real,2}, exog::AbstractArray{Float64,2}, sd::StackedTimeSolverData) for (vi, fc) in enumerate(sd.FC) assign_fc!(view(x, :, vi), exog[:, vi], vi, sd, fc) end diff --git a/test/mtkext.jl b/test/mtkext.jl new file mode 100644 index 0000000..4b57b09 --- /dev/null +++ b/test/mtkext.jl @@ -0,0 +1,99 @@ +using ModelingToolkit +using NonlinearSolve +using SymbolicIndexingInterface +using DelimitedFiles +using TimeSeriesEcon + +function test_mtk(m_in, path; atol=1e-9) + + plan_data = readdlm(path, ',', Float64) + + for fctype in [fcgiven, fclevel, fcslope, fcnatural] + + m1 = deepcopy(m_in) + m2 = deepcopy(m_in) + + if fctype === fclevel || fctype === fcslope + + # Compute steady state with MTK. + sss = steady_state_system(m1) + solve_steady_state!(m1, sss; solver = NewtonRaphson()) + + # Compute steady state with StateSpaceEcon. + clear_sstate!(m2) + sssolve!(m2) + + @test isapprox(m1.sstate.values, m2.sstate.values; atol) + + end + + # Simulate with MTK. + IDX = size(plan_data, 1) + plan = Plan(m1, 1+m1.maxlag:IDX-m1.maxlead) + s = stacked_time_system(m1, plan_data; fctype, plan) + nf = NonlinearFunction(s) + prob = NonlinearProblem(nf, zeros(length(unknowns(s)))) + sol = solve(prob, NewtonRaphson()) + result = getu(sol, reduce(vcat, map(collect, map(var -> getproperty(s, var.name), m1.variables))))(sol) + result = reshape(result, size(plan_data, 1) - m1.maxlag - m1.maxlead, :) + + # Simulate with StateSpaceEcon. + nvars = nvariables(m2) + nshks = nshocks(m2) + IDX = size(plan_data, 1) + plan = Plan(m2, 1+m2.maxlag:IDX-m2.maxlead) + data = zeroarray(m2, plan) + data[m2.maxlag+1:end-m2.maxlead, nvars.+(1:nshks)] = plan_data[m2.maxlag+1:end-m2.maxlead, nvars.+(1:nshks)] # Shocks + data[1:m2.maxlag, :] = plan_data[1:m2.maxlag, :] # Initial conditions + if fctype === fcgiven + data[end-m2.maxlead+1:end, :] = plan_data[end-m2.maxlead+1:end, :] # Final conditions + end + correct = simulate(m2, plan, data; fctype) + correct_mvts = MVTSeries(1U:(IDX)U, m2.varshks, correct) + + result_mvts = sol2data(sol, s, m1, plan) + + @test isapprox(result, correct[m2.maxlag+1:end-m2.maxlead, 1:nvars]; atol) + @test isapprox(result_mvts[begin:end-m_in.maxlag-m_in.maxlead-1, :], correct_mvts[begin:end-m_in.maxlag-m_in.maxlead-1, :]; atol) + + # test with adjusted plan + # Simulate with MTK. + IDX = size(plan_data, 1) + plan = Plan(m1, 1+m1.maxlag:IDX-m1.maxlead) + autoexogenize!(plan, m1, first(rangeof(plan)):50U) + s = stacked_time_system(m1, plan_data; fctype, plan) + nf = NonlinearFunction(s) + prob = NonlinearProblem(nf, zeros(length(unknowns(s)))) + sol = solve(prob, NewtonRaphson()) + result = getu(sol, reduce(vcat, map(collect, map(var -> getproperty(s, var.name), m1.variables))))(sol) + result_mvts2 = sol2data(sol, s, m1, plan) + # result = reshape(result, size(plan_data, 1) - m1.maxlag - m1.maxlead, :) + + # Simulate with StateSpaceEcon. + nvars = nvariables(m2) + nshks = nshocks(m2) + IDX = size(plan_data, 1) + plan = Plan(m2, 1+m2.maxlag:IDX-m2.maxlead) + autoexogenize!(plan, m2, 1U:50U) + data = zeroarray(m2, plan) + data[m2.maxlag+1:50, 1:nvars] = plan_data[m2.maxlag+1:50, 1:nvars] # Variables (exogenous part) + data[51:end-m2.maxlead, nvars.+(1:nshks)] = plan_data[51:end-m2.maxlead, nvars.+(1:nshks)] # Shocks (exogenous part) + data[1:m2.maxlag, :] = plan_data[1:m2.maxlag, :] # Initial conditions + if fctype === fcgiven + data[end-m2.maxlead+1:end, :] = plan_data[end-m2.maxlead+1:end, :] # Final conditions + end + correct = simulate(m2, plan, data; fctype) + correct_mvts2 = MVTSeries(1U:(IDX)U, m2.varshks, correct) + + @test isapprox(result_mvts2[begin:end-m_in.maxlag-m_in.maxlead-1, :], correct_mvts2[begin:end-m_in.maxlag-m_in.maxlead-1, :]; atol) + end + +end + +@testset "E3.mtkext" begin + test_mtk(getE3(), joinpath(@__DIR__, "data", "M3_TestData.csv")) +end + +@testset "E7A.mtkext" begin + test_mtk(getE7A(), joinpath(@__DIR__, "data", "M7_TestData.csv")) +end diff --git a/test/runtests.jl b/test/runtests.jl index 0e5dc46..cb4a793 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -409,6 +409,8 @@ end include("sim_solver.jl") +include("mtkext.jl") + # keep this one last because it overwrites getE?() include("modelchanges.jl") From 0be725fbea64044a6728e01d303d8560e592b6f8 Mon Sep 17 00:00:00 2001 From: Jason Jensen Date: Thu, 18 Jul 2024 13:26:51 -0400 Subject: [PATCH 2/2] add 1.9 to testrun --- .github/workflows/main.yml | 2 +- test/runtests.jl | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 8f8094b..c7e2c00 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -20,7 +20,7 @@ jobs: strategy: matrix: - julia-version: ["1.10"] + julia-version: ["!.9","1.10"] julia-arch: [x64] os: [ubuntu-latest, windows-latest, macOS-latest] diff --git a/test/runtests.jl b/test/runtests.jl index cb4a793..d0ca2a9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -409,7 +409,9 @@ end include("sim_solver.jl") -include("mtkext.jl") +if VERSION >= v"1.10" + include("mtkext.jl") +end # keep this one last because it overwrites getE?() include("modelchanges.jl")