From e1e5118373801769aa4d2576a024b4343a1488f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Villemot?= <sebastien@dynare.org>
Date: Wed, 19 Jan 2022 15:53:45 +0100
Subject: [PATCH] Refactor the code that substitutes unary ops in model
 equations

In particular, make it more visible that the substitution also occurs in PAC
equations (and not only in VAR and TCM equations).
---
 src/DynamicModel.cc | 22 +++++++---------------
 src/DynamicModel.hh | 21 +++++++++++----------
 src/ModFile.cc      | 26 +++++++++++++++-----------
 3 files changed, 33 insertions(+), 36 deletions(-)

diff --git a/src/DynamicModel.cc b/src/DynamicModel.cc
index 4a3d159c..fe3cc898 100644
--- a/src/DynamicModel.cc
+++ b/src/DynamicModel.cc
@@ -5572,17 +5572,18 @@ DynamicModel::getEquationNumbersFromTags(const set<string> &eqtags) const
   return eqnumbers;
 }
 
-void
-DynamicModel::findPacExpectationEquationNumbers(set<int> &eqnumbers) const
+set<int>
+DynamicModel::findPacExpectationEquationNumbers() const
 {
+  set<int> eqnumbers;
   int i = 0;
   for (auto &equation : equations)
     {
-      if (equation->containsPacExpectation()
-          && find(eqnumbers.begin(), eqnumbers.end(), i) == eqnumbers.end())
+      if (equation->containsPacExpectation())
         eqnumbers.insert(i);
       i++;
     }
+  return eqnumbers;
 }
 
 pair<lag_equivalence_table_t, ExprNode::subst_table_t>
@@ -5590,20 +5591,11 @@ DynamicModel::substituteUnaryOps(PacModelTable &pac_model_table)
 {
   vector<int> eqnumbers(equations.size());
   iota(eqnumbers.begin(), eqnumbers.end(), 0);
-  return substituteUnaryOps(eqnumbers, pac_model_table);
-}
-
-pair<lag_equivalence_table_t, ExprNode::subst_table_t>
-DynamicModel::substituteUnaryOps(const set<string> &var_model_eqtags, PacModelTable &pac_model_table)
-{
-  set<int> eqnumbers = getEquationNumbersFromTags(var_model_eqtags);
-  findPacExpectationEquationNumbers(eqnumbers);
-  vector<int> eqnumbers_vec(eqnumbers.begin(), eqnumbers.end());
-  return substituteUnaryOps(eqnumbers_vec, pac_model_table);
+  return substituteUnaryOps(set<int>(eqnumbers.begin(), eqnumbers.end()), pac_model_table);
 }
 
 pair<lag_equivalence_table_t, ExprNode::subst_table_t>
-DynamicModel::substituteUnaryOps(const vector<int> &eqnumbers, PacModelTable &pac_model_table)
+DynamicModel::substituteUnaryOps(const set<int> &eqnumbers, PacModelTable &pac_model_table)
 {
   lag_equivalence_table_t nodes;
   ExprNode::subst_table_t subst_table;
diff --git a/src/DynamicModel.hh b/src/DynamicModel.hh
index 85be14bb..50cd6d47 100644
--- a/src/DynamicModel.hh
+++ b/src/DynamicModel.hh
@@ -235,10 +235,6 @@ private:
   //! Create a legacy *_dynamic.m file for Matlab/Octave not yet using the temporary terms array interface
   void writeDynamicMatlabCompatLayer(const string &basename) const;
 
-  set<int> getEquationNumbersFromTags(const set<string> &eqtags) const;
-
-  void findPacExpectationEquationNumbers(set<int> &eqnumber) const;
-
   //! Internal helper for the copy constructor and assignment operator
   /*! Copies all the structures that contain ExprNode*, by the converting the
       pointers into their equivalent in the new tree */
@@ -524,14 +520,13 @@ public:
   //! Substitutes out all model-local variables
   void substituteModelLocalVariables();
 
-  //! Creates aux vars for all unary operators
+  /* Creates aux vars for all unary operators in all equations. Also makes the
+     substitution in growth terms of pac_model/pac_target_info. */
   pair<lag_equivalence_table_t, ExprNode::subst_table_t> substituteUnaryOps(PacModelTable &pac_model_table);
 
-  //! Creates aux vars for unary operators in certain equations: originally implemented for support of VARs
-  pair<lag_equivalence_table_t, ExprNode::subst_table_t> substituteUnaryOps(const set<string> &eq_tags, PacModelTable &pac_model_table);
-
-  //! Creates aux vars for unary operators in certain equations: originally implemented for support of VARs
-  pair<lag_equivalence_table_t, ExprNode::subst_table_t> substituteUnaryOps(const vector<int> &eqnumbers, PacModelTable &pac_model_table);
+  /* Creates aux vars for all unary operators in specified equations. Also makes the
+     substitution in growth terms of pac_model/pac_target_info. */
+  pair<lag_equivalence_table_t, ExprNode::subst_table_t> substituteUnaryOps(const set<int> &eqnumbers, PacModelTable &pac_model_table);
 
   //! Substitutes diff operator
   pair<lag_equivalence_table_t, ExprNode::subst_table_t> substituteDiff(PacModelTable &pac_model_table);
@@ -643,5 +638,11 @@ public:
   //! Simplify model equations: if a variable is equal to a constant, replace that variable elsewhere in the model
   /*! Equations with MCP tags are excluded, see dynare#1697 */
   void simplifyEquations();
+
+  // Converts a set of equation tags into the corresponding set of equation numbers
+  set<int> getEquationNumbersFromTags(const set<string> &eqtags) const;
+
+  // Returns the set of equations (as numbers) which have a pac_expectation operator
+  set<int> findPacExpectationEquationNumbers() const;
 };
 #endif
diff --git a/src/ModFile.cc b/src/ModFile.cc
index 1ee64c9b..3c797720 100644
--- a/src/ModFile.cc
+++ b/src/ModFile.cc
@@ -412,15 +412,19 @@ ModFile::transformPass(bool nostrict, bool stochastic, bool compute_xrefs, bool
   if (unusedEndogsIsErr)
     exit(EXIT_FAILURE);
 
-  // Get all equation tags associated with VARs and Trend Component Models
-  set<string> eqtags;
-  for (const auto &it : trend_component_model_table.getEqTags())
-    for (auto &it1 : it.second)
-      eqtags.insert(it1);
-
-  for (const auto &it : var_model_table.getEqTags())
-    for (auto &it1 : it.second)
-      eqtags.insert(it1);
+  /* Get the list of equations in which to scan for and substitute unary ops:
+     – equations which are part of VARs and Trend Component Models
+     – PAC equations (those with a pac_expectation operator) */
+  set<string> var_tcm_eqtags;
+  for (const auto &[name, tags] : trend_component_model_table.getEqTags())
+    for (auto &tag : tags)
+      var_tcm_eqtags.insert(tag);
+  for (const auto &[name, tags] : var_model_table.getEqTags())
+    for (auto &tag : tags)
+      var_tcm_eqtags.insert(tag);
+
+  set<int> unary_ops_eqs = dynamic_model.getEquationNumbersFromTags(var_tcm_eqtags);
+  unary_ops_eqs.merge(dynamic_model.findPacExpectationEquationNumbers());
 
   // Create auxiliary variables and equations for unary ops
   lag_equivalence_table_t unary_ops_nodes;
@@ -428,8 +432,8 @@ ModFile::transformPass(bool nostrict, bool stochastic, bool compute_xrefs, bool
   if (transform_unary_ops)
     tie(unary_ops_nodes, unary_ops_subst_table) = dynamic_model.substituteUnaryOps(pac_model_table);
   else
-    // substitute only those unary ops that appear in auxiliary model equations
-    tie(unary_ops_nodes, unary_ops_subst_table) = dynamic_model.substituteUnaryOps(eqtags, pac_model_table);
+    // substitute only those unary ops that appear in VAR, TCM and PAC model equations
+    tie(unary_ops_nodes, unary_ops_subst_table) = dynamic_model.substituteUnaryOps(unary_ops_eqs, pac_model_table);
 
   // Create auxiliary variable and equations for Diff operators
   auto [diff_nodes, diff_subst_table] = dynamic_model.substituteDiff(pac_model_table);
-- 
GitLab