[breaking] scalar row indexing produces column vector (#624)

* fix row vector bug

* format

* Fix docs

* Update

* Update

* Update

---------

Co-authored-by: odow <[email protected]>

Author: ericphanson

docs/src/examples/optimization_with_complex_variables/power_flow_optimization.jl

@@ -18,7 +18,7 @@ W = ComplexVariable(n, n);
 objective = real(sum(diag(W)));
 c1 = Constraint[];
 for i in 2:n
-    push!(c1, sum(W[i, :] .* (Y[i, :]')) == inj[i])
+    push!(c1, dot(Y[i, :], W[i, :]) == inj[i])
 end
 c2 = isposdef(W)
 c3 = real(W[1, 1]) == 1.06^2;

src/atoms/IndexAtom.jl

@@ -9,18 +9,14 @@ mutable struct IndexAtom <: AbstractExpr
     rows::Union{AbstractArray,Nothing}
     cols::Union{AbstractArray,Nothing}
     inds::Union{AbstractArray,Nothing}
+end
 
-    function IndexAtom(
-        x::AbstractExpr,
-        rows::AbstractArray,
-        cols::AbstractArray,
-    )
-        return new((x,), (length(rows), length(cols)), rows, cols, nothing)
-    end
+function IndexAtom(x::AbstractExpr, rows::AbstractArray, cols::AbstractArray)
+    return IndexAtom((x,), (length(rows), length(cols)), rows, cols, nothing)
+end
 
-    function IndexAtom(x::AbstractExpr, inds::AbstractArray)
-        return new((x,), (length(inds), 1), nothing, nothing, inds)
-    end
+function IndexAtom(x::AbstractExpr, inds::AbstractArray)
+    return IndexAtom((x,), (length(inds), 1), nothing, nothing, inds)
 end
 
 head(io::IO, ::IndexAtom) = print(io, "index")
@@ -33,7 +29,11 @@ curvature(::IndexAtom) = ConstVexity()
 
 function evaluate(x::IndexAtom)
     result = if x.inds === nothing
-        getindex(evaluate(x.children[1]), x.rows, x.cols)
+        # reshape to ensure we are respecting that a scalar row index
+        # creates a column vector. We can't just check `length(x.rows)`
+        # since that doesn't distinguish between `i` and `i:i`. But
+        # we had that info when we set the size, so we will use it now.
+        reshape(getindex(evaluate(x.children[1]), x.rows, x.cols), x.size)
     else
         getindex(evaluate(x.children[1]), x.inds)
     end
@@ -98,7 +98,11 @@ function Base.getindex(x::AbstractExpr, row::Real, col::Real)
 end
 
 function Base.getindex(x::AbstractExpr, row::Real, cols::AbstractVector{<:Real})
-    return getindex(x, row:row, cols)
+    # In this case, we must construct a column vector
+    # https://github.com/jump-dev/Convex.jl/issues/509
+    # Here we construct `getindex(x, row:row, cols)`
+    # except with the size set to be that of a column vector
+    return IndexAtom((x,), (length(cols), 1), row:row, cols, nothing)
 end
 
 function Base.getindex(x::AbstractExpr, rows::AbstractVector{<:Real}, col::Real)

test/test_atoms.jl

@@ -523,6 +523,26 @@ function test_IndexAtom()
     return
 end
 
+function test_IndexAtom_issue_509()
+    # *Scalar* indexing creates a column vector
+    # https://github.com/jump-dev/Convex.jl/issues/509
+    x = rand(2, 3)
+    ϕ = Variable(2, 3)
+    Convex.set_value!(ϕ, rand(2, 3))
+    i = j = 1
+    # "column-vector" in Convex.jl is a 1 column matrix
+    @test size(ϕ[i, :]) == (3, 1)
+    @test size(ϕ[i, 1:2]) == (2, 1)
+    @test size(ϕ[i, [2, 1]]) == (2, 1)
+    @test size(ϕ[i:i, :]) == (1, 3)
+    @test size(evaluate(ϕ[i, :])) == (3,)
+    @test size(evaluate(ϕ[i:i, :])) == (1, 3)
+    # "Column vector", not 3x3 matrix!
+    atom = broadcast(*, ϕ[i, :] - ϕ[j, :], x[i, :] - x[j, :])
+    @test size(evaluate(atom)) == (3, 1)
+    return
+end
+
 ### affine/MultiplyAtom
 
 function test_MultiplyAtom()