Add support for constraint with AlgebraElement (#129)

* Add support for constraint with AlgebraElement

* Fix format

* Add tests

Author: blegat

src/constraint.jl

@@ -1,5 +1,7 @@
 export moments
 
+const _AE_APL = Union{SA.AlgebraElement,MP.AbstractPolynomialLike}
+
 abstract type PolynomialSet end
 
 _in(::MIME) = Sys.iswindows() ? "in" : "∈"
@@ -13,16 +15,10 @@ struct ZeroPoly <: PolynomialSet end
 struct NonNegPoly <: PolynomialSet end
 struct PosDefPolyMatrix <: PolynomialSet end
 
-function JuMP.function_string(
-    ::MIME"text/plain",
-    p::Union{SA.AlgebraElement,MP.AbstractPolynomialLike},
-)
+function JuMP.function_string(::MIME"text/plain", p::_AE_APL)
     return sprint(show, MIME"text/plain"(), p)
 end
-function JuMP.function_string(
-    mime::MIME"text/latex",
-    p::Union{SA.AlgebraElement,MP.AbstractPolynomialLike},
-)
+function JuMP.function_string(mime::MIME"text/latex", p::_AE_APL)
     return SA.trim_LaTeX(mime, sprint(show, MIME"text/latex"(), p))
 end
 
@@ -203,7 +199,7 @@ non_constant_coefficients(p) = non_constant(MP.coefficients(p))
 ## ZeroPoly
 function JuMP.build_constraint(
     error_fn::Function,
-    p::MP.AbstractPolynomialLike,
+    p::_AE_APL,
     s::ZeroPoly;
     domain::SS.AbstractSemialgebraicSet = SS.FullSpace(),
     kws...,
@@ -238,7 +234,7 @@ function JuMP.build_constraint(
 end
 function JuMP.build_constraint(
     error_fn::Function,
-    p::MP.AbstractPolynomialLike,
+    p::_AE_APL,
     s::MOI.EqualTo;
     kws...,
 )
@@ -304,15 +300,15 @@ end
 # `NonNegPoly`
 function JuMP.build_constraint(
     error_fn::Function,
-    p::MP.AbstractPolynomialLike,
+    p::_AE_APL,
     s::MOI.GreaterThan;
     kws...,
 )
     return JuMP.build_constraint(error_fn, p - s.lower, NonNegPoly(); kws...)
 end
 function JuMP.build_constraint(
     error_fn::Function,
-    p::MP.AbstractPolynomialLike,
+    p::_AE_APL,
     s::MOI.LessThan;
     kws...,
 )
@@ -324,7 +320,7 @@ end
 # need a more specific here to avoid ambiguity
 function JuMP.build_constraint(
     error_fn::Function,
-    p::AbstractMatrix{<:MP.AbstractPolynomialLike},
+    p::AbstractMatrix{<:_AE_APL},
     s::PSDCone;
     kws...,
 )
@@ -334,7 +330,7 @@ end
 # Needed for the syntax `@constraint(model, A >= B, PSDCone())`
 function JuMP.build_constraint(
     error_fn::Function,
-    f::AbstractMatrix{<:MP.AbstractPolynomialLike},
+    f::AbstractMatrix{<:_AE_APL},
     s::MOI.GreaterThan,
     extra::PSDCone,
 )
@@ -345,7 +341,7 @@ end
 # Needed for the syntax `@constraint(model, A <= B, PSDCone())`
 function JuMP.build_constraint(
     error_fn::Function,
-    f::AbstractMatrix{<:MP.AbstractPolynomialLike},
+    f::AbstractMatrix{<:_AE_APL},
     s::MOI.LessThan,
     extra::PSDCone,
 )

src/variable.jl

@@ -15,10 +15,21 @@ function JuMP.value(
     return MP.polynomial(JuMP.value.(f, MP.terms(p)), MP.SortedUniqState())
 end
 
+function JuMP.value(
+    f::Function,
+    p::SA.AlgebraElement{A,<:JuMP.AbstractJuMPScalar},
+) where {A}
+    return SA.AlgebraElement(JuMP.value.(f, SA.coeffs(p)), parent(p))
+end
+
 function JuMP.value(p::MP.AbstractPolynomialLike{<:JuMP.AbstractJuMPScalar})
     return JuMP.value(JuMP.value, p)
 end
 
+function JuMP.value(p::SA.AlgebraElement{A,<:JuMP.AbstractJuMPScalar}) where {A}
+    return JuMP.value(JuMP.value, p)
+end
+
 abstract type AbstractPoly end
 
 """

test/variable.jl

@@ -177,8 +177,13 @@ function test_value_function(var)
     p = α * x * y + β * x^2
     JuMP.fix(α, 2)
     JuMP.fix(β, 3)
-    @test_broken JuMP.value(p) == 2x * y + 3x^2
-    @test JuMP.value(fix_value, p) == 2x * y + 3x^2
+    expected = 2x * y + 3x^2
+    @test_broken JuMP.value(p) == expected
+    @test JuMP.value(fix_value, p) == expected
+    a = MB.algebra_element(p)
+    exp_a = MB.algebra_element(expected)
+    @test_broken JuMP.value(a) == exp_a
+    @test JuMP.value(fix_value, a) == exp_a
 end
 
 function runtests(var)