From 8a2c38cf6c1645da453ee6c5cc3a8eb3f221174a Mon Sep 17 00:00:00 2001
From: Houtan Bastani <houtan@dynare.org>
Date: Wed, 24 Oct 2018 17:13:06 +0200
Subject: [PATCH] olsgibbs: save fitted values in dataset; allow use of
 dictionary as in dyn_ols

---
 matlab/olsgibbs.m                        | 43 ++++++++++++++++++++++--
 tests/estimation/univariate/bayesian.mod |  3 +-
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/matlab/olsgibbs.m b/matlab/olsgibbs.m
index 7757b5af3..2d210d728 100644
--- a/matlab/olsgibbs.m
+++ b/matlab/olsgibbs.m
@@ -1,4 +1,4 @@
-function olsgibbs(ds, eqtag, BetaPriorExpectation, BetaPriorVariance, s2, nu, ndraws, discarddraws, thin)
+function ds = olsgibbs(ds, eqtag, BetaPriorExpectation, BetaPriorVariance, s2, nu, ndraws, discarddraws, thin, fitted_names_dict)
 
 % Implements Gibbs Samipling for univariate linear model.
 %
@@ -13,9 +13,16 @@ function olsgibbs(ds, eqtag, BetaPriorExpectation, BetaPriorVariance, s2, nu, nd
 % - ndraws                      [integer]    scalar, total number of draws (Gibbs sampling)
 % - discarddraws                [integer]    scalar, number of draws to be discarded.
 % - thin                        [integer]    scalar, if thin == N, save every Nth draw (default is 1).
+% - fitted_names_dict           [cell]       Nx2 or Nx3 cell array to be used in naming fitted
+%                                            values; first column is the equation tag,
+%                                            second column is the name of the
+%                                            associated fitted value, third column
+%                                            (if it exists) is the function name of
+%                                            the transformation to perform on the
+%                                            fitted value.
 %
 % OUTPUTS
-% - none
+% - ds                          [dseries]    dataset updated with fitted value
 %
 % SPECIAL REQUIREMENTS
 %   none
@@ -47,7 +54,7 @@ global M_ oo_ options_
 
 % Check input
 
-if nargin < 7 || nargin > 9
+if nargin < 7 || nargin > 10
     error('Incorrect number of arguments passed to olsgibbs')
 end
 
@@ -109,6 +116,16 @@ else
     end
 end
 
+if nargin == 9
+    fitted_names_dict = {};
+else
+    if (~isempty(fitted_names_dict) && ...
+            (~iscell(fitted_names_dict) || ...
+            (size(fitted_names_dict, 2) < 2 || size(fitted_names_dict, 2) > 3)))
+        error('The 10th argument must be an Nx2 or Nx3 cell array');
+    end
+end
+
 % Let dyn_ols parse the equation to be estimated.
 [N, pnames, X, Y, equation, fp, lp] = dyn_ols(ds, {}, {eqtag});
 
@@ -169,6 +186,26 @@ oo_.olsgibbs.(eqtag).posterior.variance.h = var(oo_.olsgibbs.(eqtag).draws(:,n+1
 oo_.olsgibbs.(eqtag).s2 = mean(oo_.olsgibbs.(eqtag).draws(:,n+2));
 oo_.olsgibbs.(eqtag).R2 = mean(oo_.olsgibbs.(eqtag).draws(:,n+3));
 
+% Yhat
+idx = 0;
+yhatname = [eqtag '_olsgibbs_FIT'];
+if ~isempty(fitted_names_dict)
+    idx = strcmp(fitted_names_dict(:,1), eqtag);
+    if any(idx)
+        yhatname = fitted_names_dict{idx, 2};
+    end
+end
+oo_.olsgibbs.(eqtag).Yhat = dseries(X*oo_.olsgibbs.(eqtag).posterior.mean.beta, fp, yhatname);
+
+% Apply correcting function for Yhat if it was passed
+if any(idx) ...
+        && length(fitted_names_dict(idx, :)) == 3 ...
+        && ~isempty(fitted_names_dict{idx, 3})
+    oo_.olsgibbs.(eqtag).Yhat = ...
+        feval(fitted_names_dict{idx, 3}, oo_.olsgibbs.(eqtag).Yhat);
+end
+ds = [ds oo_.olsgibbs.(eqtag).Yhat];
+
 % Compute and save posterior densities.
 for i=1:n
     xx = oo_.olsgibbs.(eqtag).draws(:,i);
diff --git a/tests/estimation/univariate/bayesian.mod b/tests/estimation/univariate/bayesian.mod
index 513ff128d..d5dcce24f 100644
--- a/tests/estimation/univariate/bayesian.mod
+++ b/tests/estimation/univariate/bayesian.mod
@@ -51,8 +51,7 @@ gibbslength = 1000000; // Set the number of iterations in Gibbs
 burnin = 10000;        // Set the number of iterations to be discarded (try to remove the effects of the initial condition).
 steps = 10;            // Do not record all iterations (try to remove the dependence between the draws).
 
-
-olsgibbs(ds, 'eqols', beta0, V0, s2priormean, s2df, gibbslength, burnin, steps);
+ds = olsgibbs(ds, 'eqols', beta0, V0, s2priormean, s2df, gibbslength, burnin, steps, {'eqols', 'eqols_olsgibbs_fit'});
 
 // Since we use a diffuse prior for β, the posterior mean of β should be close to the OLS estimate.
 if max(abs(oo_.ols.eqols.beta-oo_.olsgibbs.eqols.posterior.mean.beta))>.1