diff --git a/include/pyoptinterface/cppad_interface.hpp b/include/pyoptinterface/cppad_interface.hpp index c76e46f..e1d67f8 100644 --- a/include/pyoptinterface/cppad_interface.hpp +++ b/include/pyoptinterface/cppad_interface.hpp @@ -32,11 +32,12 @@ ADFunDouble sparse_hessian(const ADFunDouble &f, const sparsity_pattern_t &patte const std::vector &p_values); // Transform ExpressionGraph to CppAD function -// selected_outputs: indices of outputs to trace, empty means all outputs +// selected: indices of outputs to trace, empty means all outputs ADFunDouble cppad_trace_graph_constraints(const ExpressionGraph &graph, - const std::vector &selected_outputs = {}); -ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, bool aggregate = true, - const std::vector &selected_outputs = {}); + const std::vector &selected = {}); +ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, + const std::vector &selected = {}, + bool aggregate = true); struct CppADAutodiffGraph { diff --git a/include/pyoptinterface/knitro_model.hpp b/include/pyoptinterface/knitro_model.hpp index c3b1711..d3413bd 100644 --- a/include/pyoptinterface/knitro_model.hpp +++ b/include/pyoptinterface/knitro_model.hpp @@ -161,93 +161,60 @@ struct CallbackEvaluator void setup() { fun.optimize(); - CppAD::sparse_rc> jac_pattern_in(fun.Range(), fun.Range(), fun.Range()); - for (size_t k = 0; k < fun.Range(); k++) - { - jac_pattern_in.set(k, k, k); - } - fun.rev_jac_sparsity(jac_pattern_in, false, false, true, jac_pattern_); - jac_pattern_in.resize(fun.Domain(), fun.Domain(), fun.Domain()); - for (size_t i = 0; i < fun.Domain(); i++) + auto nx = fun.Domain(); + auto ny = fun.Range(); + CppAD::sparse_rc> jac_pattern_in(nx, nx, nx); + for (size_t i = 0; i < nx; i++) { jac_pattern_in.set(i, i, i); } - CppAD::sparse_rc> jac_pattern_out; - fun.for_jac_sparsity(jac_pattern_in, false, false, true, jac_pattern_out); - std::vector select_rows(fun.Range(), true); + fun.for_jac_sparsity(jac_pattern_in, false, false, true, jac_pattern_); + std::vector select_rows(ny, true); fun.rev_hes_sparsity(select_rows, false, true, hess_pattern_); + auto &hess_rows = hess_pattern_.row(); + auto &hess_cols = hess_pattern_.col(); for (size_t k = 0; k < hess_pattern_.nnz(); k++) { - size_t row = hess_pattern_.row()[k]; - size_t col = hess_pattern_.col()[k]; + size_t row = hess_rows[k]; + size_t col = hess_cols[k]; if (row <= col) { hess_pattern_symm_.push_back(row, col); } } - x.resize(fun.Domain(), 0.0); - w.resize(fun.Range(), 0.0); + x.resize(nx, 0.0); + w.resize(ny, 0.0); jac_ = CppAD::sparse_rcv, std::vector>(jac_pattern_); hess_ = CppAD::sparse_rcv, std::vector>(hess_pattern_symm_); } - void eval_fun(const V *req_x, V *res_y, bool aggregate = false) + bool is_objective() const { - for (size_t i = 0; i < indexVars.size(); i++) - { - x[i] = req_x[indexVars[i]]; - } + return indexCons.empty(); + } + + void eval_fun(const V *req_x, V *res_y) + { + copy_ptr(req_x, indexVars.data(), x); auto y = fun.Forward(0, x); - for (size_t k = 0; k < fun.Range(); k++) - { - if (aggregate) - { - res_y[0] += y[k]; - } - else - { - res_y[k] = y[k]; - } - } + copy_vec(y, res_y, is_objective()); } void eval_jac(const V *req_x, V *res_jac) { - for (size_t i = 0; i < indexVars.size(); i++) - { - x[i] = req_x[indexVars[i]]; - } + copy_ptr(req_x, indexVars.data(), x); fun.sparse_jac_rev(x, jac_, jac_pattern_, jac_coloring_, jac_work_); auto &jac = jac_.val(); - for (size_t i = 0; i < jac_.nnz(); i++) - { - res_jac[i] = jac[i]; - } + copy_vec(jac, res_jac); } - void eval_hess(const V *req_x, const V *req_w, V *res_hess, bool aggregate = false) + void eval_hess(const V *req_x, const V *req_w, V *res_hess) { - for (size_t i = 0; i < indexVars.size(); i++) - { - x[i] = req_x[indexVars[i]]; - } - for (size_t k = 0; k < fun.Range(); k++) - { - if (aggregate) - { - w[k] = req_w[0]; - } - else - { - w[k] = req_w[indexCons[k]]; - } - } + copy_ptr(req_x, indexVars.data(), x); + copy_ptr(req_w, indexCons.data(), w, is_objective()); fun.sparse_hes(x, w, hess_, hess_pattern_, hess_coloring_, hess_work_); auto &hess = hess_.val(); - for (size_t i = 0; i < hess_.nnz(); i++) - { - res_hess[i] = hess[i]; - } + copy_vec(hess, res_hess); } CallbackPattern get_callback_pattern() const @@ -283,13 +250,50 @@ struct CallbackEvaluator return pattern; } + + private: + template + static void copy_ptr(const T *src, const I *idx, std::vector &dst, bool duplicate = false) + { + for (size_t i = 0; i < dst.size(); i++) + { + if (duplicate) + { + dst[i] = src[0]; + } + else + { + dst[i] = src[idx[i]]; + } + } + } + + template + static void copy_vec(const std::vector &src, T *dst, bool aggregate = false) + { + if (aggregate) + { + dst[0] = 0.0; + } + for (size_t i = 0; i < src.size(); i++) + { + if (aggregate) + { + dst[0] += src[i]; + } + else + { + dst[i] = src[i]; + } + } + } }; struct Outputs { - std::vector obj_idxs; - std::vector con_idxs; - std::vector cons; + std::vector objective_outputs; + std::vector constraint_outputs; + std::vector constraints; }; inline bool is_name_empty(const char *name) @@ -575,7 +579,7 @@ class KNITROModel : public OnesideLinearConstraintMixin, std::unordered_map m_pending_outputs; std::vector>> m_evaluators; - bool m_need_to_add_callbacks = false; + bool m_has_pending_callbacks = false; int m_solve_status = 0; bool m_is_dirty = true; @@ -596,9 +600,11 @@ class KNITROModel : public OnesideLinearConstraintMixin, void _set_quadratic_objective(const ScalarQuadraticFunction &f); void _reset_objective(); void _add_graph(ExpressionGraph &graph); - void _add_callbacks(); - void _add_constraint_callback(ExpressionGraph *graph, const Outputs &outputs); - void _add_objective_callback(ExpressionGraph *graph, const Outputs &outputs); + void _add_pending_callbacks(); + void _add_callbacks(const ExpressionGraph &graph, const Outputs &outputs); + void _add_callback(const ExpressionGraph &graph, const std::vector &outputs, + const std::vector &constraints); + void _register_callback(CallbackEvaluator *evaluator); void _update(); void _pre_solve(); void _solve(); @@ -668,48 +674,6 @@ class KNITROModel : public OnesideLinearConstraintMixin, m_is_dirty = true; } - template - void _register_callback(CallbackEvaluator *evaluator, const F f, const G g, const H h) - { - CB_context *cb = nullptr; - auto p = evaluator->get_callback_pattern(); - int error; - error = knitro::KN_add_eval_callback(m_kc.get(), p.indexCons.empty(), p.indexCons.size(), - p.indexCons.data(), f, &cb); - _check_error(error); - error = knitro::KN_set_cb_user_params(m_kc.get(), cb, evaluator); - _check_error(error); - error = knitro::KN_set_cb_grad(m_kc.get(), cb, p.objGradIndexVars.size(), - p.objGradIndexVars.data(), p.jacIndexCons.size(), - p.jacIndexCons.data(), p.jacIndexVars.data(), g); - _check_error(error); - error = knitro::KN_set_cb_hess(m_kc.get(), cb, p.hessIndexVars1.size(), - p.hessIndexVars1.data(), p.hessIndexVars2.data(), h); - _check_error(error); - } - - template - void _add_callback_impl(const ExpressionGraph &graph, const std::vector cons, - const T &trace, const F f, const G g, const H h) - { - auto evaluator_ptr = std::make_unique>(); - auto *evaluator = evaluator_ptr.get(); - evaluator->fun = trace(graph); - evaluator->indexVars.resize(graph.n_variables()); - for (size_t i = 0; i < graph.n_variables(); i++) - { - evaluator->indexVars[i] = _variable_index(graph.m_variables[i]); - } - evaluator->indexCons.resize(cons.size()); - for (size_t i = 0; i < cons.size(); i++) - { - evaluator->indexCons[i] = _constraint_index(cons[i]); - } - evaluator->setup(); - _register_callback(evaluator, f, g, h); - m_evaluators.push_back(std::move(evaluator_ptr)); - } - template using Getter = std::function; template @@ -752,12 +716,12 @@ class KNITROModel : public OnesideLinearConstraintMixin, _check_error(error); } - template - std::string _get_name(F get, KNINT index, const char *prefix) const + template + std::string _get_name(F get, K key, const char *prefix) const { char name[1024]; name[0] = '\0'; - int error = get(m_kc.get(), index, 1024, name); + int error = get(m_kc.get(), key, 1024, name); _check_error(error); if (name[0] != '\0') @@ -766,17 +730,10 @@ class KNITROModel : public OnesideLinearConstraintMixin, } else { - return fmt::format("{}{}", prefix, index); + return fmt::format("{}{}", prefix, key); } } - template - void _set_name(F set, KNINT index, const std::string &name) - { - int error = set(m_kc.get(), index, name.c_str()); - _check_error(error); - } - template KNINT _get_index(const I &index) const { diff --git a/lib/cppad_interface.cpp b/lib/cppad_interface.cpp index 7dc93b2..40c97eb 100644 --- a/lib/cppad_interface.cpp +++ b/lib/cppad_interface.cpp @@ -428,7 +428,7 @@ CppAD::AD cppad_trace_expression( } ADFunDouble cppad_trace_graph_constraints(const ExpressionGraph &graph, - const std::vector &selected_outputs) + const std::vector &selected) { ankerl::unordered_dense::map> seen_expressions; @@ -456,7 +456,7 @@ ADFunDouble cppad_trace_graph_constraints(const ExpressionGraph &graph, auto &outputs = graph.m_constraint_outputs; std::vector indices; - if (selected_outputs.empty()) + if (selected.empty()) { indices.reserve(outputs.size()); for (size_t i = 0; i < outputs.size(); i++) @@ -466,7 +466,7 @@ ADFunDouble cppad_trace_graph_constraints(const ExpressionGraph &graph, } else { - indices = selected_outputs; + indices = selected; } auto N_outputs = indices.size(); @@ -486,8 +486,8 @@ ADFunDouble cppad_trace_graph_constraints(const ExpressionGraph &graph, return f; } -ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, bool aggregate, - const std::vector &selected_outputs) +ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, + const std::vector &selected, bool aggregate) { ankerl::unordered_dense::map> seen_expressions; @@ -513,7 +513,7 @@ ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, bool aggre auto &outputs = graph.m_objective_outputs; std::vector indices; - if (selected_outputs.empty()) + if (selected.empty()) { indices.reserve(outputs.size()); for (size_t i = 0; i < outputs.size(); i++) @@ -523,7 +523,7 @@ ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, bool aggre } else { - indices = selected_outputs; + indices = selected; } auto N_outputs = indices.size(); diff --git a/lib/cppad_interface_ext.cpp b/lib/cppad_interface_ext.cpp index 60d64ab..ac554a4 100644 --- a/lib/cppad_interface_ext.cpp +++ b/lib/cppad_interface_ext.cpp @@ -2,6 +2,7 @@ #include #include #include +#include namespace nb = nanobind; @@ -183,8 +184,8 @@ NB_MODULE(cppad_interface_ext, m) .def_ro("hessian", &CppADAutodiffGraph::hessian_graph); m.def("cppad_trace_graph_constraints", cppad_trace_graph_constraints, nb::arg("graph"), - nb::arg("selected_outputs") = std::vector{}); + nb::arg("selected") = std::vector{}); m.def("cppad_trace_graph_objective", cppad_trace_graph_objective, nb::arg("graph"), - nb::arg("aggregate") = true, nb::arg("selected_outputs") = std::vector{}); + nb::arg("selected") = std::vector{}, nb::arg("aggregate") = true); m.def("cppad_autodiff", &cppad_autodiff); } diff --git a/lib/knitro_model.cpp b/lib/knitro_model.cpp index b1bc3c4..1ffc098 100644 --- a/lib/knitro_model.cpp +++ b/lib/knitro_model.cpp @@ -200,7 +200,7 @@ VariableIndex KNITROModel::add_variable(VariableDomain domain, double lb, double if (!is_name_empty(name)) { - _set_name(knitro::KN_set_var_name, indexVar, name); + _set_value(knitro::KN_set_var_name, indexVar, name); } m_n_vars++; @@ -261,7 +261,7 @@ std::string KNITROModel::get_variable_name(const VariableIndex &variable) const void KNITROModel::set_variable_name(const VariableIndex &variable, const std::string &name) { KNINT indexVar = _variable_index(variable); - _set_name(knitro::KN_set_var_name, indexVar, name); + _set_value(knitro::KN_set_var_name, indexVar, name.c_str()); } void KNITROModel::set_variable_domain(const VariableIndex &variable, VariableDomain domain) @@ -389,10 +389,10 @@ ConstraintIndex KNITROModel::add_single_nl_constraint(ExpressionGraph &graph, _add_graph(graph); graph.add_constraint_output(result); size_t i = graph.m_constraint_outputs.size() - 1; - m_pending_outputs[&graph].con_idxs.push_back(i); + m_pending_outputs[&graph].constraint_outputs.push_back(i); auto setter = [this, &graph](const ConstraintIndex &constraint) { - m_pending_outputs[&graph].cons.push_back(constraint); - m_need_to_add_callbacks = true; + m_pending_outputs[&graph].constraints.push_back(constraint); + m_has_pending_callbacks = true; }; return _add_constraint_impl(ConstraintType::NL, interval, name, setter); } @@ -551,7 +551,7 @@ void KNITROModel::delete_constraint(const ConstraintIndex &constraint) void KNITROModel::set_constraint_name(const ConstraintIndex &constraint, const std::string &name) { KNINT indexCon = _constraint_index(constraint); - _set_name(knitro::KN_set_con_name, indexCon, name); + _set_value(knitro::KN_set_con_name, indexCon, name.c_str()); } std::string KNITROModel::get_constraint_name(const ConstraintIndex &constraint) const @@ -726,8 +726,8 @@ void KNITROModel::add_single_nl_objective(ExpressionGraph &graph, const Expressi _add_graph(graph); graph.add_objective_output(result); size_t i = graph.m_objective_outputs.size() - 1; - m_pending_outputs[&graph].obj_idxs.push_back(i); - m_need_to_add_callbacks = true; + m_pending_outputs[&graph].objective_outputs.push_back(i); + m_has_pending_callbacks = true; m_obj_flag |= OBJ_NONLINEAR; _mark_dirty(); } @@ -775,7 +775,7 @@ void KNITROModel::_reset_objective() _check_error(error); for (auto &[graph, outputs] : m_pending_outputs) { - outputs.obj_idxs.clear(); + outputs.objective_outputs.clear(); } } m_obj_flag = 0; @@ -836,93 +836,135 @@ double KNITROModel::get_obj_value() const return _get_value(knitro::KN_get_obj_value); } -void KNITROModel::_add_constraint_callback(ExpressionGraph *graph, const Outputs &outputs) +void KNITROModel::_register_callback(CallbackEvaluator *evaluator) { - auto f = [](KN_context *, CB_context *, KN_eval_request *req, KN_eval_result *res, + auto f = [](KN_context *, CB_context *cb, KN_eval_request *req, KN_eval_result *res, void *data) -> int { - auto *evaluator = static_cast *>(data); - evaluator->eval_fun(req->x, res->c); + auto evaluator = static_cast *>(data); + if (evaluator->is_objective()) + { + evaluator->eval_fun(req->x, res->obj); + } + else + { + evaluator->eval_fun(req->x, res->c); + } return 0; }; - auto g = [](KN_context *, CB_context *, KN_eval_request *req, KN_eval_result *res, + + auto g = [](KN_context *, CB_context *cb, KN_eval_request *req, KN_eval_result *res, void *data) -> int { - auto *evaluator = static_cast *>(data); - evaluator->eval_jac(req->x, res->jac); + auto evaluator = static_cast *>(data); + if (evaluator->is_objective()) + { + evaluator->eval_jac(req->x, res->objGrad); + } + else + { + evaluator->eval_jac(req->x, res->jac); + } return 0; }; - auto h = [](KN_context *, CB_context *, KN_eval_request *req, KN_eval_result *res, + + auto h = [](KN_context *, CB_context *cb, KN_eval_request *req, KN_eval_result *res, void *data) -> int { - auto *evaluator = static_cast *>(data); - evaluator->eval_hess(req->x, req->lambda, res->hess); + auto evaluator = static_cast *>(data); + if (evaluator->is_objective()) + { + evaluator->eval_hess(req->x, req->sigma, res->hess); + } + else + { + evaluator->eval_hess(req->x, req->lambda, res->hess); + } return 0; }; - auto trace = [outputs](const ExpressionGraph &graph) { - return cppad_trace_graph_constraints(graph, outputs.con_idxs); - }; - _add_callback_impl(*graph, outputs.cons, trace, f, g, h); + + CB_context *cb = nullptr; + auto p = evaluator->get_callback_pattern(); + int error; + error = knitro::KN_add_eval_callback(m_kc.get(), p.indexCons.empty(), p.indexCons.size(), + p.indexCons.data(), f, &cb); + _check_error(error); + error = knitro::KN_set_cb_user_params(m_kc.get(), cb, evaluator); + _check_error(error); + error = knitro::KN_set_cb_grad(m_kc.get(), cb, p.objGradIndexVars.size(), + p.objGradIndexVars.data(), p.jacIndexCons.size(), + p.jacIndexCons.data(), p.jacIndexVars.data(), g); + _check_error(error); + error = knitro::KN_set_cb_hess(m_kc.get(), cb, p.hessIndexVars1.size(), p.hessIndexVars1.data(), + p.hessIndexVars2.data(), h); + _check_error(error); } -void KNITROModel::_add_objective_callback(ExpressionGraph *graph, const Outputs &outputs) +void KNITROModel::_add_callback(const ExpressionGraph &graph, const std::vector &outputs, + const std::vector &constraints) { - auto f = [](KN_context *, CB_context *, KN_eval_request *req, KN_eval_result *res, - void *data) -> int { - auto *evaluator = static_cast *>(data); - res->obj[0] = 0.0; - evaluator->eval_fun(req->x, res->obj, true); - return 0; - }; - auto g = [](KN_context *, CB_context *, KN_eval_request *req, KN_eval_result *res, - void *data) -> int { - auto *evaluator = static_cast *>(data); - evaluator->eval_jac(req->x, res->objGrad); - return 0; - }; - auto h = [](KN_context *, CB_context *, KN_eval_request *req, KN_eval_result *res, - void *data) -> int { - auto *evaluator = static_cast *>(data); - evaluator->eval_hess(req->x, req->sigma, res->hess, true); - return 0; - }; - auto trace = [outputs](const ExpressionGraph &graph) { - return cppad_trace_graph_objective(graph, true, outputs.obj_idxs); - }; - _add_callback_impl(*graph, {}, trace, f, g, h); + auto evaluator_ptr = std::make_unique>(); + auto *evaluator = evaluator_ptr.get(); + evaluator->indexVars.resize(graph.n_variables()); + for (size_t i = 0; i < graph.n_variables(); i++) + { + evaluator->indexVars[i] = _variable_index(graph.m_variables[i]); + } + evaluator->indexCons.resize(constraints.size()); + for (size_t i = 0; i < constraints.size(); i++) + { + evaluator->indexCons[i] = _constraint_index(constraints[i]); + } + if (evaluator->is_objective()) + { + evaluator->fun = cppad_trace_graph_objective(graph, outputs, true); + } + else + { + evaluator->fun = cppad_trace_graph_constraints(graph, outputs); + } + evaluator->setup(); + _register_callback(evaluator); + m_evaluators.push_back(std::move(evaluator_ptr)); +} + +void KNITROModel::_add_callbacks(const ExpressionGraph &graph, const Outputs &outputs) +{ + if (graph.has_constraint_output() && !outputs.constraint_outputs.empty()) + { + _add_callback(graph, outputs.constraint_outputs, outputs.constraints); + } + + if (graph.has_objective_output() && !outputs.objective_outputs.empty()) + { + _add_callback(graph, outputs.objective_outputs, {}); + } } -void KNITROModel::_add_callbacks() +void KNITROModel::_add_pending_callbacks() { - if (!m_need_to_add_callbacks) + if (!m_has_pending_callbacks) { return; } for (const auto &[graph, outputs] : m_pending_outputs) { - if (graph->has_constraint_output() && !outputs.con_idxs.empty()) - { - _add_constraint_callback(graph, outputs); - } - - if (graph->has_objective_output() && !outputs.obj_idxs.empty()) - { - _add_objective_callback(graph, outputs); - } + _add_callbacks(*graph, outputs); } + m_pending_outputs.clear(); - m_need_to_add_callbacks = false; + m_has_pending_callbacks = false; } // Solve functions void KNITROModel::_update() { - _add_callbacks(); + _add_pending_callbacks(); int error = knitro::KN_update(m_kc.get()); _check_error(error); } void KNITROModel::_pre_solve() { - _add_callbacks(); + _add_pending_callbacks(); } void KNITROModel::_solve() @@ -1045,7 +1087,7 @@ void KNITROModel::_reset_state() m_obj_flag = 0; m_pending_outputs.clear(); m_evaluators.clear(); - m_need_to_add_callbacks = false; + m_has_pending_callbacks = false; _mark_dirty(); m_solve_status = 0; }