Add dynamic equations function to model structure for direct equation evaluation

src/MacroModelling.jl

@@ -143,7 +143,7 @@ export get_irfs, get_irf, get_IRF, simulate, get_simulation, get_simulations, ge
 export get_conditional_forecast
 export get_solution, get_first_order_solution, get_perturbation_solution, get_second_order_solution, get_third_order_solution
 export get_steady_state, get_SS, get_ss, get_non_stochastic_steady_state, get_stochastic_steady_state, get_SSS, steady_state, SS, SSS, ss, sss
-export get_non_stochastic_steady_state_residuals, get_residuals, check_residuals
+export get_non_stochastic_steady_state_residuals, get_residuals, check_residuals, get_dynamic_residuals
 export get_moments, get_statistics, get_covariance, get_standard_deviation, get_variance, get_var, get_std, get_stdev, get_cov, var, std, stdev, cov, get_mean #, mean
 export get_autocorrelation, get_correlation, get_variance_decomposition, get_corr, get_autocorr, get_var_decomp, corr, autocorr
 export get_fevd, fevd, get_forecast_error_variance_decomposition, get_conditional_variance_decomposition
@@ -6982,6 +6982,63 @@ function write_functions_mapping!(𝓂::ℳ, max_perturbation_order::Int;
 
     𝓂.jacobian = buffer, func_exprs
 
+    # Generate dynamic equations function
+    # This function evaluates the dynamic equations themselves (not derivatives)
+    # and fills a pre-allocated residual vector
+    dyn_eqs_vector = collect(dyn_equations)
+
+    lennz_dyn_eqs = count(!iszero, dyn_eqs_vector)
+
+    if lennz_dyn_eqs > nnz_parallel_threshold
+        parallel_dyn = Symbolics.ShardedForm(1500,4)
+    else
+        parallel_dyn = Symbolics.SerialForm()
+    end
+
+    # First generate the core function with the standard signature
+    _, func_dyn_eqs_core = Symbolics.build_function(dyn_eqs_vector, 𝔓, 𝔙, 
+                                            cse = cse, 
+                                            skipzeros = skipzeros, 
+                                            parallel = parallel_dyn,
+                                            expression_module = @__MODULE__,
+                                            expression = Val(false))::Tuple{<:Function, <:Function}
+
+    # Create wrapper function with user-friendly signature that handles indexing internally
+    # Signature: func!(residual, parameters, past, present, future, steady_state, shocks)
+    n_params = length(𝓂.parameters)
+    n_calib_params = length(𝓂.calibration_equations_parameters)
+    n_vars = length(𝓂.var)
+    n_exo = length(𝓂.exo)
+
+    # Capture indices in closure
+    local_dyn_var_future_idx = copy(dyn_var_future_idx)
+    local_dyn_var_present_idx = copy(dyn_var_present_idx)
+    local_dyn_var_past_idx = copy(dyn_var_past_idx)
+    local_dyn_ss_idx = copy(dyn_ss_idx)
+
+    func_dyn_eqs = function(residual, parameters, past, present, future, steady_state, shocks)
+        # Build the 𝔓 vector: [parameters, calibration_parameters, steady_state_values]
+        pars_ext = vcat(parameters, zeros(n_calib_params))
+        params_and_SS = vcat(pars_ext, steady_state)
+
+        # Build the 𝔙 vector: [future[indices], present[indices], past[indices], shocks]
+        n_future = length(local_dyn_var_future_idx)
+        n_present = length(local_dyn_var_present_idx)
+        n_past = length(local_dyn_var_past_idx)
+
+        var_vec = zeros(n_future + n_present + n_past + n_exo)
+
+        var_vec[1:n_future] = future[local_dyn_var_future_idx]
+        var_vec[n_future+1:n_future+n_present] = present[local_dyn_var_present_idx]
+        var_vec[n_future+n_present+1:n_future+n_present+n_past] = past[local_dyn_var_past_idx]
+        var_vec[end-n_exo+1:end] = shocks
+
+        # Call the core generated function
+        func_dyn_eqs_core(residual, params_and_SS, var_vec)
+    end
+
+    𝓂.dyn_equations_func = func_dyn_eqs
+
 
     ∇₁_parameters = derivatives[1][2][:,1:nps]
 

src/get_functions.jl

@@ -3661,4 +3661,79 @@ get_residuals = get_non_stochastic_steady_state_residuals
 """
 See [`get_non_stochastic_steady_state_residuals`](@ref)
 """
-check_residuals = get_non_stochastic_steady_state_residuals
+check_residuals = get_non_stochastic_steady_state_residuals
+
+
+"""
+$(SIGNATURES)
+Evaluate the dynamic equations of the model and fill the pre-allocated residual vector.
+
+This function provides a convenient interface to evaluate the model's dynamic equations
+at any point in the state space. It is particularly useful for:
+- Verifying that the steady state satisfies the dynamic equations
+- Debugging model specifications
+- Custom solution algorithms
+
+# Arguments
+- `residual`: Pre-allocated vector to store the residuals (modified in-place)
+- `parameters`: Vector of parameter values
+- `past`: Vector of past variable values (t-1)
+- `present`: Vector of present variable values (t)
+- `future`: Vector of future variable values (t+1)
+- `steady_state`: Vector of steady state values
+- `shocks`: Vector of shock values (exogenous variables)
+- `𝓂`: Model object
+
+# Returns
+- Nothing (residuals are filled in the `residual` vector)
+
+# Examples
+```jldoctest
+using MacroModelling
+
+@model RBC begin
+    1  /  c[0] = (β  /  c[1]) * (α * exp(z[1]) * k[0]^(α - 1) + (1 - δ))
+    c[0] + k[0] = (1 - δ) * k[-1] + q[0]
+    q[0] = exp(z[0]) * k[-1]^α
+    z[0] = ρ * z[-1] + std_z * eps_z[x]
+end
+
+@parameters RBC begin
+    std_z = 0.01
+    ρ = 0.2
+    δ = 0.02
+    α = 0.5
+    β = 0.95
+end
+
+# Solve and get steady state
+solve!(RBC)
+SS = get_steady_state(RBC)
+
+# Get steady state values for the SS argument
+aux_idx = RBC.solution.perturbation.auxiliary_indices
+SS_for_func = SS[aux_idx.dyn_ss_idx]
+
+# Allocate residual vector
+residual = zeros(length(RBC.dyn_equations))
+
+# Evaluate at steady state with zero shocks
+get_dynamic_residuals(residual, RBC.parameter_values, SS, SS, SS, SS_for_func, zeros(length(RBC.exo)), RBC)
+
+# Residuals should be near zero at steady state
+maximum(abs.(residual))
+# output
+3.552713678800501e-15
+```
+"""
+function get_dynamic_residuals(residual::Vector{Float64},
+                               parameters::Vector{Float64},
+                               past::Vector{Float64},
+                               present::Vector{Float64},
+                               future::Vector{Float64},
+                               steady_state::Vector{Float64},
+                               shocks::Vector{Float64},
+                               𝓂::ℳ)
+    𝓂.dyn_equations_func(residual, parameters, past, present, future, steady_state, shocks)
+    return nothing
+end

src/macros.jl

@@ -901,6 +901,7 @@ macro model(𝓂,ex...)
                         (zeros(0,0), x->x), # third_order_derivatives
                         (zeros(0,0), x->x), # third_order_derivatives_parameters
                         (zeros(0,0), x->x), # third_order_derivatives_SS_and_pars
+                        x->x, # dyn_equations_func
                         # (x->x, SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0), 𝒟.prepare_jacobian(x->x, 𝒟.AutoForwardDiff(), [0]), SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)), # third_order_derivatives
                         # ([], SparseMatrixCSC{Float64, Int64}(ℒ.I, 0, 0)), # model_jacobian
                         # ([], Int[], zeros(1,1)), # model_jacobian

src/structures.jl

@@ -422,6 +422,7 @@ mutable struct ℳ
     third_order_derivatives::Tuple{AbstractMatrix{<: Real},Function}
     third_order_derivatives_parameters::Tuple{AbstractMatrix{<: Real},Function}
     third_order_derivatives_SS_and_pars::Tuple{AbstractMatrix{<: Real},Function}
+    dyn_equations_func::Function
 
     # model_jacobian::Tuple{Vector{Function}, SparseMatrixCSC{Float64}}
     # model_jacobian::Tuple{Vector{Function}, Vector{Int}, Matrix{<: Real}}