Improve printing of complex numbers

docs/src/manual/complex.md

@@ -21,7 +21,7 @@ Create a complex-valued variable using [`ComplexPlane`](@ref):
 julia> model = Model();
 
 julia> @variable(model, x in ComplexPlane())
-real(x) + (0.0 + 1.0im) imag(x)
+real(x) + imag(x) im
 ```
 
 Note that `x` is not a [`VariableRef`](@ref); instead, it is an affine
@@ -64,7 +64,7 @@ To create an anonymous variable, use the `set` keyword argument:
 julia> model = Model();
 
 julia> x = @variable(model, set = ComplexPlane())
-_[1] + (0.0 + 1.0im) _[2]
+_[1] + _[2] im
 ```
 
 ## Complex-valued variable bounds
@@ -85,7 +85,7 @@ julia> @variable(
            upper_bound = 2.0 + 3.0im,
            start = 4im,
        )
-real(x) + (0.0 + 1.0im) imag(x)
+real(x) + imag(x) im
 
 julia> vars = all_variables(model)
 2-element Vector{VariableRef}:
@@ -125,7 +125,7 @@ julia> set_silent(model)
 julia> @variable(model, x[1:2]);
 
 julia> @constraint(model, (1 + 2im) * x[1] + 3 * x[2] == 4 + 5im)
-(1.0 + 2.0im) x[1] + (3.0 + 0.0im) x[2] = (4.0 + 5.0im)
+(1 + 2im) x[1] + 3 x[2] = (4 + 5im)
 
 julia> optimize!(model)
 
@@ -168,7 +168,7 @@ julia> set_silent(model)
 julia> @variable(model, x in ComplexPlane());
 
 julia> @constraint(model, (1 + 2im) * x + 3 * x == 4 + 5im)
-(4.0 + 2.0im) real(x) + (-2.0 + 4.0im) imag(x) = (4.0 + 5.0im)
+(4 + 2im) real(x) + (-2 + 4im) imag(x) = (4 + 5im)
 
 julia> optimize!(model)
 
@@ -207,9 +207,9 @@ julia> model = Model();
 
 julia> @variable(model, H[1:3, 1:3] in HermitianPSDCone())
 3×3 LinearAlgebra.Hermitian{GenericAffExpr{ComplexF64, VariableRef}, Matrix{GenericAffExpr{ComplexF64, VariableRef}}}:
- real(H[1,1])                               …  real(H[1,3]) + (0.0 + 1.0im) imag(H[1,3])
- real(H[1,2]) + (0.0 - 1.0im) imag(H[1,2])     real(H[2,3]) + (0.0 + 1.0im) imag(H[2,3])
- real(H[1,3]) + (0.0 - 1.0im) imag(H[1,3])     real(H[3,3])
+ real(H[1,1])                    …  real(H[1,3]) + imag(H[1,3]) im
+ real(H[1,2]) - imag(H[1,2]) im     real(H[2,3]) + imag(H[2,3]) im
+ real(H[1,3]) - imag(H[1,3]) im     real(H[3,3])
 ```
 
 Behind the scenes, JuMP has created nine real-valued decision variables:
@@ -267,12 +267,12 @@ julia> import LinearAlgebra
 
 julia> H = LinearAlgebra.Hermitian([x[1] 1im; -1im -x[2]])
 2×2 LinearAlgebra.Hermitian{GenericAffExpr{ComplexF64, VariableRef}, Matrix{GenericAffExpr{ComplexF64, VariableRef}}}:
- x[1]           (0.0 + 1.0im)
- (0.0 - 1.0im)  (-1.0 - 0.0im) x[2]
+ x[1]  im
+ -im   -x[2]
 
 julia> @constraint(model, H in HermitianPSDCone())
-[x[1]           (0.0 + 1.0im);
- (0.0 - 1.0im)  (-1.0 - 0.0im) x[2]] ∈ HermitianPSDCone()
+[x[1]  im;
+ -im   -x[2]] ∈ HermitianPSDCone()
 ```
 
 !!! note

docs/src/manual/constraints.md

@@ -225,17 +225,17 @@ julia> model = Model();
 
 julia> @variable(model, X[1:2, 1:2] in HermitianPSDCone())
 2×2 Hermitian{GenericAffExpr{ComplexF64, VariableRef}, Matrix{GenericAffExpr{ComplexF64, VariableRef}}}:
- real(X[1,1])                               …  real(X[1,2]) + (0.0 + 1.0im) imag(X[1,2])
- real(X[1,2]) + (0.0 - 1.0im) imag(X[1,2])     real(X[2,2])
+ real(X[1,1])                    real(X[1,2]) + imag(X[1,2]) im
+ real(X[1,2]) - imag(X[1,2]) im  real(X[2,2])
 
 julia> @constraint(model, X == LinearAlgebra.I)
-[real(X[1,1]) + (-1.0 - 0.0im)              real(X[1,2]) + (0.0 + 1.0im) imag(X[1,2]);
- real(X[1,2]) + (0.0 - 1.0im) imag(X[1,2])  real(X[2,2]) + (-1.0 - 0.0im)] ∈ Zeros()
+[real(X[1,1]) - 1                real(X[1,2]) + imag(X[1,2]) im;
+ real(X[1,2]) - imag(X[1,2]) im  real(X[2,2]) - 1] ∈ Zeros()
 
 julia> @constraint(model, X .== LinearAlgebra.I)
 2×2 Matrix{ConstraintRef{Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{ComplexF64}, MathOptInterface.EqualTo{ComplexF64}}, ScalarShape}}:
- real(X[1,1]) = (1.0 - 0.0im)                               …  real(X[1,2]) + (0.0 + 1.0im) imag(X[1,2]) = (0.0 - 0.0im)
- real(X[1,2]) + (0.0 - 1.0im) imag(X[1,2]) = (0.0 + 0.0im)     real(X[2,2]) = (1.0 - 0.0im)
+ real(X[1,1]) = 1                    real(X[1,2]) + imag(X[1,2]) im = 0
+ real(X[1,2]) - imag(X[1,2]) im = 0  real(X[2,2]) = 1
 ```
 
 ## Containers of constraints

docs/src/manual/variables.md

@@ -1247,8 +1247,8 @@ julia> model = Model();
 
 julia> @variable(model, H[1:2, 1:2] in HermitianPSDCone())
 2×2 LinearAlgebra.Hermitian{GenericAffExpr{ComplexF64, VariableRef}, Matrix{GenericAffExpr{ComplexF64, VariableRef}}}:
- real(H[1,1])                               …  real(H[1,2]) + (0.0 + 1.0im) imag(H[1,2])
- real(H[1,2]) + (0.0 - 1.0im) imag(H[1,2])     real(H[2,2])
+ real(H[1,1])                    real(H[1,2]) + imag(H[1,2]) im
+ real(H[1,2]) - imag(H[1,2]) im  real(H[2,2])
 ```
 
 This adds 4 real variables in the [`MOI.HermitianPositiveSemidefiniteConeTriangle`](@ref):
@@ -1266,8 +1266,8 @@ julia> model = Model();
 
 julia> @variable(model, x[1:2, 1:2], Hermitian)
 2×2 LinearAlgebra.Hermitian{GenericAffExpr{ComplexF64, VariableRef}, Matrix{GenericAffExpr{ComplexF64, VariableRef}}}:
- real(x[1,1])                               …  real(x[1,2]) + (0.0 + 1.0im) imag(x[1,2])
- real(x[1,2]) + (0.0 - 1.0im) imag(x[1,2])     real(x[2,2])
+ real(x[1,1])                    real(x[1,2]) + imag(x[1,2]) im
+ real(x[1,2]) - imag(x[1,2]) im  real(x[2,2])
 ```
 
 This is equivalent to declaring the variable in [`HermitianMatrixSpace`](@ref):
@@ -1276,8 +1276,8 @@ julia> model = Model();
 
 julia> @variable(model, x[1:2, 1:2] in HermitianMatrixSpace())
 2×2 LinearAlgebra.Hermitian{GenericAffExpr{ComplexF64, VariableRef}, Matrix{GenericAffExpr{ComplexF64, VariableRef}}}:
- real(x[1,1])                               …  real(x[1,2]) + (0.0 + 1.0im) imag(x[1,2])
- real(x[1,2]) + (0.0 - 1.0im) imag(x[1,2])     real(x[2,2])
+ real(x[1,1])                    real(x[1,2]) + imag(x[1,2]) im
+ real(x[1,2]) - imag(x[1,2]) im  real(x[2,2])
 ```
 
 ### Why use variables constrained on creation?

src/print.jl

@@ -51,13 +51,23 @@ function _is_one_for_printing(coef)
     return _is_zero_for_printing(abs(coef) - oneunit(coef))
 end
 
-_is_one_for_printing(coef::Complex{T}) where {T} = coef == one(T)
+function _is_one_for_printing(coef::Complex{T}) where {T}
+    r, i = reim(coef)
+    return _is_one_for_printing(r) && _is_zero_for_printing(i)
+end
 
 function _is_zero_for_printing(coef::Complex)
     return _is_zero_for_printing(real(coef)) &&
            _is_zero_for_printing(imag(coef))
 end
 
+_is_im_for_printing(coef) = false
+
+function _is_im_for_printing(coef::Complex)
+    r, i = reim(coef)
+    return _is_zero_for_printing(r) && _is_one_for_printing(i)
+end
+
 # Helper function that rounds carefully for the purposes of printing Reals
 # e.g.   5.3  =>  5.3
 #        1.0  =>  1
@@ -67,6 +77,28 @@ _string_round(::typeof(abs), x::Real) = _string_round(abs(x))
 
 _sign_string(x::Real) = x < zero(x) ? " - " : " + "
 
+function _string_round(::typeof(abs), x::Complex)
+    r, i = reim(x)
+    if _is_zero_for_printing(r)
+        return _string_round(Complex(r, abs(i)))
+    elseif _is_zero_for_printing(i)
+        return _string_round(Complex(abs(r), i))
+    else
+        return _string_round(x)
+    end
+end
+
+function _sign_string(x::Complex)
+    r, i = reim(x)
+    if _is_zero_for_printing(r)
+        return _sign_string(i)
+    elseif _is_zero_for_printing(i)
+        return _sign_string(r)
+    else
+        return " + "
+    end
+end
+
 # Fallbacks for other number types
 
 _string_round(x::Any) = string(x)
@@ -75,7 +107,29 @@ _string_round(::typeof(abs), x::Any) = _string_round(x)
 
 _sign_string(::Any) = " + "
 
-_string_round(x::Complex) = string("(", x, ")")
+function _string_round(x::Complex)
+    r, i = reim(x)
+    r_str = _string_round(r)
+    if _is_zero_for_printing(i)
+        return r_str
+    elseif _is_zero_for_printing(r)
+        if _is_one_for_printing(i)
+            if i < 0
+                return "-im"
+            else
+                return "im"
+            end
+        else
+            return string(_string_round(i), "im")
+        end
+    end
+    if _is_one_for_printing(i)
+        i_str = "im"
+    else
+        i_str = string(_string_round(abs, i), "im")
+    end
+    return string("(", r_str, _sign_string(i_str), i_str, ")")
+end
 
 # REPL-specific symbols
 # Anything here: https://en.wikipedia.org/wiki/Windows-1252
@@ -527,6 +581,16 @@ function function_string(mode::MIME"text/latex", v::AbstractVariableRef)
     return var_name
 end
 
+function _term_string(coef, factor)
+    if _is_one_for_printing(coef)
+        return factor
+    elseif _is_im_for_printing(coef)
+        return string(factor, " ", _string_round(abs, coef))
+    else
+        return string(_string_round(abs, coef), " ", factor)
+    end
+end
+
 # TODO(odow): remove show_constant in JuMP 1.0
 function function_string(mode, a::GenericAffExpr, show_constant = true)
     if length(linear_terms(a)) == 0
@@ -535,12 +599,7 @@ function function_string(mode, a::GenericAffExpr, show_constant = true)
     terms = fill("", 2 * length(linear_terms(a)))
     for (elm, (coef, var)) in enumerate(linear_terms(a))
         terms[2*elm-1] = _sign_string(coef)
-        v = function_string(mode, var)
-        if _is_one_for_printing(coef)
-            terms[2*elm] = v
-        else
-            terms[2*elm] = string(_string_round(abs, coef), " ", v)
-        end
+        terms[2*elm] = _term_string(coef, function_string(mode, var))
     end
     terms[1] = terms[1] == " - " ? "-" : ""
     ret = join(terms)
@@ -563,17 +622,13 @@ function function_string(mode, q::GenericQuadExpr)
         x = function_string(mode, var1)
         y = function_string(mode, var2)
         terms[2*elm-1] = _sign_string(coef)
-        if _is_one_for_printing(coef)
-            terms[2*elm] = "$x"
-        else
-            terms[2*elm] = string(_string_round(abs, coef), " ", x)
-        end
         if x == y
-            terms[2*elm] *= _math_symbol(mode, :sq)
+            factor = x * _math_symbol(mode, :sq)
         else
             times = mode == MIME("text/latex") ? "\\times " : "*"
-            terms[2*elm] *= string(times, y)
+            factor = string(x, times, y)
         end
+        terms[2*elm] = _term_string(coef, factor)
     end
     terms[1] = terms[1] == " - " ? "-" : ""
     ret = join(terms)

test/test_complex.jl

@@ -158,9 +158,9 @@ function test_complex_print()
     model = Model()
     @variable(model, x)
     y = (1 + 2im) * x + 1
-    @test sprint(show, y) == "(1.0 + 2.0im) x + (1.0 + 0.0im)"
+    @test sprint(show, y) == "(1 + 2im) x + 1"
     y = im * x
-    @test sprint(show, y) == "(0.0 + 1.0im) x"
+    @test sprint(show, y) == "x im"
     return
 end
 

test/test_constraint.jl

@@ -1476,7 +1476,7 @@ function test_extension_HermitianPSDCone_errors(
         "In `@constraint(model, H in HermitianPSDCone(), unknown_kw = 1)`:" *
         " Unrecognized constraint building format. Tried to invoke " *
         "`build_constraint(error, $aff_str[" *
-        "x (0.0 + 1.0im); (0.0 - 1.0im) (-1.0 - 0.0im) y], $(HermitianPSDCone()); unknown_kw = 1)`, but no " *
+        "x im; -im -y], $(HermitianPSDCone()); unknown_kw = 1)`, but no " *
         "such method exists. This is due to specifying an unrecognized " *
         "function, constraint set, and/or extra positional/keyword " *
         "arguments.\n\nIf you're trying to create a JuMP extension, you " *

test/test_operator.jl

@@ -282,15 +282,15 @@ function test_extension_uniform_scaling(
     @test_expression_with_string 2 * LinearAlgebra.I + x "x + 2"
     @test_expression_with_string LinearAlgebra.I + (x + 1) "x + 2"
     @test_expression_with_string 2 * LinearAlgebra.I - x "-x + 2"
-    @test_expression_with_string (2im * LinearAlgebra.I) - x "(-1.0 - 0.0im) x + (0.0 + 2.0im)"
+    @test_expression_with_string (2im * LinearAlgebra.I) - x "-x + 2im"
     @test_expression_with_string LinearAlgebra.I - (x - 1) "-x + 2"
-    @test_expression_with_string (LinearAlgebra.I * im) - (x - 1) "(-1.0 + 0.0im) x + (1.0 + 1.0im)"
+    @test_expression_with_string (LinearAlgebra.I * im) - (x - 1) "-x + (1 + im)"
     @test_expression_with_string LinearAlgebra.I * x "x"
-    @test_expression_with_string (LinearAlgebra.I * im) * x "(0.0 + 1.0im) x"
+    @test_expression_with_string (LinearAlgebra.I * im) * x "x im"
     @test_expression_with_string LinearAlgebra.I * (x + 1) "x + 1"
-    @test_expression_with_string (LinearAlgebra.I * im) * (x + 1) "(0.0 + 1.0im) x + (0.0 + 1.0im)"
+    @test_expression_with_string (LinearAlgebra.I * im) * (x + 1) "x im + im"
     @test_expression_with_string (x + 1) * LinearAlgebra.I "x + 1"
-    @test_expression_with_string (x + 1) * (LinearAlgebra.I * im) "(0.0 + 1.0im) x + (0.0 + 1.0im)"
+    @test_expression_with_string (x + 1) * (LinearAlgebra.I * im) "x im + im"
     return
 end
 

test/test_print.jl

@@ -850,11 +850,17 @@ function test_show_latex_parameter()
     return
 end
 
-function test_minus_one_complex_aff_expr()
+function test_complex_expr()
     model = Model()
     @variable(model, x)
+    @variable(model, y)
     f = 1.0im * x + 1.0im
-    @test sprint(show, im * f) == "(-1.0 + 0.0im) x + (-1.0 + 0.0im)"
+    @test sprint(show, im * f) == "-x - 1"
+    @test sprint(show, -f) == "-x im - im"
+    @test sprint(show, f * x) == "x² im + x im"
+    @test sprint(show, f * y) == "x*y im + y im"
+    @test sprint(show, f + y) == "x im + y + im"
+    @test sprint(show, 2f) == "2im x + 2im"
     return
 end
 
@@ -865,9 +871,9 @@ function test_print_hermitian_psd_cone()
     H = Hermitian([x[1] 1im; -1im x[2]])
     c = @constraint(model, H in HermitianPSDCone())
     @test sprint(io -> show(io, MIME("text/plain"), c)) ==
-          "[x[1]           (0.0 + 1.0im);\n (0.0 - 1.0im)  x[2]] $in_sym $(HermitianPSDCone())"
+          "[x[1]  im;\n -im   x[2]] $in_sym $(HermitianPSDCone())"
     @test sprint(io -> show(io, MIME("text/latex"), c)) ==
-          "\$\$ \\begin{bmatrix}\nx_{1} & (0.0 + 1.0im)\\\\\n(0.0 - 1.0im) & x_{2}\\\\\n\\end{bmatrix} \\in \\text{$(HermitianPSDCone())} \$\$"
+          "\$\$ \\begin{bmatrix}\nx_{1} & im\\\\\n-im & x_{2}\\\\\n\\end{bmatrix} \\in \\text{$(HermitianPSDCone())} \$\$"
     return
 end
 

test/test_variable.jl

@@ -1390,10 +1390,10 @@ function test_Hermitian_PSD_anon()
     x = parent(y)
     @test sprint(show, x[1, 1]) == "_[1]"
     @test sprint(show, y[1, 1]) == "_[1]"
-    @test sprint(show, x[1, 2]) == "_[2] + (0.0 + 1.0im) _[4]"
-    @test sprint(show, y[1, 2]) == "_[2] + (0.0 + 1.0im) _[4]"
-    @test sprint(show, x[2, 1]) == "_[2] + (-0.0 - 1.0im) _[4]"
-    @test sprint(show, y[2, 1]) == "_[2] + (0.0 - 1.0im) _[4]"
+    @test sprint(show, x[1, 2]) == "_[2] + _[4] im"
+    @test sprint(show, y[1, 2]) == "_[2] + _[4] im"
+    @test sprint(show, x[2, 1]) == "_[2] - _[4] im"
+    @test sprint(show, y[2, 1]) == "_[2] - _[4] im"
     @test sprint(show, x[2, 2]) == "_[3]"
     @test sprint(show, y[2, 2]) == "_[3]"
     return