Add support for complex variable (#14)

* Add support for complex variable

* Fix format

Author: blegat

README.md

@@ -19,8 +19,9 @@ const COI = ComplexOptInterface
 
 model = Model()
 COI.add_all_bridges(model)
-@variable(model, x[1:2, 1:2] in COI.HermitianPSDCone())
-@constraint(model, x[1, 1] + x[2, 2] * im == 1 + 2im)
+@variable(model, x in COI.ComplexPlane(), start = 5 + 6im, lower_bound = 1 + 2im, upper_bound = 3 + 4im)
+@variable(model, Q[1:2, 1:2] in COI.HermitianPSDCone())
+@constraint(model, (1 + 2im) * x + Q[1, 1] + Q[2, 2] * im == 1 + 2im)
 ```
 
 ## Design considerations

src/jump.jl

@@ -7,6 +7,33 @@ function add_all_bridges(model::JuMP.Model)
     return JuMP.add_bridge(model, Bridges.Constraint.SplitZeroBridge)
 end
 
+struct ComplexPlane end
+
+struct ComplexVariable{S,T,U,V} <: JuMP.AbstractVariable
+    info::JuMP.VariableInfo{S,T,U,V}
+end
+
+function JuMP.build_variable(::Function, v::JuMP.ScalarVariable, ::ComplexPlane)
+    return ComplexVariable(v.info)
+end
+
+function JuMP.add_variable(
+    model::JuMP.Model,
+    v::ComplexVariable,
+    name::String = "",
+)
+    model.is_model_dirty = true
+    var = JuMP.ScalarVariable(v.info)
+    real_part = JuMP.add_variable(model, _real(var), _real(name))
+    imag_part = JuMP.add_variable(model, _imag(var), _imag(name))
+    # Efficiently build `real_part + imag_part * im`
+    return JuMP.GenericAffExpr{ComplexF64,JuMP.VariableRef}(
+        zero(ComplexF64),
+        real_part => one(ComplexF64),
+        imag_part => convert(ComplexF64, im),
+    )
+end
+
 struct HermitianPSDCone end
 
 struct HermitianMatrixShape <: JuMP.AbstractShape
@@ -64,8 +91,8 @@ function _mapinfo(f::Function, v::JuMP.ScalarVariable)
     )
 end
 
-_real(s::String) = s
-_imag(s::String) = s
+_real(s::String) = string("real(", s, ")")
+_imag(s::String) = string("imag(", s, ")")
 
 _real(v::JuMP.ScalarVariable) = _mapinfo(real, v)
 _imag(v::JuMP.ScalarVariable) = _mapinfo(imag, v)

test/jump.jl

@@ -9,15 +9,39 @@ using Test
 function test_readme_example()
     model = Model()
     COI.add_all_bridges(model)
-    @variable(model, x[1:2, 1:2] in COI.HermitianPSDCone())
-    @test num_variables(model) == 4
+    @variable(
+        model,
+        x in COI.ComplexPlane(),
+        start = 5 + 6im,
+        lower_bound = 1 + 2im,
+        upper_bound = 3 + 4im
+    )
+    xr = first(x.terms).first
+    xi = collect(x.terms)[2].first
+    @test lower_bound(xr) == 1
+    @test upper_bound(xr) == 3
+    @test start_value(xr) == 5
+    @test lower_bound(xi) == 2
+    @test upper_bound(xi) == 4
+    @test start_value(xi) == 6
+    @variable(model, Q[1:2, 1:2] in COI.HermitianPSDCone())
+    @test num_variables(model) == 6
     v = all_variables(model)
-    @test x[1, 1] == 1v[1]
-    @test x[1, 2] == v[2] + v[4] * im
-    @test x[2, 2] == 1v[3]
+    @test xr == v[1]
+    @test name(v[1]) == "real(x)"
+    @test xi == v[2]
+    @test name(v[2]) == "imag(x)"
+    @test x == v[1] + v[2] * im
+    @test name(v[3]) == "real(Q[1,1])"
+    @test Q[1, 1] == 1v[3]
+    @test name(v[4]) == "real(Q[1,2])"
+    @test name(v[6]) == "imag(Q[1,2])"
+    @test Q[1, 2] == v[4] + v[6] * im
+    @test name(v[5]) == "real(Q[2,2])"
+    @test Q[2, 2] == 1v[5]
     # FIXME needs https://github.com/jump-dev/JuMP.jl/pull/2899
-    #@test x[2, 1] == conj(x[1, 2])
-    @constraint(model, x[1, 1] + x[2, 2] * im == 1 + 2im)
+    #@test Q[2, 1] == conj(x[1, 2])
+    @constraint(model, Q[1, 1] + Q[2, 2] * im == 1 + 2im)
 end
 
 function test_simple_equality()