From db3a1a59a77c4ad0571f3c12fe569096882afe44 Mon Sep 17 00:00:00 2001
From: Nicholas Labelle St-Pierre <nlabellest-pierre@bank-banque-canada.ca>
Date: Fri, 4 Oct 2024 16:07:08 -0400
Subject: [PATCH] eval_equation and tests

---
 src/evaluation.jl | 55 +++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl  | 28 ++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/src/evaluation.jl b/src/evaluation.jl
index b628165..d6f68fc 100644
--- a/src/evaluation.jl
+++ b/src/evaluation.jl
@@ -496,6 +496,61 @@ function eval_RJ(point::Matrix{Float64}, slmed::SelectiveLinearizationMED)
     return res, jac
 end
 
+"""
+    eval_equation(model::AbstractModel, eqn::AbstractEquation, sim_data::AbstractMatrix{Float64}, rng::UnitRange{Int64} = 1:size(sim_data,1))
+
+Evaluate the residuals of a given equation over a range of time points.
+
+This function calculates the residuals of the provided equation `eqn` for each time step in the range `rng` from the simulated data `sim_data`. The model's lag and lead structure is respected during evaluation.
+
+# Arguments
+- `model::AbstractModel`: The model containing the equation to be evaluated.
+- `eqn::AbstractEquation`: The equation for which residuals are to be calculated.
+- `sim_data::AbstractMatrix{Float64}`: The simulated data, with rows representing time points and columns representing model.allvars (variables, shocks and auxiliary variables).
+- `rng::UnitRange{Int64}`: The range of time points over which to evaluate the equation. By default, evaluates over all time points in `sim_data`.
+
+# Returns
+- `res::Vector{Float64}`: A vector of residuals for each time point in the range `rng`. Entries for time points where residuals cannot be computed (due to insufficient lags or leads) are filled with `NaN`.
+"""
+function eval_equation(model::AbstractModel, eqn::AbstractEquation, sim_data::AbstractMatrix{Float64}, rng::UnitRange{Int64} = 1:size(sim_data,1))
+    # Check bounds
+    @assert rng[begin] >= 1 && rng[end] <= size(sim_data, 1) "Error: The range specified is out of bounds. Ensure that the range starts from 1 or higher and ends within the size of the data."
+
+    # Map the model variables to their respective indices
+    var_to_idx = _make_var_to_idx(model.allvars)
+
+    # Calculate t_start based on the model's maximum lag
+    t_start = 1 + model.maxlag
+
+    # Create index mapping for the equation's time series references
+    inds = [CartesianIndex((t_start + ti, var_to_idx[var])) for (var, ti) in keys(eqn.tsrefs)]
+
+    # Account for steady state values in case they are used
+    ed = DynEqnEvalData(eqn, model, var_to_idx)
+
+    # Initialize the residual vector with NaN values
+    res = fill(NaN, length(rng))
+
+    # Iterate over the specified time range
+    for (idx, t) = enumerate(rng)
+        # Define the range of data points required for evaluation, including lags and leads
+        rng_sub = t - model.maxlag : t + model.maxlead
+
+        # Ensure the subrange is within bounds of the data
+        if rng_sub[begin] >= 1 && rng_sub[end] <= size(sim_data, 1)
+            # Extract the relevant data points for the current time step
+            point = @view sim_data[rng_sub, :]
+
+            # Evaluate the residual for the current data point using the evaluation data structure
+            res[idx] = eval_resid(eqn, point[inds], ed)
+        end
+    end
+
+    # Return the vector of residuals
+    return res
+end
+export eval_equation
+
 """
     selective_linearize!(model)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 74adf30..24c1974 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1743,4 +1743,32 @@ end
         true
     end
 end
+
+@testset "eval_equation" begin
+    model = Model()
+    @parameters model p = 0.1
+    @variables model x
+    @shocks model x_shk
+    @equations model begin
+        :EQ01 => x[t] = (1-0.50)*@sstate(x) + 0.25*x[t-1] + 0.25*x[t+1] + x_shk[t]
+    end
+    @initialize model
+    @steadystate model x = 2.0
+    model.sstate.x.level = 2.0
+    sim_data = [0.5    0.0;
+    1.5980861244019138 0.0;
+    1.8923444976076556 0.0;
+    1.9712918660287082 0.0;
+    1.992822966507177  0.0;
+    2.0                0.0]
+
+    eqtn = model.equations[:EQ01]
+    res = eval_equation(model, eqtn, sim_data)
+    @test isnan(res[1]) && isapprox(res[2:5],[0.0, 0.0, 0.0, 0.0]; atol=1e-12) && isnan(res[6])
+
+    sim_data[3,1] += 1
+    @test eval_equation(model, eqtn, sim_data,3:4) ≈ [1.0, -0.25]
+
+    @test_throws AssertionError eval_equation(model, eqtn, sim_data, 1:7)
+end
 nothing
\ No newline at end of file