Implement Base.isreal for complex expressions (#3252)

jump-dev · Feb 27, 2023 · 6329762 · 6329762
1 parent 2976346
commit 6329762
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/src/aff_expr.jl b/src/aff_expr.jl
@@ -203,6 +203,7 @@ Base.conj(a::GenericAffExpr{<:Real}) = a
 Base.real(a::GenericAffExpr{<:Real}) = a
 Base.imag(a::GenericAffExpr{<:Real}) = zero(a)
 Base.abs2(a::GenericAffExpr{<:Real}) = a^2
+Base.isreal(x::GenericAffExpr{<:Real}) = true
 
 Base.conj(a::GenericAffExpr{<:Complex}) = map_coefficients(conj, a)
 
@@ -224,6 +225,10 @@ function Base.abs2(a::GenericAffExpr{<:Complex})
     return add_to_expression!(real(a)^2, imag_a, imag_a)
 end
 
+function Base.isreal(x::GenericAffExpr{<:Complex})
+    return isreal(x.constant) && all(isreal, values(x.terms))
+end
+
 # Needed for cases when Julia uses `x == 0` instead of `iszero(x)` (e.g., in the
 # stdlib).
 Base.:(==)(x::GenericAffExpr, y::Number) = isempty(x.terms) && x.constant == y

diff --git a/src/quad_expr.jl b/src/quad_expr.jl
@@ -114,11 +114,16 @@ Base.broadcastable(q::GenericQuadExpr) = Ref(q)
 Base.conj(a::GenericQuadExpr{<:Real}) = a
 Base.real(a::GenericQuadExpr{<:Real}) = a
 Base.imag(a::GenericQuadExpr{<:Real}) = a
+Base.isreal(::GenericQuadExpr{<:Real}) = true
 
 Base.conj(a::GenericQuadExpr{<:Complex}) = map_coefficients(conj, a)
 Base.real(a::GenericQuadExpr{<:Complex}) = map_coefficients(real, a)
 Base.imag(a::GenericQuadExpr{<:Complex}) = map_coefficients(imag, a)
 
+function Base.isreal(x::GenericQuadExpr{<:Complex})
+    return isreal(x.aff) && all(isreal, values(x.terms))
+end
+
 # Needed for cases when Julia uses `x == 0` instead of `iszero(x)` (e.g., in the
 # stdlib).
 Base.:(==)(x::GenericQuadExpr, y::Number) = isempty(x.terms) && x.aff == y

diff --git a/src/variables.jl b/src/variables.jl
@@ -200,6 +200,7 @@ Base.conj(v::AbstractVariableRef) = v
 Base.real(v::AbstractVariableRef) = v
 Base.imag(v::AbstractVariableRef) = zero(v)
 Base.abs2(v::AbstractVariableRef) = v^2
+Base.isreal(::AbstractVariableRef) = true
 
 """
     VariableRef <: AbstractVariableRef

diff --git a/test/test_complex.jl b/test/test_complex.jl
@@ -257,4 +257,20 @@ function test_complex_hermitian_constraint()
     return
 end
 
+function test_isreal()
+    model = Model()
+    @variable(model, x[1:2])
+    @test isreal(x[1])
+    @test isreal(2 * x[1] + 3) == true
+    @test isreal(2 * x[1] + 3 * x[2]) == true
+    @test isreal(2 * x[1] + 3 * x[2] + 4) == true
+    @test isreal(2 * x[1] + 0im) == true
+    @test isreal(2 * x[1] + 1im) == false
+    @test isreal(2im * x[1] + 1) == false
+    @test isreal(x[1] * x[2] + 1) == true
+    @test isreal(x[1] * x[2] + 1im) == false
+    @test isreal(x[1] * x[2] * 1im + 2) == false
+    return
+end
+
 end