Add support for collections of variables in ComplexPlane (#3184)

jump-dev · Jan 11, 2023 · 5c0b0d8 · 5c0b0d8
1 parent 3bec576
commit 5c0b0d8
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/src/variables.jl b/src/variables.jl
@@ -1388,6 +1388,22 @@ function JuMP.add_variable(
     )
 end
 
+function build_variable(
+    _error::Function,
+    variables::AbstractArray{<:ScalarVariable},
+    set::ComplexPlane,
+)
+    return build_variable.(_error, variables, Ref(set))
+end
+
+function add_variable(
+    model::Model,
+    variables::AbstractArray{<:ComplexVariable},
+    name::Union{<:AbstractArray{String},String} = "",
+)
+    return add_variable.(model, variables, name)
+end
+
 """
     reduced_cost(x::VariableRef)::Float64
 

diff --git a/test/test_variable.jl b/test/test_variable.jl
@@ -1212,6 +1212,15 @@ function test_complex_variable()
     return
 end
 
+function test_complex_variable()
+    model = Model()
+    @variable(model, x[1:2] in ComplexPlane())
+    @test x[1] isa GenericAffExpr{ComplexF64,VariableRef}
+    @test x[2] isa GenericAffExpr{ComplexF64,VariableRef}
+    @test num_variables(model) == 4
+    return
+end
+
 function test_extension_complex_variable_errors(
     ModelType = Model,
     VariableRefType = VariableRef,