Skip to content

Commit

Permalink
Improve printing of complex numbers (#3332)
Browse files Browse the repository at this point in the history 
* Improve printing of complex numbers

* Doctest fixes for complex printing update

* Fix test

* Fix printing tests

* Improve printing with pure imaginary

* Fix doctests

---------

Co-authored-by: odow <o.dowson@gmail.com>
  • Loading branch information
@blegat @odow
blegat and odow authored May 4, 2023
1 parent 468f5bf commit b5e96b3
Show file tree
Hide file tree
Showing 9 changed files with 116 additions and 55 deletions.
24 changes: 12 additions & 12 deletions docs/src/manual/complex.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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}:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
12 changes: 6 additions & 6 deletions docs/src/manual/constraints.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 6 additions & 6 deletions docs/src/manual/variables.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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):
Expand All @@ -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?
Expand Down
85 changes: 70 additions & 15 deletions src/print.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions test/test_complex.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion test/test_constraint.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 " *
Expand Down
10 changes: 5 additions & 5 deletions test/test_operator.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
14 changes: 10 additions & 4 deletions test/test_print.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand Down
8 changes: 4 additions & 4 deletions test/test_variable.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit b5e96b3

Please sign in to comment.