Abstraction of NonlinearEvaluator

Some new codes for cache_model

Author: metab0t

CMakeLists.txt

@@ -41,10 +41,12 @@ set(POI_INSTALL_DIR ${SKBUILD_PLATLIB_DIR}/pyoptinterface/_src)
 
 add_library(core STATIC)
 target_sources(core PRIVATE
+  include/pyoptinterface/cache_model.hpp
   include/pyoptinterface/core.hpp
   include/pyoptinterface/container.hpp
   include/pyoptinterface/dylib.hpp
   include/pyoptinterface/solver_common.hpp
+  lib/cache_model.cpp
   lib/core.cpp
 )
 target_include_directories(core PUBLIC include thirdparty)
@@ -63,7 +65,7 @@ target_sources(nleval PRIVATE
   include/pyoptinterface/nleval.hpp
   lib/nleval.cpp
 )
-target_link_libraries(nleval PUBLIC core)
+target_link_libraries(nleval PUBLIC nlexpr core)
 
 add_library(cppad_interface STATIC)
 target_sources(cppad_interface PRIVATE

include/pyoptinterface/cache_model.hpp

@@ -0,0 +1,57 @@
+#pragma once
+
+#include <concepts>
+#include <vector>
+#include <span>
+
+// This file defines some common utilities to store the optimization model in a compact way
+
+template <std::integral ColumnIndexT, std::integral VariableIndexT,
+          std::floating_point CoefficientT>
+struct LinearExpressionCache
+{
+	std::vector<ColumnIndexT> column_ptr = {0};
+	std::vector<VariableIndexT> variables;
+	std::vector<CoefficientT> coefficients;
+
+	template <std::integral IT, std::floating_point FT>
+	void add_row(std::span<const IT> row_variables, std::span<const FT> row_coefficients)
+	{
+		variables.insert(variables.end(), row_variables.begin(), row_variables.end());
+		coefficients.insert(coefficients.end(), row_coefficients.begin(), row_coefficients.end());
+		column_ptr.push_back(variables.size());
+	}
+};
+
+template <std::integral ColumnIndexT, std::integral VariableIndexT,
+          std::floating_point CoefficientT>
+struct QuadraticExpressionCache
+{
+	std::vector<ColumnIndexT> column_ptr = {0};
+	std::vector<VariableIndexT> variable_1s;
+	std::vector<VariableIndexT> variable_2s;
+	std::vector<CoefficientT> coefficients;
+
+	std::vector<ColumnIndexT> lin_column_ptr = {0};
+	std::vector<VariableIndexT> lin_variables;
+	std::vector<CoefficientT> lin_coefficients;
+
+	template <std::integral IT, std::floating_point FT>
+	void add_row(std::span<const IT> row_variable_1s, std::span<const IT> row_variable_2s,
+	             std::span<const FT> row_quadratic_coefficients,
+	             std::span<const IT> row_lin_variables, std::span<const FT> row_lin_coefficients)
+	{
+		variable_1s.insert(variable_1s.end(), row_variable_1s.begin(), row_variable_1s.end());
+		variable_2s.insert(variable_2s.end(), row_variable_2s.begin(), row_variable_2s.end());
+		coefficients.insert(coefficients.end(), row_quadratic_coefficients.begin(),
+		                    row_quadratic_coefficients.end());
+		column_ptr.push_back(variable_1s.size());
+
+		lin_variables.insert(lin_variables.end(), row_lin_variables.begin(),
+		                     row_lin_variables.end());
+		lin_coefficients.insert(lin_coefficients.end(), row_lin_coefficients.begin(),
+		                        row_lin_coefficients.end());
+		lin_column_ptr.push_back(lin_variables.size());
+	}
+};
+

include/pyoptinterface/ipopt_model.hpp

@@ -152,20 +152,20 @@ struct IpoptModel : public OnesideLinearConstraintMixin<IpoptModel>,
 	} nl_objective_groups;
 
 	int add_graph_index();
-	void record_graph_hash(size_t graph_index, const ExpressionGraph &graph);
-	int aggregate_graph_constraint_groups();
-	int get_graph_constraint_group_representative(int group_index) const;
-	int aggregate_graph_objective_groups();
-	int get_graph_objective_group_representative(int group_index) const;
-
-	void assign_constraint_group_autodiff_structure(int group_index,
-	                                                const AutodiffSymbolicStructure &structure);
-	void assign_constraint_group_autodiff_evaluator(int group_index,
-	                                                const ConstraintAutodiffEvaluator &evaluator);
-	void assign_objective_group_autodiff_structure(int group_index,
-	                                               const AutodiffSymbolicStructure &structure);
-	void assign_objective_group_autodiff_evaluator(int group_index,
-	                                               const ObjectiveAutodiffEvaluator &evaluator);
+	void finalize_graph_instance(size_t graph_index, const ExpressionGraph &graph);
+	int aggregate_nl_constraint_groups();
+	int get_nl_constraint_group_representative(int group_index) const;
+	int aggregate_nl_objective_groups();
+	int get_nl_objective_group_representative(int group_index) const;
+
+	void assign_nl_constraint_group_autodiff_structure(int group_index,
+	                                                   const AutodiffSymbolicStructure &structure);
+	void assign_nl_constraint_group_autodiff_evaluator(
+	    int group_index, const ConstraintAutodiffEvaluator &evaluator);
+	void assign_nl_objective_group_autodiff_structure(int group_index,
+	                                                  const AutodiffSymbolicStructure &structure);
+	void assign_nl_objective_group_autodiff_evaluator(int group_index,
+	                                                  const ObjectiveAutodiffEvaluator &evaluator);
 
 	ConstraintIndex add_single_nl_constraint(size_t graph_index, const ExpressionGraph &graph,
 	                                         double lb, double ub);
@@ -193,13 +193,18 @@ struct IpoptModel : public OnesideLinearConstraintMixin<IpoptModel>,
 	 * constraint) to the reordered one (linear, quadratic, NL group 0 -> con0, con1 ,..., conN0, NL
 	 * group1 -> con0, con1,..., conN1)
 	 */
-	// these two vectors are maintained when adding NL constraint
-	// which graph instance this constraint belongs to
-	std::vector<int> nl_constraint_graph_instance_indices;
-	// the order of this constraint in the graph instance
-	std::vector<int> nl_constraint_graph_instance_orders;
+	// this is maintained when adding NL constraint
+	struct ConstraintGraphMembership
+	{
+		// which graph it belongs to
+		int graph;
+		// the rank in that graph
+		int rank;
+	};
+	// record the graph a nonlinear constraint belongs to
+	std::vector<ConstraintGraphMembership> nl_constraint_graph_memberships;
 
-	// these two vectors are constructed before optimization
+	// this vector is constructed before optimization
 	// ext means the external monotonic order
 	// int means the internal order that passes to Ipopt
 	std::vector<int> nl_constraint_map_ext2int;
@@ -227,6 +232,8 @@ struct IpoptModel : public OnesideLinearConstraintMixin<IpoptModel>,
 	std::optional<LinearEvaluator> m_linear_obj_evaluator;
 	std::optional<QuadraticEvaluator> m_quadratic_obj_evaluator;
 
+	NonlinearEvaluator m_nl_evaluator;
+
 	// The options of the Ipopt solver, we cache them before constructing the m_problem
 	Hashmap<std::string, int> m_options_int;
 	Hashmap<std::string, double> m_options_num;

include/pyoptinterface/nleval.hpp

@@ -3,7 +3,8 @@
 #include <cstdint>
 #include <vector>
 
-#include "core.hpp"
+#include "pyoptinterface/core.hpp"
+#include "pyoptinterface/nlexpr.hpp"
 
 enum class HessianSparsityType
 {
@@ -157,4 +158,120 @@ struct QuadraticEvaluator
 	                               Hashmap<std::tuple<int, int>, int> &hessian_index_map,
 	                               HessianSparsityType hessian_type);
 	void eval_lagrangian_hessian(const double *restrict lambda, double *restrict hessian) const;
+};
+
+struct NonlinearEvaluator
+{
+	// How many graph instances are there
+	size_t n_graph_instances = 0;
+	// record the inputs of graph instances
+	struct GraphInput
+	{
+		std::vector<int> variables;
+		std::vector<double> constants;
+	};
+	std::vector<GraphInput> graph_inputs;
+	// record graph instances with constraint output and objective output
+	struct GraphHash
+	{
+		// hash of this graph instance
+		uint64_t hash;
+		// index of this graph instance
+		int index;
+	};
+	struct GraphHashes
+	{
+		std::vector<GraphHash> hashes;
+		size_t n_hashes_since_last_aggregation;
+	} constraint_graph_hashes, objective_graph_hashes;
+
+	// length = n_graph_instances
+	// record which group this graph instance belongs to
+	struct GraphGroupMembership
+	{
+		// which group it belongs to
+		int group;
+		// the rank in that group
+		int rank;
+	};
+	std::vector<GraphGroupMembership> constraint_group_memberships, objective_group_memberships;
+	// record which index of constraint this graph starts
+	std::vector<int> constraint_indices_offsets;
+
+	// graph groups
+	struct ConstraintGraphGroup
+	{
+		std::vector<int> instance_indices;
+		AutodiffSymbolicStructure autodiff_structure;
+		ConstraintAutodiffEvaluator autodiff_evaluator;
+
+		// where to store the hessian matrix
+		// length = instance_indices.size() * hessian_nnz
+		std::vector<int> hessian_indices;
+	};
+	std::vector<ConstraintGraphGroup> constraint_groups;
+	Hashmap<uint64_t, int> hash_to_constraint_group;
+
+	struct ObjectiveGraphGroup
+	{
+		std::vector<int> instance_indices;
+		AutodiffSymbolicStructure autodiff_structure;
+		ObjectiveAutodiffEvaluator autodiff_evaluator;
+		// where to store the gradient vector
+		// length = instance_indices.size() * jacobian_nnz
+		std::vector<int> gradient_indices;
+		// where to store the hessian matrix
+		// length = instance_indices.size() * hessian_nnz
+		std::vector<int> hessian_indices;
+	};
+	std::vector<ObjectiveGraphGroup> objective_groups;
+	Hashmap<uint64_t, int> hash_to_objective_group;
+
+	int add_graph_instance();
+	void finalize_graph_instance(size_t graph_index, const ExpressionGraph &graph);
+	int aggregate_constraint_groups();
+	int get_constraint_group_representative(int group_index) const;
+	int aggregate_objective_groups();
+	int get_objective_group_representative(int group_index) const;
+
+	void assign_constraint_group_autodiff_structure(int group_index,
+	                                                const AutodiffSymbolicStructure &structure);
+	void assign_constraint_group_autodiff_evaluator(int group_index,
+	                                                const ConstraintAutodiffEvaluator &evaluator);
+	void assign_objective_group_autodiff_structure(int group_index,
+	                                               const AutodiffSymbolicStructure &structure);
+	void assign_objective_group_autodiff_evaluator(int group_index,
+	                                               const ObjectiveAutodiffEvaluator &evaluator);
+
+	void calculate_constraint_graph_instances_offset();
+
+	// functions to evaluate the nonlinear constraints and objectives
+
+	// f
+	void eval_constraints(const double *restrict x, double *restrict f) const;
+	double eval_objective(const double *restrict x) const;
+
+	// first order derivative
+	void analyze_constraints_jacobian_structure(size_t row_base, size_t &global_jacobian_nnz,
+	                                            std::vector<int> &global_jacobian_rows,
+	                                            std::vector<int> &global_jacobian_cols);
+	void analyze_objective_gradient_structure(std::vector<int> &global_gradient_cols,
+	                                          Hashmap<int, int> &sparse_gradient_map);
+
+	void eval_constraints_jacobian(const double *restrict x, double *restrict jacobian) const;
+	void eval_objective_gradient(const double *restrict x, double *restrict grad_f) const;
+
+	// second order derivative
+	void analyze_constraints_hessian_structure(
+	    size_t &global_hessian_nnz, std::vector<int> &global_hessian_rows,
+	    std::vector<int> &global_hessian_cols,
+	    Hashmap<std::tuple<int, int>, int> &hessian_index_map, HessianSparsityType hessian_type);
+	void analyze_objective_hessian_structure(size_t &global_hessian_nnz,
+	                                         std::vector<int> &global_hessian_rows,
+	                                         std::vector<int> &global_hessian_cols,
+	                                         Hashmap<std::tuple<int, int>, int> &hessian_index_map,
+	                                         HessianSparsityType hessian_type);
+
+	void eval_lagrangian_hessian(const double *restrict x, const double *restrict lambda,
+	                             const double sigma, double *restrict hessian) const;
 };