From 3694fc40b3a06933ac738ba694a9b13df9be31ec Mon Sep 17 00:00:00 2001
From: Houtan Bastani <houtan@dynare.org>
Date: Wed, 4 Dec 2019 15:48:33 +0100
Subject: [PATCH] detrend epilogue equations

---
 src/DynamicModel.hh       | 10 ++++++++++
 src/ModFile.cc            |  3 +++
 src/ModelEquationBlock.cc | 25 +++++++++++++++++++++++++
 src/ModelEquationBlock.hh |  6 +++++-
 4 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/DynamicModel.hh b/src/DynamicModel.hh
index 31ee8c01..575c72ae 100644
--- a/src/DynamicModel.hh
+++ b/src/DynamicModel.hh
@@ -486,6 +486,16 @@ public:
   //! Transforms the model by removing trends specified by the user
   void detrendEquations();
 
+  inline const nonstationary_symbols_map_t & getNonstationarySymbolsMap() const
+  {
+    return nonstationary_symbols_map;
+  }
+
+  inline const map<int, expr_t> & getTrendSymbolsMap() const
+  {
+    return trend_symbols_map;
+  }
+
   //! Substitutes adl operator
   void substituteAdl();
 
diff --git a/src/ModFile.cc b/src/ModFile.cc
index 837de22f..50a560a3 100644
--- a/src/ModFile.cc
+++ b/src/ModFile.cc
@@ -514,6 +514,9 @@ ModFile::transformPass(bool nostrict, bool stochastic, bool compute_xrefs, const
       dynamic_model.detrendEquations();
       trend_dynamic_model = dynamic_model;
       dynamic_model.removeTrendVariableFromEquations();
+      const auto & trend_symbols = dynamic_model.getTrendSymbolsMap();
+      const auto & nonstationary_symbols = dynamic_model.getNonstationarySymbolsMap();
+      epilogue.detrend(trend_symbols, nonstationary_symbols);
     }
 
   mod_file_struct.orig_eq_nbr = dynamic_model.equation_number();
diff --git a/src/ModelEquationBlock.cc b/src/ModelEquationBlock.cc
index cd2e7031..a4eaf92b 100644
--- a/src/ModelEquationBlock.cc
+++ b/src/ModelEquationBlock.cc
@@ -335,6 +335,31 @@ Epilogue::checkPass(WarningConsolidation &warnings) const
       so_far_defined.push_back(it.first);
 }
 
+void
+Epilogue::detrend(const map<int, expr_t> & trend_symbols_map,
+                  const nonstationary_symbols_map_t & nonstationary_symbols_map)
+{
+  for (auto it = nonstationary_symbols_map.crbegin();
+       it != nonstationary_symbols_map.crend(); it++)
+    for (auto & [symb_id, expr] : def_table)
+      {
+        expr = expr->detrend(it->first, it->second.first, it->second.second);
+        assert(expr != nullptr);
+      }
+
+  for (auto & [symb_id, expr] : def_table)
+    {
+      expr = expr->removeTrendLeadLag(trend_symbols_map);
+      assert(expr != nullptr);
+    }
+
+  for (auto & [symb_id, expr] : def_table)
+    {
+      expr = expr->replaceTrendVar();
+      assert(expr != nullptr);
+    }
+}
+
 void
 Epilogue::writeEpilogueFile(const string &basename) const
 {
diff --git a/src/ModelEquationBlock.hh b/src/ModelEquationBlock.hh
index f01605eb..9feb71af 100644
--- a/src/ModelEquationBlock.hh
+++ b/src/ModelEquationBlock.hh
@@ -1,5 +1,5 @@
 /*
- * Copyright © 2010-2018 Dynare Team
+ * Copyright © 2010-2019 Dynare Team
  *
  * This file is part of Dynare.
  *
@@ -89,6 +89,10 @@ public:
   //! Checks that no variable is declared twice
   void checkPass(WarningConsolidation &warnings) const;
 
+  //! Deal with trend variables in the epilogue block
+  void detrend(const map<int, expr_t> & trend_symbols_map,
+               const nonstationary_symbols_map_t & nonstationary_symbols_map);
+
   //! Write the steady state file
   void writeEpilogueFile(const string &basename) const;
 };