wip

ericphanson · ericphanson · commit dc3f4f3827b5 · 2023-07-16T14:22:50.000+02:00
diff --git a/Project.toml b/Project.toml
@@ -5,7 +5,6 @@ version = "0.16.0"
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
diff --git a/src/Context.jl b/src/Context.jl
@@ -18,7 +18,8 @@ mutable struct Context{T,M}
     detected_infeasible_during_formulation::Ref{Bool}
 
     # Cache
-    conic_form_cache::DataStructures.WeakKeyIdDict{Any, Any}
+    # conic_form_cache::DataStructures.WeakKeyIdDict{Any, Any}
+    conic_form_cache::IdDict{Any, Any}
 end
 
 function Context{T}(optimizer) where {T}
@@ -38,3 +39,4 @@ function Context{T}(optimizer) where {T}
         IdDict{Any,Any}()
     )
 end
+
diff --git a/src/Convex.jl b/src/Convex.jl
@@ -7,7 +7,7 @@ using SparseArrays
 using LDLFactorizations
 using AbstractTrees: AbstractTrees, children
 using SuiteSparseGraphBLAS
-using DataStructures
+# using DataStructures
 
 using MathOptInterface
 const MOI = MathOptInterface
@@ -130,6 +130,12 @@ function vec_tril(M)
     return M[inds]
 end
 
+const SPARSE_VECTOR{T} = GBVector{T,T}
+const SPARSE_MATRIX{T} = GBMatrix{T,T}
+spzeros(T, d) = GBVector{T,T}(d)
+spzeros(T, n, m) = GBMatrix{T,T}(n, m)
+spidentity(T, d) = GBMatrix{T,T}(Diagonal(ones(T, d)))
+
 include("Context.jl")
 ### modeling framework
 include("dcp.jl")
diff --git a/src/SparseTape.jl b/src/SparseTape.jl
@@ -1,6 +1,6 @@
 struct SparseAffineOperation{T}
-    matrix::Union{GBMatrix{T,T}}
-    vector::Union{GBVector{T,T}}
+    matrix::SPARSE_MATRIX{T}
+    vector::SPARSE_VECTOR{T}
 end
 
 # function Base.convert(
@@ -16,7 +16,7 @@ function SparseAffineOperation(
     A::AbstractMatrix{T},
     b::AbstractVector{T},
 ) where {T}
-    return SparseAffineOperation{T}(GBMatrix{T,T}(A), GBVector{T, T}(b))
+    return SparseAffineOperation{T}(SPARSE_MATRIX{T}(A), GBVector{T, T}(b))
 end
 
 
diff --git a/src/atoms/affine/diag.jl b/src/atoms/affine/diag.jl
@@ -78,7 +78,7 @@ function _conic_form!(context::Context{T}, x::DiagAtom) where {T}
         sz_diag = Base.min(num_rows + k, num_cols)
     end
 
-    select_diag = GBMatrix{T,T}(sz_diag, length(x.children[1]))
+    select_diag = SPARSE_MATRIX{T}(sz_diag, length(x.children[1]))
     for i in 1:sz_diag
         select_diag[i, start_index] = 1
         start_index += num_rows + 1
diff --git a/src/atoms/affine/diagm.jl b/src/atoms/affine/diagm.jl
@@ -63,8 +63,8 @@ function _conic_form!(context::Context{T}, x::DiagMatrixAtom) where {T}
     obj = conic_form!(context, only(children(x)))
 
     sz = x.size[1]
-    I = 1:sz+1:sz*sz
-    J = 1:sz
+    I = collect(1:sz+1:sz*sz)
+    J = collect(1:sz)
     V = one(T)
     coeff = GBMatrix{T, T}(I, J, V, sz * sz, sz)
     # coeff = sparse(, 1:sz, one(T),
diff --git a/src/atoms/affine/index.jl b/src/atoms/affine/index.jl
@@ -71,9 +71,9 @@ function _conic_form!(context::Context{T}, x::IndexAtom) where {T}
             end
         end
 
-        index_matrix = GBMatrix{T,T}(collect(1:sz), J, one(T), m, n)
+        index_matrix = SPARSE_MATRIX{T}(collect(1:sz), J, one(T), m, n)
     else
-        index_matrix = GBMatrix{T,T}(collect(1:length(x.inds)), collect(x.inds), one(T), m, n)
+        index_matrix = SPARSE_MATRIX{T}(collect(1:length(x.inds)), collect(x.inds), one(T), m, n)
     end
 
     return operate(add_operation, T, sign(x), index_matrix, obj)
diff --git a/src/atoms/affine/multiply_divide.jl b/src/atoms/affine/multiply_divide.jl
@@ -66,18 +66,18 @@ function real_convert(::Type{T}, x::Number) where {T}
     return T(x)
 end
 function real_convert(::Type{T}, x::AbstractMatrix) where {T}
-    return GBMatrix{T, T}(x)
+    return GBMatrix{T,T}(x)
 end
 
-function real_convert(::Type{T}, x::GBMatrix{T,T}) where {T}
+function real_convert(::Type{T}, x::SPARSE_MATRIX{T}) where {T}
     return x
 end
 
-function real_convert(::Type{T}, x::GBVector{T,T}) where {T}
+function real_convert(::Type{T}, x::SPARSE_VECTOR{T}) where {T}
     return x
 end
 function real_convert(::Type{T}, x::AbstractVector) where {T}
-    return GBVector{T,T}(x)
+    return SPARSE_VECTOR{T}(x)
 end
 
 function _conic_form!(context::Context{T}, x::MultiplyAtom) where {T}
@@ -158,7 +158,7 @@ function _conic_form!(context::Context{T}, x::MultiplyAtom) where {T}
 end
 
 # _id(T, n) = Diagonal(one(T)*I, n)
-_id(T, n) = gbidentity(T, n)
+_id(T, n) = spidentity(T, n)
 
 function *(x::AbstractExpr, y::AbstractExpr)
     if isequal(x, y) && x.size == (1, 1)
@@ -171,6 +171,32 @@ end
 *(x::AbstractExpr, y::Value) = MultiplyAtom(x, constant(y))
 /(x::AbstractExpr, y::Value) = MultiplyAtom(x, constant(1 ./ y))
 
+# ambiguity
+function Base.:(*)(
+    x::Convex.AbstractExpr,
+    y::Union{
+        LinearAlgebra.Transpose{
+            <:Any,
+            <:SuiteSparseGraphBLAS.AbstractGBArray{T,F,O},
+        },
+        SuiteSparseGraphBLAS.AbstractGBArray{T,F,O},
+    },
+) where {T,F,O}
+    return MultiplyAtom(x, constant(y))
+end
+function Base.:(*)(
+    x::Union{
+        LinearAlgebra.Transpose{
+            <:Any,
+            <:SuiteSparseGraphBLAS.AbstractGBArray{T,F,O},
+        },
+        SuiteSparseGraphBLAS.AbstractGBArray{T,F,O},
+    },
+    y::Convex.AbstractExpr,
+) where {T,F,O}
+    return MultiplyAtom(constant(x), y)
+end
+
 function dotmultiply(x, y)
     if vexity(x) != ConstVexity()
         if vexity(y) != ConstVexity()
diff --git a/src/atoms/affine/partialtrace.jl b/src/atoms/affine/partialtrace.jl
@@ -5,18 +5,18 @@
 # This function returns the jth term in the sum, namely
 # (I ⊗ <j| ⊗ I) x (I ⊗ |j> ⊗ I).
 function _term(x, j::Int, sys, dims)
-    a = sparse(1.0I, 1, 1)
-    b = sparse(1.0I, 1, 1)
+    a = spidentity(Float64, 1)
+    b = spidentity(Float64, 1)
     for (i_sys, dim) in enumerate(dims)
         if i_sys == sys
             # create a vector that is only 1 at its jth component
-            v = spzeros(dim, 1)
+            v = spzeros(Float64, dim, 1)
             v[j] = 1
             a = kron(a, v')
             b = kron(b, v)
         else
-            a = kron(a, sparse(1.0I, dim, dim))
-            b = kron(b, sparse(1.0I, dim, dim))
+            a = kron(a, spidentity(Float64, dim))
+            b = kron(b, spidentity(Float64, dim))
         end
     end
     return a * x * b
diff --git a/src/atoms/exp_+_sdp_cone/logdet.jl b/src/atoms/exp_+_sdp_cone/logdet.jl
@@ -42,7 +42,7 @@ function _conic_form!(context::Context{T}, x::LogDetAtom) where {T}
     X = conic_form!(context, v)
 
     t = conic_form!(context, Variable())
-    f = operate(vcat, T, sign(x), t, [one(T)], X)
+    f = operate(vcat, T, sign(x), t, SPARSE_VECTOR{T}([one(T)]), X)
     side_dimension = size(only(children(x)), 1)
 
     set = MOI.LogDetConeTriangle(side_dimension)
diff --git a/src/atoms/sdp_cone/operatornorm.jl b/src/atoms/sdp_cone/operatornorm.jl
@@ -73,7 +73,7 @@ function _conic_form!(context::Context{T}, x::OperatorNormAtom) where {T}
         t = Variable()
         p = minimize(
             t,
-            [t*gbidentity(T, m) A; A' t*gbidentity(T, n)] ⪰ 0,
+            [t*spidentity(T, m) A; A' t*spidentity(T, n)] ⪰ 0,
         )
         return conic_form!(context, p)
     else
diff --git a/src/constant.jl b/src/constant.jl
@@ -75,8 +75,8 @@ function evaluate(c::ComplexConstant)
 end
 
 mutable struct ComplexStructOfVec{T<:Real}
-    real_vec::Vector{T}
-    imag_vec::Vector{T}
+    real_vec::SPARSE_VECTOR{T}
+    imag_vec::SPARSE_VECTOR{T}
 end
 
 Base.real(c::ComplexStructOfVec) = c.real_vec
@@ -118,6 +118,13 @@ function output(x::Value)
     end
 end
 
+vec(x) = Base.vec(x)
+function vec(x::GBMatrix)
+    # Hacks to try to get `vec` to work
+    x = reshape(x, length(x), 1)
+    return x[:, 1]
+end
+
 evaluate(x::Constant) = output(x.value)
 
 sign(x::Constant) = x.sign
@@ -129,8 +136,11 @@ length(x::Constant) = length(x.value)
 function _conic_form!(::Context{T}, C::Constant) where {T}
     # this should happen at `Constant` creation?
     # No, we don't have access to `T` yet; that's problem-specific
-    if eltype(C.value) != T
-        C = Constant(convert(Matrix{T}, C.value))
-    end
-    return vec(C.value)
+    # if eltype(C.value) != T
+    #     C = Constant{T}(convert(GBMatrix{T}, C.value))
+    # end
+
+    x = SPARSE_VECTOR{T}(vec(C.value))
+
+    return x
 end
diff --git a/src/operate.jl b/src/operate.jl
@@ -2,8 +2,8 @@ complex_promote(tape::ComplexTape) = tape
 complex_promote(tape::SparseTape) = tape
 complex_promote(v::ComplexStructOfVec) = v
 
-function complex_promote(v::Vector)
-    return ComplexStructOfVec(v, zero(v))
+function complex_promote(v::AbstractVector{T}) where {T}
+    return ComplexStructOfVec(v, SPARSE_VECTOR{T}(length(v)))
 end
 
 # Here we run `complex_promote` and dispatch to either `real_operate`
@@ -17,30 +17,26 @@ function operate(op::F, ::Type{T}, sign::Sign, args...) where {F,T}
         if op === add_operation
             @assert length(args) == 2
             x, y = args
-            # if !(
-            #     typeof(x) in (
-            #         T,
-            #         Complex{T},
-            #         SparseMatrixCSC{T,Int},
-            #         SparseMatrixCSC{Complex{T},Int},
-            #     )
-            # )
-            #     error(
-            #         "Convex.jl internal error: unexpected type $(typeof(x)) for first argument of operation $op of with sign=$sign",
-            #     )
-            # end
-            # if !(
-            #     typeof(y) in (
-            #         SparseTape{T},
-            #         Vector{T},
-            #         ComplexTape{T},
-            #         ComplexStructOfVec{T},
-            #     )
-            # )
-            #     error(
-            #         "Convex.jl internal error: unexpected type $(typeof(y)) for second argument of operation $op of with sign=$sign",
-            #     )
-            # end
+            if !(
+                typeof(x) in
+                (T, Complex{T}, SPARSE_MATRIX{T}, GBMatrix{Complex{T},Complex{T}})
+            )
+                error(
+                    "Convex.jl internal error: unexpected type $(typeof(x)) for first argument of operation $op of with sign=$sign",
+                )
+            end
+            if !(
+                typeof(y) in (
+                    SparseTape{T},
+                    SPARSE_VECTOR{T},
+                    ComplexTape{T},
+                    ComplexStructOfVec{T},
+                )
+            )
+                error(
+                    "Convex.jl internal error: unexpected type $(typeof(y)) for second argument of operation $op of with sign=$sign",
+                )
+            end
         else
             # Everything should be either be
             # ComplexTape{T}, SparseTape{T}, or ComplexStructOfVec{T}
@@ -79,22 +75,25 @@ function operate(op::F, ::Type{T}, sign::Sign, args...) where {F,T}
                 @assert length(args) == 2
 
                 x, y = args
-                # if !(typeof(x) in (T, SparseMatrixCSC{T,Int}))
-                #     error(
-                #         "Convex.jl internal error: unexpected type $(typeof(x)) for first argument of operation $op of with sign=$sign",
-                #     )
-                # end
-                # if !(typeof(y) in (SparseTape{T}, Vector{T}))
-                #     error(
-                #         "Convex.jl internal error: unexpected type $(typeof(y)) for second argument of operation $op of with sign=$sign",
-                #     )
-                # end
+                if !(typeof(x) in (T, SPARSE_MATRIX{T}))
+                    error(
+                        "Convex.jl internal error: unexpected type $(typeof(x)) for first argument of operation $op of with sign=$sign",
+                    )
+                end
+                if !(typeof(y) in (SparseTape{T}, SPARSE_VECTOR{T}))
+                    error(
+                        "Convex.jl internal error: unexpected type $(typeof(y)) for second argument of operation $op of with sign=$sign",
+                    )
+                end
             end
         else
             # Everything should be either a
-            # SparseTape{T} or Vector{T}
+            # SparseTape{T} or GBVector{T, T}
             for arg in args
-                if !(typeof(arg) == SparseTape{T} || typeof(arg) == Vector{T} || typeof(arg) == GBVector{T, T})
+                if !(
+                    typeof(arg) == SparseTape{T} ||
+                    typeof(arg) == SPARSE_VECTOR{T}
+                )
                     error("Internal error: unexpected type $(typeof(arg))")
                 end
             end
diff --git a/src/problem_depot/problems/affine.jl b/src/problem_depot/problems/affine.jl
@@ -802,13 +802,13 @@ end
 ) where {T,test}
     dims = [2, 3, 4]
     d = prod(dims)
-    A = rand(ComplexF64, 2, 2)
-    B = rand(ComplexF64, 3, 3)
+    A = rand( 2, 2)
+    B = rand( 3, 3)
     C = rand(ComplexF64, 4, 4)
-    M = kron(A, B, C)
+    M = kron(A, kron(B, C))
     Mt1 = kron(transpose(A), B, C)
     Mt2 = kron(A, transpose(B), C)
-    Mt3 = kron(A, B, transpose(C))
+    Mt3 = kron(A, kron(B, transpose(C)))
 
     Rt1 = ComplexVariable(d, d)
     Rt2 = ComplexVariable(d, d)
diff --git a/src/real_operate.jl b/src/real_operate.jl
diff --git a/src/variable_template.jl b/src/variable_template.jl
diff --git a/test/runtests.jl b/test/runtests.jl
