Implement more nonlinear operators

Author: metab0t

lib/ipopt_model_ext.cpp

@@ -24,6 +24,7 @@ NB_MODULE(ipopt_model_ext, m)
 	    .value("Restoration_Failed", ApplicationReturnStatus::Restoration_Failed)
 	    .value("Error_In_Step_Computation", ApplicationReturnStatus::Error_In_Step_Computation)
 	    .value("Maximum_CpuTime_Exceeded", ApplicationReturnStatus::Maximum_CpuTime_Exceeded)
+	    .value("Maximum_WallTime_Exceeded", ApplicationReturnStatus::Maximum_WallTime_Exceeded)
 	    .value("Not_Enough_Degrees_Of_Freedom",
 	           ApplicationReturnStatus::Not_Enough_Degrees_Of_Freedom)
 	    .value("Invalid_Problem_Definition", ApplicationReturnStatus::Invalid_Problem_Definition)

lib/main.cpp

@@ -318,19 +318,28 @@ auto test_highs() -> void
 	model.optimize();
 }
 
-auto test_highs_capi() -> void
+auto debug_highs_passname(int N) -> void
 {
 	void *highs = highs::Highs_create();
+	
+	auto start = std::chrono::high_resolution_clock::now();
+
+	HighsModelMixin model;
+	for (auto i = 0; i < N; i++)
+	{
+		highs::Highs_addCol(highs, 0.0, 0.0, 0.0, 0, nullptr, nullptr);
 
-	double c[] = {0.0, 1.0, 0.0};
-	double lower[] = {0, 0, 0};
-	double upper[] = {1.0, 1.0, 1.0};
+		auto col = highs::Highs_getNumCol(highs);
 
-	// highs::Highs_addCols(highs, 3, c, lower, upper, 0, nullptr, nullptr, nullptr);
+		assert(col == i + 1);
 
-	auto hessian_nz = highs::Highs_getHessianNumNz(highs);
-	highs::Highs_run(highs);
-	hessian_nz = highs::Highs_getHessianNumNz(highs);
+		auto name = fmt::format("x{}", i);
+		highs::Highs_passColName(highs, i, name.c_str());
+	}
+
+	auto end = std::chrono::high_resolution_clock::now();
+	auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
+	fmt::print("N={}, time={} milliseconds\n", N, duration.count());
 }
 
 void debug_highs(int N)
@@ -350,16 +359,18 @@ void debug_highs(int N)
 
 void test_highs_add_variable()
 {
-	std::vector<int> Ns{100, 1000, 5000, 10000};
+	std::vector<int> Ns{1000, 10000, 100000, 1000000};
 
 	for (const auto& N : Ns)
 	{
-		debug_highs(N);
+		debug_highs_passname(N);
 	}
 }
 
 auto main() -> int
 {
+	highs::load_library("E:\\HiGHS\\install\\bin\\highs.dll");
+
 	test_highs_add_variable();
 	return 0;
 }

lib/nlcore_ext.cpp

@@ -8,6 +8,7 @@
 namespace nb = nanobind;
 
 #include "pyoptinterface/nlcore.hpp"
+#include "cppad/utility/pow_int.hpp"
 
 using a_double = CppAD::AD<double>;
 using advec = std::vector<a_double>;
@@ -45,12 +46,6 @@ NB_MODULE(nlcore_ext, m)
 	    .def(nb::self *= double())
 	    .def(nb::self /= double());
 
-	m.def("sqrt", [](const a_double &x) { return CppAD::sqrt(x); });
-	m.def("sin", [](const a_double &x) { return CppAD::sin(x); });
-	m.def("cos", [](const a_double &x) { return CppAD::cos(x); });
-	m.def("exp", [](const a_double &x) { return CppAD::exp(x); });
-	m.def("log", [](const a_double &x) { return CppAD::log(x); });
-
 	nb::bind_vector<advec, nb::rv_policy::reference_internal>(m, "advec");
 
 	m.def("Independent", nb::overload_cast<advec &>(&CppAD::Independent<advec>), nb::arg("x"));
@@ -129,6 +124,31 @@ NB_MODULE(nlcore_ext, m)
 	m.def("sparse_jacobian", &sparse_jacobian<double>);
 	m.def("sparse_hessian", &sparse_hessian<double>);
 
+	m.def("abs", [](const a_double &x) { return CppAD::abs(x); });
+	m.def("acos", [](const a_double &x) { return CppAD::acos(x); });
+	m.def("acosh", [](const a_double &x) { return CppAD::acosh(x); });
+	m.def("asin", [](const a_double &x) { return CppAD::asin(x); });
+	m.def("asinh", [](const a_double &x) { return CppAD::asinh(x); });
+	m.def("atan", [](const a_double &x) { return CppAD::atan(x); });
+	m.def("atanh", [](const a_double &x) { return CppAD::atanh(x); });
+	m.def("cos", [](const a_double &x) { return CppAD::cos(x); });
+	m.def("cosh", [](const a_double &x) { return CppAD::cosh(x); });
+	m.def("erf", [](const a_double &x) { return CppAD::erf(x); });
+	m.def("erfc", [](const a_double &x) { return CppAD::erfc(x); });
+	m.def("exp", [](const a_double &x) { return CppAD::exp(x); });
+	m.def("expm1", [](const a_double &x) { return CppAD::expm1(x); });
+	m.def("log1p", [](const a_double &x) { return CppAD::log1p(x); });
+	m.def("log", [](const a_double &x) { return CppAD::log(x); });
+	m.def("pow", [](const a_double &x, const int &y) { return CppAD::pow(x, y); });
+	m.def("pow", [](const a_double &x, const double &y) { return CppAD::pow(x, y); });
+	m.def("pow", [](const double &x, const a_double &y) { return CppAD::pow(x, y); });
+	m.def("pow", [](const a_double &x, const a_double &y) { return CppAD::pow(x, y); });
+	m.def("sin", [](const a_double &x) { return CppAD::sin(x); });
+	m.def("sinh", [](const a_double &x) { return CppAD::sinh(x); });
+	m.def("sqrt", [](const a_double &x) { return CppAD::sqrt(x); });
+	m.def("tan", [](const a_double &x) { return CppAD::tan(x); });
+	m.def("tanh", [](const a_double &x) { return CppAD::tanh(x); });
+
 	nb::enum_<graph_op_enum>(m, "graph_op")
 	    .value("abs", graph_op_enum::abs_graph_op)
 	    .value("acos", graph_op_enum::acos_graph_op)

src/pyoptinterface/__init__.py

@@ -26,7 +26,29 @@
 
 from pyoptinterface._src.aml import make_nd_variable, quicksum, quicksum_
 
-from pyoptinterface._src.nlcore_ext import sqrt, sin, cos, exp, log
+from pyoptinterface._src.nlcore_ext import (
+    abs,
+    acos,
+    acosh,
+    asin,
+    asinh,
+    atan,
+    atanh,
+    cos,
+    cosh,
+    erf,
+    erfc,
+    exp,
+    expm1,
+    log1p,
+    log,
+    pow,
+    sin,
+    sinh,
+    sqrt,
+    tan,
+    tanh,
+)
 
 # Alias of ConstraintSense
 Eq = ConstraintSense.Equal

src/pyoptinterface/_src/codegen_c.py

@@ -6,15 +6,31 @@
 from typing import IO
 
 op2name = {
+    graph_op.abs: "fabs",
+    graph_op.acos: "acos",
+    graph_op.acosh: "acosh",
+    graph_op.asin: "asin",
+    graph_op.asinh: "asinh",
+    graph_op.atan: "atan",
+    graph_op.atanh: "atanh",
+    graph_op.cos: "cos",
+    graph_op.cosh: "cosh",
+    graph_op.erf: "erf",
+    graph_op.erfc: "erfc",
+    graph_op.exp: "exp",
+    graph_op.expm1: "expm1",
+    graph_op.log1p: "log1p",
+    graph_op.log: "log",
+    graph_op.pow: "pow",
+    graph_op.sin: "sin",
+    graph_op.sinh: "sinh",
+    graph_op.sqrt: "sqrt",
+    graph_op.tan: "tan",
+    graph_op.tanh: "tanh",
     graph_op.add: "+",
     graph_op.sub: "-",
     graph_op.mul: "*",
     graph_op.div: "/",
-    graph_op.sqrt: "sqrt",
-    graph_op.sin: "sin",
-    graph_op.cos: "cos",
-    graph_op.exp: "exp",
-    graph_op.log: "log",
     graph_op.azmul: "*",
     graph_op.neg: "-",
 }

src/pyoptinterface/_src/codegen_llvm.py

@@ -205,17 +205,46 @@ def create_indirect_load_store(module: ir.Module):
     builder.ret_void()
 
 
+op2name = {
+    graph_op.abs: "fabs",
+    graph_op.acos: "acos",
+    graph_op.acosh: "acosh",
+    graph_op.asin: "asin",
+    graph_op.asinh: "asinh",
+    graph_op.atan: "atan",
+    graph_op.atanh: "atanh",
+    graph_op.cos: "cos",
+    graph_op.cosh: "cosh",
+    graph_op.erf: "erf",
+    graph_op.erfc: "erfc",
+    graph_op.exp: "exp",
+    graph_op.expm1: "expm1",
+    graph_op.log1p: "log1p",
+    graph_op.log: "log",
+    graph_op.pow: "pow",
+    graph_op.sin: "sin",
+    graph_op.sinh: "sinh",
+    graph_op.sqrt: "sqrt",
+    graph_op.tan: "tan",
+    graph_op.tanh: "tanh"
+}
+
+math_ops = set(op2name.keys())
+
+binary_ops = set([graph_op.pow])
+
 def create_llvmir_basic_functions(module: ir.Module):
     create_azmul(module)
     create_sign(module)
     create_direct_load_store(module)
     create_indirect_load_store(module)
 
-    sqrt = ir.Function(module, ir.FunctionType(D, [D]), name="sqrt")
-    sin = ir.Function(module, ir.FunctionType(D, [D]), name="sin")
-    cos = ir.Function(module, ir.FunctionType(D, [D]), name="cos")
-    exp = ir.Function(module, ir.FunctionType(D, [D]), name="exp")
-    log = ir.Function(module, ir.FunctionType(D, [D]), name="log")
+    for (op, op_name) in op2name.items():
+        if op in binary_ops:
+            func_type = ir.FunctionType(D, [D, D])
+        else:
+            func_type = ir.FunctionType(D, [D])
+        ir.Function(module, func_type, name=op_name)
 
     # sin = module.declare_intrinsic('llvm.sin', [D])
     # cos = module.declare_intrinsic('llvm.cos', [D])
@@ -319,11 +348,13 @@ def generate_llvmir_from_graph(
 
     # sin = module.get_global("llvm.sin.f64")
     # cos = module.get_global("llvm.cos.f64")
-    sqrt = module.get_global("sqrt")
-    sin = module.get_global("sin")
-    cos = module.get_global("cos")
-    exp = module.get_global("exp")
-    log = module.get_global("log")
+    math_functions = dict()
+    for op_name in op2name.values():
+        op_function = module.get_global(op_name)
+        if op_function is None:
+            raise ValueError(f"Math function {op_name} not found in module")
+        math_functions[op_name] = op_function
+
     azmul = module.get_global("azmul")
     sign = module.get_global("sign")
     load_direct = module.get_global("load_direct")
@@ -403,7 +434,7 @@ def get_node_value(node: int):
         return val
 
     for iter in graph_obj:
-        op_enum = iter.op_enum
+        op = iter.op_enum
         n_result = iter.n_result
         arg_node = iter.arg_node
 
@@ -414,31 +445,30 @@ def get_node_value(node: int):
         if len(arg_node) == 2:
             arg2 = get_node_value(arg_node[1])
 
-        if op_enum == graph_op.add:
+        if op == graph_op.add:
             ret_val = builder.fadd(arg1, arg2)
-        elif op_enum == graph_op.sub:
+        elif op == graph_op.sub:
             ret_val = builder.fsub(arg1, arg2)
-        elif op_enum == graph_op.mul:
+        elif op == graph_op.mul:
             ret_val = builder.fmul(arg1, arg2)
-        elif op_enum == graph_op.div:
+        elif op == graph_op.div:
             ret_val = builder.fdiv(arg1, arg2)
-        elif op_enum == graph_op.sqrt:
-            ret_val = builder.call(sqrt, [arg1])
-        elif op_enum == graph_op.sin:
-            ret_val = builder.call(sin, [arg1])
-        elif op_enum == graph_op.cos:
-            ret_val = builder.call(cos, [arg1])
-        elif op_enum == graph_op.exp:
-            ret_val = builder.call(exp, [arg1])
-        elif op_enum == graph_op.log:
-            ret_val = builder.call(log, [arg1])
-        elif op_enum == graph_op.azmul:
+        elif op == graph_op.azmul:
             ret_val = builder.fmul(arg1, arg2)
             # ret_val = builder.call(azmul, [arg1, arg2])
-        elif op_enum == graph_op.neg:
+        elif op == graph_op.neg:
             ret_val = builder.fneg(arg1)
+        elif op == graph_op.sign:
+            ret_val = builder.call(sign, [arg1])
+        elif op in math_ops:
+            op_name = op2name[op]
+            op_function = math_functions[op_name]
+            if op in binary_ops:
+                ret_val = builder.call(op_function, [arg1, arg2])
+            else:
+                ret_val = builder.call(op_function, [arg1])
         else:
-            raise ValueError(f"Unknown op_enum: {op_enum}")
+            raise ValueError(f"Unknown op_enum: {op}")
 
         ret_val.name = f"v[{result_node}]"
         v_dict[result_node] = ret_val

src/pyoptinterface/_src/ipopt.py

@@ -10,6 +10,9 @@
 from .jit_llvm import LLJITCompiler
 from .tracefun import trace_adfun
 
+from .core_ext import ConstraintIndex
+from .nlcore_ext import NLConstraintIndex
+
 from .attributes import (
     VariableAttribute,
     ConstraintAttribute,
@@ -282,13 +285,33 @@ def get_terminationstatus(model):
     ),
 }
 
+
+def get_constraint_primal(model, constraint):
+    if isinstance(constraint, ConstraintIndex):
+        return model.get_constraint_primal(constraint.index)
+    elif isinstance(constraint, NLConstraintIndex):
+        index = constraint.index
+        dim = constraint.dim
+        values = [model.get_constraint_primal(index + i) for i in range(dim)]
+        return values
+    
+    raise ValueError(f"Unknown constraint type: {type(constraint)}")
+
+def get_constraint_dual(model, constraint):
+    if isinstance(constraint, ConstraintIndex):
+        return model.get_constraint_dual(constraint.index)
+    elif isinstance(constraint, NLConstraintIndex):
+        index = constraint.index
+        dim = constraint.dim
+        values = [model.get_constraint_dual(index + i) for i in range(dim)]
+        return values
+    
+    raise ValueError(f"Unknown constraint type: {type(constraint)}")
+
+
 constraint_attribute_get_func_map = {
-    ConstraintAttribute.Primal: lambda model, constraint: model.get_constraint_primal(
-        constraint
-    ),
-    ConstraintAttribute.Dual: lambda model, constraint: model.get_constraint_dual(
-        constraint
-    ),
+    ConstraintAttribute.Primal: get_constraint_primal,
+    ConstraintAttribute.Dual: get_constraint_dual,
 }
 
 constraint_attribute_set_func_map = {}
@@ -332,8 +355,8 @@ def optimize(self, jit_engine="LLVM"):
 
         super().optimize()
 
-    def register_function(self, f, /, x, name, p=()):
-        adfun = trace_adfun(f, x, p)
+    def register_function(self, f, /, var, name, param=()):
+        adfun = trace_adfun(f, var, param)
         return super().register_function(adfun, name)
 
     def get_variable_attribute(self, variable, attribute: VariableAttribute):