Systematically compute recursive ordering of auxiliary equations

Auxiliary equations appearing in set_auxiliary_variables.m and
dynamic_set_auxiliary_series.m need to appear in recursive ordering, since
those files are used for sequential evaluation.

Previously, the recursive ordering was guaranteed by a set of ad hoc rules and
workarounds, but that would not cover certain edge cases.

With this commit, the recursive ordering is systematically computed, using a
topological sort on the directed acyclic graph whose vertices are auxiliary
equations and whose edges are dependency relationships.

Closes: #22

src/DynamicModel.cc

@@ -6430,13 +6430,10 @@ DynamicModel::substituteLeadLagInternal(AuxVarType type, bool deterministic_mode
 
   // Add new equations
   for (auto & neweq : neweqs)
-    addEquation(neweq, -1);
-
-  // Order of auxiliary variable definition equations:
-  //  - expectation (entered before this function is called)
-  //  - lead variables from lower lead to higher lead
-  //  - lag variables from lower lag to higher lag
-  copy(neweqs.begin(), neweqs.end(), back_inserter(aux_equations));
+    {
+      addEquation(neweq, -1);
+      aux_equations.push_back(neweq);
+    }
 
   if (neweqs.size() > 0)
     {
@@ -6563,9 +6560,10 @@ DynamicModel::substituteUnaryOps(const vector<int> &eqnumbers)
 
   // Add new equations
   for (auto & neweq : neweqs)
-    addEquation(neweq, -1);
-
-  copy(neweqs.begin(), neweqs.end(), back_inserter(aux_equations));
+    {
+      addEquation(neweq, -1);
+      aux_equations.push_back(neweq);
+    }
 
   if (subst_table.size() > 0)
     cout << "Substitution of Unary Ops: added " << neweqs.size() << " auxiliary variables and equations." << endl;
@@ -6600,28 +6598,6 @@ DynamicModel::substituteDiff(vector<expr_t> &pac_growth)
     if (gv != nullptr)
       gv->findDiffNodes(diff_nodes);
 
-  /* Ensure that all diff operators appear once with their argument at current
-     period (i.e. index 0 in the equivalence class, see comment above
-     lag_equivalence_table_t in ExprNode.hh for details on the concepts).
-     If it is not the case, generate the corresponding expressions.
-     This is necessary to avoid lags of more than one in the auxiliary
-     equation, which would then be modified by subsequent transformations
-     (removing lags > 1), which in turn would break the recursive ordering
-     of auxiliary equations. See issue McModelTeam/McModelProject#95 */
-  for (auto &it : diff_nodes)
-    {
-      auto iterator_max_index = it.second.rbegin();
-      int max_index = iterator_max_index->first;
-      expr_t max_index_expr = iterator_max_index->second;
-
-      while (max_index < 0)
-        {
-          max_index++;
-          max_index_expr = max_index_expr->decreaseLeadsLags(-1);
-          it.second[max_index] = max_index_expr;
-        }
-    }
-
   // Substitute in model local variables
   vector<BinaryOpNode *> neweqs;
   for (auto & it : local_variables_table)
@@ -6642,9 +6618,10 @@ DynamicModel::substituteDiff(vector<expr_t> &pac_growth)
 
   // Add new equations
   for (auto & neweq : neweqs)
-    addEquation(neweq, -1);
-
-  copy(neweqs.begin(), neweqs.end(), back_inserter(aux_equations));
+    {
+      addEquation(neweq, -1);
+      aux_equations.push_back(neweq);
+    }
 
   if (diff_subst_table.size() > 0)
     cout << "Substitution of Diff operator: added " << neweqs.size() << " auxiliary variables and equations." << endl;
@@ -6672,10 +6649,10 @@ DynamicModel::substituteExpectation(bool partial_information_model)
 
   // Add new equations
   for (auto & neweq : neweqs)
-    addEquation(neweq, -1);
-
-  // Add the new set of equations at the *beginning* of aux_equations
-  copy(neweqs.rbegin(), neweqs.rend(), front_inserter(aux_equations));
+    {
+      addEquation(neweq, -1);
+      aux_equations.push_back(neweq);
+    }
 
   if (subst_table.size() > 0)
     {

src/ModFile.cc

@@ -652,6 +652,8 @@ ModFile::transformPass(bool nostrict, bool stochastic, bool compute_xrefs, const
         exit(EXIT_FAILURE);
       }
 
+  dynamic_model.reorderAuxiliaryEquations();
+
   // Freeze the symbol table
   symbol_table.freeze();
 

src/ModelTree.cc

@@ -2425,3 +2425,47 @@ ModelTree::compileDll(const string &basename, const string &static_or_dynamic, c
       exit(EXIT_FAILURE);
     }
 }
+
+void
+ModelTree::reorderAuxiliaryEquations()
+{
+  using namespace boost;
+
+  // Create the mapping between auxiliary variables and auxiliary equations
+  int n = static_cast<int>(aux_equations.size());
+  map<int, int> auxEndoToEq;
+  for (int i = 0; i < n; i++)
+    {
+      auto varexpr = dynamic_cast<VariableNode *>(aux_equations[i]->arg1);
+      assert(varexpr && symbol_table.getType(varexpr->symb_id) == SymbolType::endogenous);
+      auxEndoToEq[varexpr->symb_id] = i;
+    }
+  assert(static_cast<int>(auxEndoToEq.size()) == n);
+
+  /* Construct the directed acyclic graph where auxiliary equations are
+     vertices and edges represent dependency relationships. */
+  using Graph = adjacency_list<vecS, vecS, directedS>;
+  Graph g(n);
+  for (int i = 0; i < n; i++)
+    {
+      set<int> endos;
+      aux_equations[i]->collectVariables(SymbolType::endogenous, endos);
+      for (int endo : endos)
+        {
+          auto it = auxEndoToEq.find(endo);
+          if (it != auxEndoToEq.end() && it->second != i)
+            add_edge(i, it->second, g);
+        }
+    }
+
+  // Topological sort of the graph
+  using Vertex = graph_traits<Graph>::vertex_descriptor;
+  vector<Vertex> ordered;
+  topological_sort(g, back_inserter(ordered));
+
+  // Reorder auxiliary equations accordingly
+  auto aux_equations_old = aux_equations;
+  auto index = get(vertex_index, g); // Maps vertex descriptors to their index
+  for (int i = 0; i < n; i++)
+    aux_equations[i] = aux_equations_old[index[ordered[i]]];
+}

src/ModelTree.hh

@@ -95,7 +95,7 @@ protected:
       For example, such a divergence appears when there is an expectation
       operator in a ramsey model, see
       tests/optimal_policy/nk_ramsey_expectation.mod */
-  deque<BinaryOpNode *> aux_equations;
+  vector<BinaryOpNode *> aux_equations;
 
   //! Stores derivatives
   /*! Index 0 is not used, index 1 contains first derivatives, ...
@@ -363,6 +363,9 @@ public:
   void set_cutoff_to_zero();
   //! Simplify model equations: if a variable is equal to a constant, replace that variable elsewhere in the model
   void simplifyEquations();
+  /*! Reorder auxiliary variables so that they appear in recursive order in
+      set_auxiliary_variables.m and dynamic_set_auxiliary_series.m */
+  void reorderAuxiliaryEquations();
   //! Find equations where variable is equal to a constant
   void findConstantEquations(map<VariableNode *, NumConstNode *> &subst_table) const;
   //! Helper for writing the Jacobian elements in MATLAB and C