diff --git a/src/DynamicModel.cc b/src/DynamicModel.cc index 14135914a7fc5b0cf101f5710fb53620ef3f10d4..6116b85a34320df2823670b9f18b0fce3ecb2edc 100644 --- a/src/DynamicModel.cc +++ b/src/DynamicModel.cc @@ -5516,13 +5516,25 @@ DynamicModel::writeAuxVarRecursiveDefinitions(ostream &output, ExprNodeOutputTyp } } +void +DynamicModel::clearEquations() +{ + equations.clear(); + equations_lineno.clear(); + equation_tags.clear(); + equation_tags_xref.clear(); +} + void DynamicModel::replaceMyEquations(DynamicModel &dynamic_model) const { - dynamic_model.equations.clear(); + dynamic_model.clearEquations(); + for (size_t i = 0; i < equations.size(); i++) - dynamic_model.addEquation(equations[i]->clone(dynamic_model), - equations_lineno[i]); + dynamic_model.addEquation(equations[i]->clone(dynamic_model), equations_lineno[i]); + + dynamic_model.equation_tags = equation_tags; + dynamic_model.equation_tags_xref = equation_tags_xref; } void @@ -5539,9 +5551,9 @@ DynamicModel::computeRamseyPolicyFOCs(const StaticModel &static_model) } cout << "Ramsey Problem: added " << i << " Multipliers." << endl; - // Add Planner Objective to equations to include in computeDerivIDs + // Add Planner Objective to equations so that it appears in Lagrangian assert(static_model.equations.size() == 1); - addEquation(static_model.equations[0]->clone(*this), static_model.equations_lineno[0]); + addEquation(static_model.equations[0]->clone(*this), -1); // Get max endo lead and max endo lag set<pair<int, int>> dynvars; @@ -5550,9 +5562,8 @@ DynamicModel::computeRamseyPolicyFOCs(const StaticModel &static_model) for (auto &equation : equations) equation->collectDynamicVariables(SymbolType::endogenous, dynvars); - for (const auto &dynvar : dynvars) + for (const auto &[symb_id, lag] : dynvars) { - int lag = dynvar.second; if (max_eq_lead < lag) max_eq_lead = lag; else if (-max_eq_lag > lag) @@ -5584,21 +5595,47 @@ DynamicModel::computeRamseyPolicyFOCs(const StaticModel &static_model) equations[i]->getNonZeroPartofEquation()->decreaseLeadsLags(lag)), lagrangian); } - equations.clear(); + // Save line numbers and tags, see below + auto old_equations_lineno = equations_lineno; + auto old_equation_tags = equation_tags; + + // Prepare derivation of the Lagrangian + clearEquations(); addEquation(AddEqual(lagrangian, Zero), -1); computeDerivIDs(); - //Compute derivatives and overwrite equations + /* Compute Lagrangian derivatives. + Also restore line numbers and tags for FOCs w.r.t. a Lagrange multiplier + (i.e. a FOC identical to an equation of the original model) */ vector<expr_t> neweqs; - for (auto &it : deriv_id_table) - // For all endogenous variables with zero lag - if (symbol_table.getType(it.first.first) == SymbolType::endogenous && it.first.second == 0) - neweqs.push_back(AddEqual(equations[0]->getNonZeroPartofEquation()->getDerivative(it.second), Zero)); + vector<int> neweqs_lineno; + map<int, vector<pair<string, string>>> neweqs_tags; + for (auto &[symb_id_and_lag, deriv_id] : deriv_id_table) + { + auto &[symb_id, lag] = symb_id_and_lag; + if (symbol_table.getType(symb_id) == SymbolType::endogenous && lag == 0) + { + neweqs.push_back(AddEqual(equations[0]->getNonZeroPartofEquation()->getDerivative(deriv_id), Zero)); + if (int i = symbol_table.getEquationNumberForMultiplier(symb_id); + i != -1) + { + // This is a derivative w.r.t. a Lagrange multiplier + neweqs_lineno.push_back(old_equations_lineno[i]); + vector<pair<string, string>> tags; + for (auto &[j, tagpair] : old_equation_tags) + if (j == i) + tags.emplace_back(tagpair); + neweqs_tags[neweqs.size()-1] = tags; + } + else + neweqs_lineno.push_back(-1); + } + } - // Add new equations - equations.clear(); - for (auto &neweq : neweqs) - addEquation(neweq, -1); + // Overwrite equations with the Lagrangian derivatives + clearEquations(); + for (size_t i = 0; i < neweqs.size(); i++) + addEquation(neweqs[i], neweqs_lineno[i], neweqs_tags[i]); } void diff --git a/src/DynamicModel.hh b/src/DynamicModel.hh index 9a2db5c758da370fe523a91fa20734c2532d68eb..487017a65eb40fc3c2cb5c780cc12865d194102d 100644 --- a/src/DynamicModel.hh +++ b/src/DynamicModel.hh @@ -446,6 +446,10 @@ public: //! Replaces model equations with derivatives of Lagrangian w.r.t. endogenous void computeRamseyPolicyFOCs(const StaticModel &static_model); + + //! Clears all equations + void clearEquations(); + //! Replaces the model equations in dynamic_model with those in this model void replaceMyEquations(DynamicModel &dynamic_model) const; diff --git a/src/ModelTree.hh b/src/ModelTree.hh index d04f5d4d9d7c1ad82700f9ee6423c995ba6476b6..040b26965ba3bb7e2d439298845092b6e3257369 100644 --- a/src/ModelTree.hh +++ b/src/ModelTree.hh @@ -77,6 +77,8 @@ protected: * The following structures keep track of the model equations and must all be updated * when adding or removing an equation. Hence, if a new parallel structure is added * in the future, it must be maintained whereever these structures are updated + * See in particular methods clearEquations(), replaceMyEquations() and + * computeRamseyPolicyFOCs() of DynamicModel class. * NB: This message added with the introduction of the `exclude_eqs` option, hence * that's a place to update future structures. */ diff --git a/src/SymbolTable.cc b/src/SymbolTable.cc index 5363026fcd94fd7a4b96e72bfe610db51ea6f89d..895a2849497c8cdb0d5f372ecaee3895010f6586 100644 --- a/src/SymbolTable.cc +++ b/src/SymbolTable.cc @@ -1096,3 +1096,12 @@ SymbolTable::getUltimateOrigSymbID(int symb_id) const } return symb_id; } + +int +SymbolTable::getEquationNumberForMultiplier(int symb_id) const +{ + for (const auto &aux_var : aux_vars) + if (aux_var.get_symb_id() == symb_id && aux_var.get_type() == AuxVarType::multiplier) + return aux_var.get_equation_number_for_multiplier(); + return -1; +} diff --git a/src/SymbolTable.hh b/src/SymbolTable.hh index d601f6f19ff230b483c136a9fe8a30d5ac5fd5b7..1d31964a23bf9fe5bea2232ae00e6d59f0dcedcc 100644 --- a/src/SymbolTable.hh +++ b/src/SymbolTable.hh @@ -412,6 +412,8 @@ public: repeatedly call getOrigSymbIDForAuxVar() until an original (non-auxiliary) variable is found. */ int getUltimateOrigSymbID(int symb_id) const; + //! If this is a Lagrange multiplier, return its associated equation number; otherwise return -1 + int getEquationNumberForMultiplier(int symb_id) const; }; inline void