Refactor cppad_trace_graph functions to use 'selected' parameter for output indices

Author: yemine0x01

include/pyoptinterface/cppad_interface.hpp

@@ -32,11 +32,12 @@ ADFunDouble sparse_hessian(const ADFunDouble &f, const sparsity_pattern_t &patte
                            const std::vector<double> &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<size_t> &selected_outputs = {});
-ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, bool aggregate = true,
-                                        const std::vector<size_t> &selected_outputs = {});
+                                          const std::vector<size_t> &selected = {});
+ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph,
+                                        const std::vector<size_t> &selected = {},
+                                        bool aggregate = true);
 
 struct CppADAutodiffGraph
 {

lib/cppad_interface.cpp

@@ -428,7 +428,7 @@ CppAD::AD<double> cppad_trace_expression(
 }
 
 ADFunDouble cppad_trace_graph_constraints(const ExpressionGraph &graph,
-                                          const std::vector<size_t> &selected_outputs)
+                                          const std::vector<size_t> &selected)
 {
 	ankerl::unordered_dense::map<ExpressionHandle, CppAD::AD<double>> seen_expressions;
 
@@ -456,7 +456,7 @@ ADFunDouble cppad_trace_graph_constraints(const ExpressionGraph &graph,
 	auto &outputs = graph.m_constraint_outputs;
 
 	std::vector<size_t> 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<size_t> &selected_outputs)
+ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph,
+                                        const std::vector<size_t> &selected, bool aggregate)
 {
 	ankerl::unordered_dense::map<ExpressionHandle, CppAD::AD<double>> seen_expressions;
 
@@ -513,7 +513,7 @@ ADFunDouble cppad_trace_graph_objective(const ExpressionGraph &graph, bool aggre
 	auto &outputs = graph.m_objective_outputs;
 
 	std::vector<size_t> 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();

lib/cppad_interface_ext.cpp

@@ -2,6 +2,7 @@
 #include <nanobind/make_iterator.h>
 #include <nanobind/stl/vector.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/optional.h>
 
 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<size_t>{});
+	      nb::arg("selected") = std::vector<size_t>{});
 	m.def("cppad_trace_graph_objective", cppad_trace_graph_objective, nb::arg("graph"),
-	      nb::arg("aggregate") = true, nb::arg("selected_outputs") = std::vector<size_t>{});
+	      nb::arg("selected") = std::vector<size_t>{}, nb::arg("aggregate") = true);
 	m.def("cppad_autodiff", &cppad_autodiff);
 }

lib/knitro_model.cpp

@@ -914,7 +914,7 @@ void KNITROModel::_add_callback(const ExpressionGraph &graph, const std::vector<
 	}
 	if (evaluator->is_objective())
 	{
-		evaluator->fun = cppad_trace_graph_objective(graph, true, outputs);
+		evaluator->fun = cppad_trace_graph_objective(graph, outputs, true);
 	}
 	else
 	{