Rename compare_ops and add test for ifelse

Author: metab0t

include/pyoptinterface/nlexpr.hpp

@@ -68,12 +68,12 @@ enum class BinaryOperator
 	Pow,
 
 	// compare
-	Lessthan,
-	Lessequal,
+	LessThan,
+	LessEqual,
 	Equal,
-	Notequal,
-	Greaterequal,
-	Greaterthan,
+	NotEqual,
+	GreaterEqual,
+	GreaterThan,
 };
 
 bool is_binary_compare_op(BinaryOperator op);

lib/cppad_interface.cpp

@@ -262,12 +262,12 @@ CppAD::AD<double> cppad_build_binary_expression(BinaryOperator op, const CppAD::
 	case BinaryOperator::Pow: {
 		return CppAD::pow(left, right);
 	}
-	case BinaryOperator::Lessthan:
-	case BinaryOperator::Lessequal:
+	case BinaryOperator::LessThan:
+	case BinaryOperator::LessEqual:
 	case BinaryOperator::Equal:
-	case BinaryOperator::Notequal:
-	case BinaryOperator::Greaterequal:
-	case BinaryOperator::Greaterthan: {
+	case BinaryOperator::NotEqual:
+	case BinaryOperator::GreaterEqual:
+	case BinaryOperator::GreaterThan: {
 		throw std::runtime_error("Currently comparision operator can only be used with ifelse "
 		                         "function and cannot be evaluated as value");
 	}
@@ -285,22 +285,22 @@ CppAD::AD<double> cppad_build_ternary_expression(BinaryOperator compare_op,
 {
 	switch (compare_op)
 	{
-	case BinaryOperator::Lessthan: {
+	case BinaryOperator::LessThan: {
 		return CppAD::CondExpLt(compare_left, compare_right, then_result, else_result);
 	}
-	case BinaryOperator::Lessequal: {
+	case BinaryOperator::LessEqual: {
 		return CppAD::CondExpLe(compare_left, compare_right, then_result, else_result);
 	}
 	case BinaryOperator::Equal: {
 		return CppAD::CondExpEq(compare_left, compare_right, then_result, else_result);
 	}
-	case BinaryOperator::Notequal: {
+	case BinaryOperator::NotEqual: {
 		return CppAD::CondExpEq(compare_left, compare_right, else_result, then_result);
 	}
-	case BinaryOperator::Greaterequal: {
+	case BinaryOperator::GreaterEqual: {
 		return CppAD::CondExpGe(compare_left, compare_right, then_result, else_result);
 	}
-	case BinaryOperator::Greaterthan: {
+	case BinaryOperator::GreaterThan: {
 		return CppAD::CondExpGt(compare_left, compare_right, then_result, else_result);
 	}
 	default: {

lib/nlexpr.cpp

@@ -83,5 +83,5 @@ NaryOperator ExpressionGraph::get_nary_operator(const ExpressionHandle &expressi
 
 bool is_binary_compare_op(BinaryOperator op)
 {
-	return (op >= BinaryOperator::Lessthan) && (op <= BinaryOperator::Greaterthan);
+	return (op >= BinaryOperator::LessThan) && (op <= BinaryOperator::GreaterThan);
 }

lib/nlexpr_ext.cpp

@@ -46,12 +46,12 @@ NB_MODULE(nlexpr_ext, m)
 	    .value("Div", BinaryOperator::Div)
 	    .value("Pow", BinaryOperator::Pow)
 	    // compare ops
-	    .value("Lessthan", BinaryOperator::Lessthan)
-	    .value("Lessequal", BinaryOperator::Lessequal)
+	    .value("LessThan", BinaryOperator::LessThan)
+	    .value("LessEqual", BinaryOperator::LessEqual)
 	    .value("Equal", BinaryOperator::Equal)
-	    .value("Notequal", BinaryOperator::Notequal)
-	    .value("Greaterequal", BinaryOperator::Greaterequal)
-	    .value("Greaterthan", BinaryOperator::Greaterthan);
+	    .value("NotEqual", BinaryOperator::NotEqual)
+	    .value("GreaterEqual", BinaryOperator::GreaterEqual)
+	    .value("GreaterThan", BinaryOperator::GreaterThan);
 
 	nb::enum_<TernaryOperator>(m, "TernaryOperator")
 	    .value("IfThenElse", TernaryOperator::IfThenElse);

src/pyoptinterface/_src/function_tracing.py

@@ -29,13 +29,13 @@ def pow_int(graph, expr, N):
 
     M, r = divmod(N, 2)
 
-    pow_2 = graph.add_binary(BinaryOperator.Mul, [expr, expr])
+    pow_2 = graph.add_nary(NaryOperator.Mul, [expr, expr])
     pow_2M = pow_int(graph, pow_2, M)
 
     if r == 0:
         return pow_2M
     else:
-        return graph.add_binary(BinaryOperator.Mul, [pow_2M, expr])
+        return graph.add_nary(NaryOperator.Mul, [pow_2M, expr])
 
 
 @dataclass

tests/test_nlp.py

@@ -49,19 +49,6 @@ def con(vars):
         nlfunc.sqrt,
         nlfunc.tan,
     ]
-    py_funcs = [
-        abs,
-        math.acos,
-        math.asin,
-        math.atan,
-        math.cos,
-        math.exp,
-        math.log,
-        math.pow,
-        math.sin,
-        math.sqrt,
-        math.tan,
-    ]
 
     def all_nlfuncs(vars):
         x = vars.x
@@ -102,13 +89,9 @@ def all_nlfuncs(vars):
         all_nlfuncs_con, poi.ConstraintAttribute.Primal
     )
 
-    correct_con_values = []
-    for f in py_funcs:
-        if f == math.pow:
-            v = f(x_value, 2)
-        else:
-            v = f(x_value)
-        correct_con_values.append(v)
+    vars = nlfunc.Vars(x=x_value)
+    correct_con_values = all_nlfuncs(vars)
+
     assert con_values == pytest.approx(correct_con_values)
 
 
@@ -157,6 +140,30 @@ def con(vars, params):
     assert x_values == pytest.approx(correct_x_values)
 
 
+def test_nlfunc_ifelse():
+    if not ipopt.is_library_loaded():
+        pytest.skip("Ipopt library is not loaded")
+
+    for x_, fx in zip([0.2, 0.5, 1.0, 2.0, 3.0], [0.2, 0.5, 1.0, 4.0, 9.0]):
+        model = ipopt.Model()
+
+        x = model.add_variable(lb=0.0, ub=10.0, start=1.0)
+
+        def con(vars):
+            x = vars.x
+            return nlfunc.ifelse(x > 1.0, x**2, x)
+
+        con_f = model.register_function(con)
+        model.add_nl_constraint(con_f, vars=nlfunc.Vars(x=x), lb=[fx])
+
+        model.set_objective(x)
+
+        model.optimize()
+
+        x_value = model.get_value(x)
+        assert x_value == pytest.approx(x_)
+
+
 if __name__ == "__main__":
     test_ipopt()
     test_nlp_param()