diff --git a/src/operators.jl b/src/operators.jl index 0ba9ade6df2..883c252e362 100644 --- a/src/operators.jl +++ b/src/operators.jl @@ -56,13 +56,13 @@ Base.:-(lhs::AbstractVariableRef, rhs::Number) = (+)(-rhs,lhs) Base.:*(lhs::AbstractVariableRef, rhs::Number) = (*)(rhs,lhs) Base.:/(lhs::AbstractVariableRef, rhs::Number) = (*)(1.0/rhs,lhs) # AbstractVariableRef--AbstractVariableRef -Base.:+(lhs::V, rhs::V) where V <: AbstractVariableRef = GenericAffExpr(0.0, lhs => 1.0, rhs => 1.0) -Base.:-(lhs::V, rhs::V) where V <: AbstractVariableRef = GenericAffExpr(0.0, lhs => 1.0, rhs => -1.0) -function Base.:*(lhs::V, rhs::V) where V <: AbstractVariableRef +Base.:+(lhs::V, rhs::V) where {V <: AbstractVariableRef} = GenericAffExpr(0.0, lhs => 1.0, rhs => 1.0) +Base.:-(lhs::V, rhs::V) where {V <: AbstractVariableRef} = GenericAffExpr(0.0, lhs => 1.0, rhs => -1.0) +function Base.:*(lhs::V, rhs::V) where {V <: AbstractVariableRef} GenericQuadExpr(GenericAffExpr{Float64, V}(), UnorderedPair(lhs, rhs) => 1.0) end # AbstractVariableRef--GenericAffExpr -function Base.:+(lhs::V, rhs::GenericAffExpr{C,V}) where {C, V <: AbstractVariableRef} +function Base.:+(lhs::V, rhs::GenericAffExpr{C, V}) where {C, V <: AbstractVariableRef} # For the variables to have the proper order in the result, we need to add the lhs first. result = zero(rhs) result.constant = rhs.constant @@ -86,11 +86,11 @@ function Base.:-(lhs::V, rhs::GenericAffExpr{C,V}) where {C,V <: AbstractVariabl return result end -function Base.:*(lhs::V, rhs::GenericAffExpr{C,V}) where {C, V <: AbstractVariableRef} +function Base.:*(lhs::V, rhs::GenericAffExpr{C, V}) where {C, V <: AbstractVariableRef} if !iszero(rhs.constant) - result = GenericQuadExpr{C,V}(lhs*rhs.constant) + result = GenericQuadExpr{C, V}(lhs*rhs.constant) else - result = zero(GenericQuadExpr{C,V}) + result = zero(GenericQuadExpr{C, V}) end for (coef, var) in linear_terms(rhs) add_to_expression!(result, coef, lhs, var) @@ -116,7 +116,7 @@ Base.:+(lhs::GenericAffExpr, rhs::Number) = (+)(+rhs,lhs) Base.:-(lhs::GenericAffExpr, rhs::Number) = (+)(-rhs,lhs) Base.:*(lhs::GenericAffExpr, rhs::Number) = (*)(rhs,lhs) Base.:/(lhs::GenericAffExpr, rhs::Number) = map_coefficients(c -> c/rhs, lhs) -function Base.:^(lhs::Union{AbstractVariableRef,GenericAffExpr}, rhs::Integer) +function Base.:^(lhs::Union{AbstractVariableRef, GenericAffExpr}, rhs::Integer) if rhs == 2 return lhs*lhs elseif rhs == 1 @@ -127,21 +127,21 @@ function Base.:^(lhs::Union{AbstractVariableRef,GenericAffExpr}, rhs::Integer) error("Only exponents of 0, 1, or 2 are currently supported. Are you trying to build a nonlinear problem? Make sure you use @NLconstraint/@NLobjective.") end end -Base.:^(lhs::Union{AbstractVariableRef,GenericAffExpr}, rhs::Number) = error("Only exponents of 0, 1, or 2 are currently supported. Are you trying to build a nonlinear problem? Make sure you use @NLconstraint/@NLobjective.") +Base.:^(lhs::Union{AbstractVariableRef, GenericAffExpr}, rhs::Number) = error("Only exponents of 0, 1, or 2 are currently supported. Are you trying to build a nonlinear problem? Make sure you use @NLconstraint/@NLobjective.") # GenericAffExpr--AbstractVariableRef -function Base.:+(lhs::GenericAffExpr{C,V}, rhs::V) where {C, V <: AbstractVariableRef} +function Base.:+(lhs::GenericAffExpr{C, V}, rhs::V) where {C, V <: AbstractVariableRef} return add_to_expression!(copy(lhs), one(C), rhs) end -function Base.:-(lhs::GenericAffExpr{C,V}, rhs::V) where {C, V <: AbstractVariableRef} +function Base.:-(lhs::GenericAffExpr{C, V}, rhs::V) where {C, V <: AbstractVariableRef} return add_to_expression!(copy(lhs), -one(C), rhs) end # Don't fall back on AbstractVariableRef*GenericAffExpr to preserve lhs/rhs # consistency (appears in printing). -function Base.:*(lhs::GenericAffExpr{C,V}, rhs::V) where {C, V <: AbstractVariableRef} +function Base.:*(lhs::GenericAffExpr{C, V}, rhs::V) where {C, V <: AbstractVariableRef} if !iszero(lhs.constant) - result = GenericQuadExpr{C,V}(lhs.constant*rhs) + result = GenericQuadExpr{C, V}(lhs.constant*rhs) else - result = zero(GenericQuadExpr{C,V}) + result = zero(GenericQuadExpr{C, V}) end for (coef, var) in linear_terms(lhs) add_to_expression!(result, coef, var, rhs) @@ -150,7 +150,7 @@ function Base.:*(lhs::GenericAffExpr{C,V}, rhs::V) where {C, V <: AbstractVariab end Base.:/(lhs::GenericAffExpr, rhs::AbstractVariableRef) = error("Cannot divide affine expression by a variable") # AffExpr--AffExpr -function Base.:+(lhs::GenericAffExpr{C,V}, rhs::GenericAffExpr{C,V}) where {C,V<:JuMPTypes} +function Base.:+(lhs::GenericAffExpr{C, V}, rhs::GenericAffExpr{C, V}) where {C, V<:JuMPTypes} if length(linear_terms(lhs)) > 50 || length(linear_terms(rhs)) > 50 if length(linear_terms(lhs)) > 1 operator_warn(owner_model(first(linear_terms(lhs))[2])) @@ -163,7 +163,7 @@ function Base.:+(lhs::GenericAffExpr{C,V}, rhs::GenericAffExpr{C,V}) where {C,V< return GenericAffExpr(lhs.constant + rhs.constant, result_terms) end -function Base.:-(lhs::GenericAffExpr{C,V}, rhs::GenericAffExpr{C,V}) where {C,V<:JuMPTypes} +function Base.:-(lhs::GenericAffExpr{C, V}, rhs::GenericAffExpr{C, V}) where {C, V<:JuMPTypes} result = copy(lhs) result.constant -= rhs.constant sizehint!(result, length(linear_terms(lhs)) + length(linear_terms(rhs))) @@ -173,8 +173,8 @@ function Base.:-(lhs::GenericAffExpr{C,V}, rhs::GenericAffExpr{C,V}) where {C,V< return result end -function Base.:*(lhs::GenericAffExpr{C,V}, rhs::GenericAffExpr{C,V}) where {C,V<:JuMPTypes} - result = zero(GenericQuadExpr{C,V}) +function Base.:*(lhs::GenericAffExpr{C, V}, rhs::GenericAffExpr{C, V}) where {C, V<:JuMPTypes} + result = zero(GenericQuadExpr{C, V}) lhs_length = length(linear_terms(lhs)) rhs_length = length(linear_terms(rhs)) @@ -267,8 +267,8 @@ function Base.:-(q1::GenericQuadExpr, q2::GenericQuadExpr) return result end -Base.:(==)(lhs::GenericAffExpr,rhs::GenericAffExpr) = (lhs.terms == rhs.terms) && (lhs.constant == rhs.constant) -Base.:(==)(lhs::GenericQuadExpr,rhs::GenericQuadExpr) = (lhs.terms == rhs.terms) && (lhs.aff == rhs.aff) +Base.:(==)(lhs::GenericAffExpr, rhs::GenericAffExpr) = (lhs.terms == rhs.terms) && (lhs.constant == rhs.constant) +Base.:(==)(lhs::GenericQuadExpr, rhs::GenericQuadExpr) = (lhs.terms == rhs.terms) && (lhs.aff == rhs.aff) ############################################################################# # Helpers to initialize memory for GenericAffExpr/GenericQuadExpr @@ -301,7 +301,7 @@ function Base.sum(array::AbstractArray{<:AbstractVariableRef}) return result_expression end # TODO: Specialize for iterables. -function Base.sum(affs::AbstractArray{T}) where T<:GenericAffExpr +function Base.sum(affs::AbstractArray{T}) where {T <: GenericAffExpr} new_aff = zero(T) for aff in affs add_to_expression!(new_aff, aff) @@ -313,18 +313,18 @@ end # for scalars, so instead of defining them one-by-one, we will # fallback to the multiplication operator Compat.LinearAlgebra.dot(lhs::JuMPTypes, rhs::JuMPTypes) = lhs*rhs -Compat.LinearAlgebra.dot(lhs::JuMPTypes, rhs::Number) = lhs*rhs -Compat.LinearAlgebra.dot(lhs::Number, rhs::JuMPTypes) = lhs*rhs +Compat.LinearAlgebra.dot(lhs::JuMPTypes, rhs::Number) = lhs*rhs +Compat.LinearAlgebra.dot(lhs::Number, rhs::JuMPTypes) = lhs*rhs -Compat.dot(lhs::AbstractVector{T},rhs::AbstractVector{S}) where {T<:JuMPTypes,S<:JuMPTypes} = _dot(lhs,rhs) -Compat.dot(lhs::AbstractVector{T},rhs::AbstractVector{S}) where {T<:JuMPTypes,S} = _dot(lhs,rhs) -Compat.dot(lhs::AbstractVector{T},rhs::AbstractVector{S}) where {T,S<:JuMPTypes} = _dot(lhs,rhs) +Compat.dot(lhs::AbstractVector{T}, rhs::AbstractVector{S}) where {T <: JuMPTypes, S <: JuMPTypes} = _dot(lhs,rhs) +Compat.dot(lhs::AbstractVector{T}, rhs::AbstractVector{S}) where {T <: JuMPTypes, S} = _dot(lhs,rhs) +Compat.dot(lhs::AbstractVector{T}, rhs::AbstractVector{S}) where {T, S <: JuMPTypes} = _dot(lhs,rhs) -Compat.dot(lhs::AbstractArray{T,N},rhs::AbstractArray{S,N}) where {T<:JuMPTypes,S,N} = _dot(lhs,rhs) -Compat.dot(lhs::AbstractArray{T,N},rhs::AbstractArray{S,N}) where {T<:JuMPTypes,S<:JuMPTypes,N} = _dot(lhs,rhs) -Compat.dot(lhs::AbstractArray{T,N},rhs::AbstractArray{S,N}) where {T,S<:JuMPTypes,N} = _dot(lhs,rhs) +Compat.dot(lhs::AbstractArray{T, N}, rhs::AbstractArray{S, N}) where {T <: JuMPTypes, S, N} = _dot(lhs,rhs) +Compat.dot(lhs::AbstractArray{T, N}, rhs::AbstractArray{S, N}) where {T <: JuMPTypes, S <: JuMPTypes, N} = _dot(lhs,rhs) +Compat.dot(lhs::AbstractArray{T, N}, rhs::AbstractArray{S, N}) where {T, S <: JuMPTypes, N} = _dot(lhs,rhs) -function _dot(lhs::AbstractArray{T}, rhs::AbstractArray{S}) where {T,S} +function _dot(lhs::AbstractArray{T}, rhs::AbstractArray{S}) where {T, S} size(lhs) == size(rhs) || error("Incompatible dimensions") ret = zero(one(T)*one(S)) for (x,y) in zip(lhs,rhs) @@ -337,18 +337,18 @@ end # A bunch of operator junk to make matrix multiplication and friends act # reasonably sane with JuMP types -Base.promote_rule(V::Type{<:AbstractVariableRef},R::Type{<:Real} ) = GenericAffExpr{Float64, V} -Base.promote_rule(V::Type{<:AbstractVariableRef}, ::Type{<:GenericAffExpr{T}} ) where T = GenericAffExpr{T, V} -Base.promote_rule(V::Type{<:AbstractVariableRef}, ::Type{<:GenericQuadExpr{T}} ) where T = GenericQuadExpr{T, V} -Base.promote_rule(::Type{GenericAffExpr{S, V}}, R::Type{<:Real} ) where {S, V} = GenericAffExpr{promote_type(S, R), V} +Base.promote_rule(V::Type{<:AbstractVariableRef}, R::Type{<:Real}) = GenericAffExpr{Float64, V} +Base.promote_rule(V::Type{<:AbstractVariableRef}, ::Type{<:GenericAffExpr{T}}) where {T} = GenericAffExpr{T, V} +Base.promote_rule(V::Type{<:AbstractVariableRef}, ::Type{<:GenericQuadExpr{T}}) where {T} = GenericQuadExpr{T, V} +Base.promote_rule(::Type{GenericAffExpr{S, V}}, R::Type{<:Real}) where {S, V} = GenericAffExpr{promote_type(S, R), V} Base.promote_rule(::Type{<:GenericAffExpr{S, V}}, ::Type{<:GenericQuadExpr{T, V}}) where {S, T, V} = GenericQuadExpr{promote_type(S, T), V} -Base.promote_rule(::Type{GenericQuadExpr{S, V}}, R::Type{<:Real} ) where {S, V} = GenericQuadExpr{promote_type(S, R), V} +Base.promote_rule(::Type{GenericQuadExpr{S, V}}, R::Type{<:Real}) where {S, V} = GenericQuadExpr{promote_type(S, R), V} Base.transpose(x::AbstractJuMPScalar) = x # Can remove the following code once == overloading is removed -function Compat.LinearAlgebra.issymmetric(x::Matrix{T}) where T<:JuMPTypes +function Compat.LinearAlgebra.issymmetric(x::Matrix{T}) where {T <: JuMPTypes} (n = size(x,1)) == size(x,2) || return false for i in 1:n, j in (i+1):n isequal(x[i,j], x[j,i]) || return false @@ -360,204 +360,168 @@ end one_indexed(A) = all(x -> isa(x, Base.OneTo), Compat.axes(A)) function Compat.LinearAlgebra.diagm(x::AbstractVector{<:AbstractVariableRef}) @assert one_indexed(x) # Base.diagm doesn't work for non-one-indexed arrays in general. - diagm(0=>copyto!(similar(x, GenericAffExpr{Float64,eltype(x)}), x)) -end - -############### -# The _multiply!(buf,y,z) adds the results of y*z into the buffer buf. No bounds/size -# checks are performed; it is expected that the caller has done this, has ensured -# that the eltype of buf is appropriate, and has zeroed the elements of buf (if desired). - -function _multiply!(ret::Array{T}, lhs::Array, rhs::Array) where T<:JuMPTypes - m, n = size(lhs,1), size(lhs,2) - r, s = size(rhs,1), size(rhs,2) - for i ∈ 1:m, j ∈ 1:s - q = ret[i,j] - _sizehint_expr!(q, n) - for k ∈ 1:n - tmp = convert(T, lhs[i,k]*rhs[k,j]) - add_to_expression!(q, tmp) - end + diagm(0=>copyto!(similar(x, GenericAffExpr{Float64, eltype(x)}), x)) +end + +############################################################################### +# Matrix/Vector Arithmetic with JuMP eltypes +############################################################################### + +############################################################################### +# convenience/utility definitions + +const GenericAffOrQuadExpr = Union{GenericAffExpr, GenericQuadExpr} + +densify_with_jump_eltype(x::AbstractMatrix) = convert(Matrix, x) + +function densify_with_jump_eltype(x::SparseMatrixCSC{V}) where {V <: AbstractVariableRef} + return convert(Matrix{GenericAffExpr{Float64, V}}, x) +end + +# See https://github.com/JuliaLang/julia/pull/18218. +_A_mul_B_eltype(::Type{R}, ::Type{S}) where {R, S} = typeof(one(R) * one(S) + one(R) * one(S)) + +_A_mul_B_ret_dims(A::AbstractMatrix, B::AbstractVector) = (size(A, 1),) +_A_mul_B_ret_dims(A::AbstractMatrix, B::AbstractMatrix) = (size(A, 1), size(B, 2)) + +# Don't do size checks here; defer that to `_A_mul_B(A, B)`. +function _A_mul_B_ret(A, B, dims...) + T = _A_mul_B_eltype(eltype(A), eltype(B)) + ret = Array{T}(undef, _A_mul_B_ret_dims(A, B)) + return fill_with_zeros!(ret, T) +end + +if VERSION < v"0.7-" + _At_mul_B_ret_dims(A::AbstractMatrix, B::AbstractVector) = (size(A, 2),) + _At_mul_B_ret_dims(A::AbstractMatrix, B::AbstractMatrix) = (size(A, 2), size(B, 2)) + function _At_mul_B_ret(A, B, dims...) + T = _A_mul_B_eltype(eltype(A), eltype(B)) + ret = Array{T}(undef, _At_mul_B_ret_dims(A, B)) + return fill_with_zeros!(ret, T) end - ret end -# this computes lhs.'*rhs and places it in ret -function _multiplyt!(ret::Array{T}, lhs::Array, rhs::Array) where T<:JuMPTypes - m, n = size(lhs,2), size(lhs,1) # transpose - r, s = size(rhs,1), size(rhs,2) - for i ∈ 1:m, j ∈ 1:s - q = ret[i,j] - _sizehint_expr!(q, n) - for k ∈ 1:n - tmp = convert(T, lhs[k,i]*rhs[k,j]) # transpose +function fill_with_zeros!(A, ::Type{T}) where {T} + for I in eachindex(A) + A[I] = zero(T) + end + return A +end + +############################################################################### +# `_A_mul_B!(ret, A, B)` adds the result of `A*B` into the buffer `ret`. We roll our own +# matmul here (instead of using Julia's generic fallbacks) because doing so allows us to +# accumulate the expressions for the inner loops in-place. Additionally, Julia's generic +# fallbacks can be finnicky when your array elements aren't `<:Number`. +# +# No bounds/size checks are performed; it is expected that the caller has done this, has +# ensured that the eltype of `ret` is appropriate, and has zeroed the elements of `ret` (if +# desired). + +function _A_mul_B!(ret::AbstractArray{T}, A, B) where {T <: JuMPTypes} + for i ∈ 1:size(A, 1), j ∈ 1:size(B, 2) + q = ret[i, j] + _sizehint_expr!(q, size(A, 2)) + for k ∈ 1:size(A, 2) + tmp = convert(T, A[i, k] * B[k, j]) add_to_expression!(q, tmp) end end ret end -function _multiply!(ret::Array{T}, lhs::SparseMatrixCSC, rhs::Array) where T<:Union{GenericAffExpr,GenericQuadExpr} - nzv = nonzeros(lhs) - rv = rowvals(lhs) - for col ∈ 1:lhs.n +function _A_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::SparseMatrixCSC, B) + nzv = nonzeros(A) + rv = rowvals(A) + for col ∈ 1:size(A, 2) for k ∈ 1:size(ret, 2) - for j ∈ nzrange(lhs, col) - add_to_expression!(ret[rv[j], k], nzv[j] * rhs[col,k]) + for j ∈ nzrange(A, col) + add_to_expression!(ret[rv[j], k], nzv[j] * B[col, k]) end end end ret end -# this computes lhs.'*rhs and places it in ret -function _multiplyt!(ret::Array{T}, lhs::SparseMatrixCSC, rhs::Array) where T<:Union{GenericAffExpr,GenericQuadExpr} - _multiply!(ret, transpose(lhs), rhs) # TODO fully implement -end - -function _multiply!(ret::Array{T}, lhs::Matrix, rhs::SparseMatrixCSC) where T<:Union{GenericAffExpr,GenericQuadExpr} - rowval = rowvals(rhs) - nzval = nonzeros(rhs) - for multivec_row in 1:size(lhs,1) - for col ∈ 1:rhs.n - idxset = nzrange(rhs, col) +function _A_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::AbstractMatrix, B::SparseMatrixCSC) + rowval = rowvals(B) + nzval = nonzeros(B) + for multivec_row in 1:size(A, 1) + for col ∈ 1:size(B, 2) + idxset = nzrange(B, col) q = ret[multivec_row, col] _sizehint_expr!(q, length(idxset)) for k ∈ idxset - add_to_expression!(q, lhs[multivec_row, rowval[k]] * nzval[k]) + add_to_expression!(q, A[multivec_row, rowval[k]] * nzval[k]) end end end ret end -# this computes lhs.'*rhs and places it in ret -function _multiplyt!(ret::Array{T}, lhs::Matrix, rhs::SparseMatrixCSC) where T<:Union{GenericAffExpr,GenericQuadExpr} - rowval = rowvals(rhs) - nzval = nonzeros(rhs) - for multivec_row ∈ 1:size(lhs,2) # transpose - for col ∈ 1:rhs.n - idxset = nzrange(rhs, col) - q = ret[multivec_row, col] - _sizehint_expr!(q, length(idxset)) - for k ∈ idxset - add_to_expression!(q, lhs[rowval[k], multivec_row] * nzval[k]) # transpose - end - end - end +# TODO: Implement sparse * sparse code as in base/sparse/linalg.jl (spmatmul). +function _A_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::SparseMatrixCSC, B::SparseMatrixCSC) + return _A_mul_B!(ret, A, densify_with_jump_eltype(B)) +end + +function _A_mul_B(A, B) + size(A, 2) == size(B, 1) || error("Incompatible sizes") + ret = _A_mul_B_ret(A, B) + _A_mul_B!(ret, A, B) ret end -# See https://github.com/JuliaLang/julia/issues/27015 -function Base.Matrix(S::SparseMatrixCSC{V}) where V<:AbstractVariableRef - A = zeros(GenericAffExpr{Float64, V}, S.m, S.n) - for Sj in 1:S.n - for Sk in nzrange(S, Sj) - Si = S.rowval[Sk] - Sv = S.nzval[Sk] - A[Si, Sj] = Sv +############################################################################### +# `_At_mul_B!(ret, A, B)` stores the result of `Aᵀ * B` into the buffer `ret`. We roll our +# own version here (instead of working with Julia's generic fallbacks) for the same reasons +# as above. + +function _At_mul_B!(ret::AbstractArray{T}, A, B) where {T <: JuMPTypes} + for i ∈ 1:size(A, 2), j ∈ 1:size(B, 2) + q = ret[i, j] + _sizehint_expr!(q, size(A, 1)) + for k ∈ 1:size(A, 1) + tmp = convert(T, A[k, i] * B[k, j]) # transpose + add_to_expression!(q, tmp) end end - return A + ret end -# TODO: implement sparse * sparse code as in base/sparse/linalg.jl (spmatmul) -_multiply!(ret::AbstractArray{T}, lhs::SparseMatrixCSC, rhs::SparseMatrixCSC) where {T<:JuMPTypes} = _multiply!(ret, lhs, Matrix(rhs)) -_multiplyt!(ret::AbstractArray{T}, lhs::SparseMatrixCSC, rhs::SparseMatrixCSC) where {T<:JuMPTypes} = _multiplyt!(ret, lhs, Matrix(rhs)) -function Base.:*(A::Union{Matrix{T}, SparseMatrixCSC{T}}, - x::Union{Matrix, Vector, SparseMatrixCSC} - ) where {T <: JuMPTypes} - return _matmul(A, x) -end -function Base.:*(A::Union{Matrix{T}, SparseMatrixCSC{T}}, - x::Union{Matrix{R}, Vector{R}, SparseMatrixCSC{R}} - ) where {T <: JuMPTypes, R <: JuMPTypes} - return _matmul(A, x) -end -function Base.:*(A::Union{Matrix, SparseMatrixCSC}, - x::Union{Matrix{T}, Vector{T}, SparseMatrixCSC{T}} - ) where {T <: JuMPTypes} - return _matmul(A, x) -end - -if VERSION >= v"0.7-" - # This is a stopgap solution to get (most) tests passing on Julia 0.7. A lot - # of cases probably still don't work. (Like A * A where A is a sparse matrix - # of a JuMP type). This code needs a big refactor. - function Base.:*(adjA::Adjoint{<:JuMPTypes,<:SparseMatrixCSC}, - x::Vector) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, adjoint(A), x) - end - function Base.:*(adjA::Adjoint{<:Any,<:SparseMatrixCSC}, - x::Vector{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, adjoint(A), x) - end - function Base.:*(adjA::Adjoint{<:JuMPTypes,<:SparseMatrixCSC}, - x::Vector{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, adjoint(A), x) - end - function Base.:*(adjA::Transpose{<:JuMPTypes,<:SparseMatrixCSC}, - x::Vector) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, transpose(A), x) - end - function Base.:*(adjA::Transpose{<:Any,<:SparseMatrixCSC}, - x::Vector{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, transpose(A), x) - end - function Base.:*(adjA::Transpose{<:JuMPTypes,<:SparseMatrixCSC}, - x::Vector{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, transpose(A), x) - end - # Matrix versions. - function Base.:*(adjA::Adjoint{<:JuMPTypes,<:SparseMatrixCSC}, - x::Matrix) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, adjoint(A), x) - end - function Base.:*(adjA::Adjoint{<:Any,<:SparseMatrixCSC}, - x::Matrix{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, adjoint(A), x) - end - function Base.:*(adjA::Adjoint{<:JuMPTypes,<:SparseMatrixCSC}, - x::Matrix{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, adjoint(A), x) - end - function Base.:*(adjA::Transpose{<:JuMPTypes,<:SparseMatrixCSC}, - x::Matrix) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, transpose(A), x) - end - function Base.:*(adjA::Transpose{<:Any,<:SparseMatrixCSC}, - x::Matrix{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, transpose(A), x) - end - function Base.:*(adjA::Transpose{<:JuMPTypes,<:SparseMatrixCSC}, - x::Matrix{<:JuMPTypes}) - A = adjA.parent - ret = _return_arrayt(A, x) - return mul!(ret, transpose(A), x) - end - # mul! is adapted from upstream Julia. +if VERSION < v"0.7-" + function _At_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::SparseMatrixCSC, B) + _A_mul_B!(ret, transpose(A), B) # TODO Fully implement this. + end + function _At_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::AbstractMatrix, B::SparseMatrixCSC) + rowval = rowvals(B) + nzval = nonzeros(B) + for multivec_row ∈ 1:size(A, 2) # transpose + for col ∈ 1:size(B, 2) + idxset = nzrange(B, col) + q = ret[multivec_row, col] + _sizehint_expr!(q, length(idxset)) + for k ∈ idxset + add_to_expression!(q, A[rowval[k], multivec_row] * nzval[k]) # transpose + end + end + end + ret + end + function _At_mul_B(A, B) + size(A, 1) == size(B, 1) || error("Incompatible sizes") + ret = _At_mul_B_ret(A, B) + _At_mul_B!(ret, A, B) + ret + end +else + function _At_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::SparseMatrixCSC, B) + _A_mul_B!(ret, copy(transpose(A)), B) # TODO fully implement + end + function _At_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::AbstractMatrix, B::SparseMatrixCSC) + _A_mul_B!(ret, transpose(A), B) + end + # This method of `_At_mul_B!` is adapted from upstream Julia. Note that we + # confuse transpose with adjoint because they're the same for all JuMP types. #= > Copyright (c) 2009-2018: Jeff Bezanson, Stefan Karpinski, Viral B. Shah, > and other contributors: @@ -583,113 +547,165 @@ if VERSION >= v"0.7-" > OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION > WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. =# - # We confuse transpose with adjoint because they're the same for all JuMP - # types. Note this doesn't extend the LinearAlgebra version. It assumes - # that C is already filled with zeros. - function mul!(C::StridedVecOrMat, - adjA::Union{Adjoint{<:Any,<:SparseMatrixCSC}, - Transpose{<:Any,<:SparseMatrixCSC}}, - B::StridedVecOrMat) - A = adjA.parent - A.n == size(C, 1) || throw(DimensionMismatch()) + function _At_mul_B!(ret::StridedVecOrMat{<:GenericAffOrQuadExpr}, A::SparseMatrixCSC, B::StridedVecOrMat) + A.n == size(ret, 1) || throw(DimensionMismatch()) A.m == size(B, 1) || throw(DimensionMismatch()) - size(B, 2) == size(C, 2) || throw(DimensionMismatch()) + size(B, 2) == size(ret, 2) || throw(DimensionMismatch()) nzv = A.nzval rv = A.rowval - # C is already filled with zeros by _return_arrayt. - for k = 1:size(C, 2) + # ret is already filled with zeros by _return_arrayt. + for k = 1:size(ret, 2) @inbounds for col = 1:A.n - tmp = zero(eltype(C)) + tmp = zero(eltype(ret)) for j = A.colptr[col]:(A.colptr[col + 1] - 1) - tmp += adjoint(nzv[j])*B[rv[j],k] + tmp += adjoint(nzv[j]) * B[rv[j],k] end - C[col,k] += tmp + ret[col,k] += tmp end end - C + ret + end + function _At_mul_B(A, B) + size(A, 1) == size(B, 1) || error("Incompatible sizes") + ret = _A_mul_B_ret(transpose(A), B) + _At_mul_B!(ret, A, B) + ret end -else - _multiply!(ret, lhs, rhs) = A_mul_B!(ret, lhs, rhs) - _multiplyt!(ret, lhs, rhs) = At_mul_B!(ret, lhs, rhs) - - # these methods are called when one does A.'*v or A'*v respectively - Base.At_mul_B(A::Union{Matrix{T},SparseMatrixCSC{T}}, x::Union{Matrix, Vector, SparseMatrixCSC}) where {T<:JuMPTypes} = _matmult(A, x) - Base.At_mul_B(A::Union{Matrix{T},SparseMatrixCSC{T}}, x::Union{Matrix{R}, Vector{R}, SparseMatrixCSC{R}}) where {T<:JuMPTypes,R<:JuMPTypes} = _matmult(A, x) - Base.At_mul_B(A::Union{Matrix,SparseMatrixCSC}, x::Union{Matrix{T}, Vector{T}, SparseMatrixCSC{T}}) where {T<:JuMPTypes} = _matmult(A, x) - # these methods are the same as above since complex numbers are not implemented in JuMP - Base.Ac_mul_B(A::Union{Matrix{T},SparseMatrixCSC{T}}, x::Union{Matrix, Vector, SparseMatrixCSC}) where {T<:JuMPTypes} = _matmult(A, x) - Base.Ac_mul_B(A::Union{Matrix{T},SparseMatrixCSC{T}}, x::Union{Matrix{R}, Vector{R}, SparseMatrixCSC{R}}) where {T<:JuMPTypes,R<:JuMPTypes} = _matmult(A, x) - Base.Ac_mul_B(A::Union{Matrix,SparseMatrixCSC}, x::Union{Matrix{T}, Vector{T}, SparseMatrixCSC{T}}) where {T<:JuMPTypes} = _matmult(A, x) -end - -function _matmul(A, x) - m, n = size(A,1), size(A,2) - r, s = size(x,1), size(x,2) - n == r || error("Incompatible sizes") - ret = _return_array(A, x) - _multiply!(ret, A, x) - ret end -function _matmult(A, x) - m, n = size(A,2), size(A,1) # transpose - r, s = size(x,1), size(x,2) - n == r || error("Incompatible sizes") - ret = _return_arrayt(A, x) - _multiplyt!(ret, A, x) - ret +# TODO: Implement sparse * sparse code as in base/sparse/linalg.jl (spmatmul). +function _At_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::SparseMatrixCSC, B::SparseMatrixCSC) + return _At_mul_B!(ret, A, densify_with_jump_eltype(B)) end -# See https://github.com/JuliaLang/julia/pull/18218 -_matprod_type(R, S) = typeof(one(R) * one(S) + one(R) * one(S)) -# Don't do size checks here in _return_array, defer that to (*) -_return_array(A::AbstractMatrix{R}, x::AbstractVector{S}) where {R,S} = _fillwithzeros(Array{_matprod_type(R,S)}(undef,size(A,1))) -_return_array(A::AbstractMatrix{R}, x::AbstractMatrix{S}) where {R,S} = _fillwithzeros(Array{_matprod_type(R,S)}(undef,size(A,1), size(x,2))) -# these are for transpose return matrices -_return_arrayt(A::AbstractMatrix{R}, x::AbstractVector{S}) where {R,S} = _fillwithzeros(Array{_matprod_type(R,S)}(undef,size(A,2))) -_return_arrayt(A::AbstractMatrix{R}, x::AbstractMatrix{S}) where {R,S} = _fillwithzeros(Array{_matprod_type(R,S)}(undef,size(A,2), size(x, 2))) +############################################################################### +# Interception of Base's matrix/vector arithmetic machinery -# helper so we don't fill the buffer array with the same object -function _fillwithzeros(arr::AbstractArray{T}) where T - for I in eachindex(arr) - arr[I] = zero(T) - end - arr +# TODO: Intercepting "externally owned" method calls by dispatching on type parameters +# (rather than outermost wrapper type) is generally bad practice, but refactoring this code +# to use a different mechanism would be a lot of work. In the future, this interception code +# would be more easily/robustly replaced by using a tool like +# https://github.com/jrevels/Cassette.jl. + +function Base.:*(A::Union{Matrix{<:JuMPTypes}, SparseMatrixCSC{<:JuMPTypes}}, + B::Union{Matrix, Vector, SparseMatrixCSC}) + return _A_mul_B(A, B) end -# Special-case sparse matrix scalar multiplication/division -Base.:*(lhs::Number, rhs::SparseMatrixCSC{T}) where {T<:JuMPTypes} = - SparseMatrixCSC(rhs.m, rhs.n, copy(rhs.colptr), copy(rhs.rowval), lhs .* rhs.nzval) -Base.:*(lhs::JuMPTypes, rhs::SparseMatrixCSC) = - SparseMatrixCSC(rhs.m, rhs.n, copy(rhs.colptr), copy(rhs.rowval), lhs .* rhs.nzval) -Base.:*(lhs::SparseMatrixCSC{T}, rhs::Number) where {T<:JuMPTypes} = - SparseMatrixCSC(lhs.m, lhs.n, copy(lhs.colptr), copy(lhs.rowval), lhs.nzval .* rhs) -Base.:*(lhs::SparseMatrixCSC, rhs::JuMPTypes) = - SparseMatrixCSC(lhs.m, lhs.n, copy(lhs.colptr), copy(lhs.rowval), lhs.nzval .* rhs) -Base.:/(lhs::SparseMatrixCSC{T}, rhs::Number) where {T<:JuMPTypes} = - SparseMatrixCSC(lhs.m, lhs.n, copy(lhs.colptr), copy(lhs.rowval), lhs.nzval ./ rhs) +function Base.:*(A::Union{Matrix{<:JuMPTypes}, SparseMatrixCSC{<:JuMPTypes}}, + B::Union{Matrix{<:JuMPTypes}, Vector{<:JuMPTypes}, SparseMatrixCSC{<:JuMPTypes}}) + return _A_mul_B(A, B) +end +function Base.:*(A::Union{Matrix, SparseMatrixCSC}, + B::Union{Matrix{<:JuMPTypes}, Vector{<:JuMPTypes}, SparseMatrixCSC{<:JuMPTypes}}) + return _A_mul_B(A, B) +end + +if VERSION < v"0.7-" + # These methods are called when one does A.'*B. + function Base.At_mul_B(A::Union{Matrix{T}, SparseMatrixCSC{T}}, + B::Union{Matrix, Vector, SparseMatrixCSC}) where {T <: JuMPTypes} + return _At_mul_B(A, B) + end + + function Base.At_mul_B(A::Union{Matrix{T}, SparseMatrixCSC{T}}, + B::Union{Matrix{R}, Vector{R}, SparseMatrixCSC{R}}) where {T <: JuMPTypes,R <: JuMPTypes} + return _At_mul_B(A, B) + end + + function Base.At_mul_B(A::Union{Matrix, SparseMatrixCSC}, + B::Union{Matrix{T}, Vector{T}, SparseMatrixCSC{T}}) where {T <: JuMPTypes} + return _At_mul_B(A, B) + end + + # These methods are called when one does A'*B. + # These methods are the same as above since complex numbers are not implemented in JuMP. + function Base.Ac_mul_B(A::Union{Matrix{T}, SparseMatrixCSC{T}}, + B::Union{Matrix, Vector, SparseMatrixCSC}) where {T <: JuMPTypes} + return _At_mul_B(A, B) + end + + function Base.Ac_mul_B(A::Union{Matrix{T}, SparseMatrixCSC{T}}, + B::Union{Matrix{R}, Vector{R}, SparseMatrixCSC{R}}) where {T <: JuMPTypes,R <: JuMPTypes} + return _At_mul_B(A, B) + end -for (op,opsymbol) in [(+,:+), (-,:-), (*,:*), (/,:/)] - @eval begin - Base.broadcast(::typeof($op),lhs::Number,rhs::JuMPTypes) = $opsymbol(lhs,rhs) - Base.broadcast(::typeof($op),lhs::JuMPTypes,rhs::Number) = $opsymbol(lhs,rhs) + function Base.Ac_mul_B(A::Union{Matrix, SparseMatrixCSC}, + B::Union{Matrix{T}, Vector{T}, SparseMatrixCSC{T}}) where {T <: JuMPTypes} + return _At_mul_B(A, B) end + +else + # TODO: This is a stopgap solution to get (most) tests passing on Julia 0.7. A lot of + # cases probably still don't work. (Like A * A where A is a sparse matrix of a JuMP + # type). This code needs a big refactor. + Base.:*(A::Adjoint{<:JuMPTypes, <:SparseMatrixCSC}, B::Vector) = _At_mul_B(parent(A), B) + Base.:*(A::Adjoint{<:Any, <:SparseMatrixCSC}, B::Vector{<:JuMPTypes}) = _At_mul_B(parent(A), B) + Base.:*(A::Adjoint{<:JuMPTypes, <:SparseMatrixCSC}, B::Vector{<:JuMPTypes}) = _At_mul_B(parent(A), B) + + Base.:*(A::Transpose{<:JuMPTypes, <:SparseMatrixCSC}, B::Vector) = _At_mul_B(parent(A), B) + Base.:*(A::Transpose{<:Any, <:SparseMatrixCSC}, B::Vector{<:JuMPTypes}) = _At_mul_B(parent(A), B) + Base.:*(A::Transpose{<:JuMPTypes, <:SparseMatrixCSC}, B::Vector{<:JuMPTypes}) = _At_mul_B(parent(A), B) + + Base.:*(A::Adjoint{<:JuMPTypes, <:SparseMatrixCSC}, B::Matrix) = _At_mul_B(parent(A), B) + Base.:*(A::Adjoint{<:Any, <:SparseMatrixCSC}, B::Matrix{<:JuMPTypes}) = _At_mul_B(parent(A), B) + Base.:*(A::Adjoint{<:JuMPTypes, <:SparseMatrixCSC}, B::Matrix{<:JuMPTypes}) = _At_mul_B(parent(A), B) + + Base.:*(A::Transpose{<:JuMPTypes, <:SparseMatrixCSC}, B::Matrix) = _At_mul_B(parent(A), B) + Base.:*(A::Transpose{<:Any, <:SparseMatrixCSC}, B::Matrix{<:JuMPTypes}) = _At_mul_B(parent(A), B) + Base.:*(A::Transpose{<:JuMPTypes, <:SparseMatrixCSC}, B::Matrix{<:JuMPTypes}) = _At_mul_B(parent(A), B) +end + +# Base doesn't define efficient fallbacks for sparse array arithmetic involving +# non-`<:Number` scalar elements, so we define some of these for `<:JuMPType` scalar +# elements here. + +function Base.:*(A::Number, B::SparseMatrixCSC{T}) where {T <: JuMPTypes} + return SparseMatrixCSC(B.m, B.n, copy(B.colptr), copy(B.rowval), A .* B.nzval) +end + +function Base.:*(A::SparseMatrixCSC{T}, B::Number) where {T <: JuMPTypes} + return SparseMatrixCSC(A.m, A.n, copy(A.colptr), copy(A.rowval), A.nzval .* B) end +function Base.:*(A::JuMPTypes, B::SparseMatrixCSC) + return SparseMatrixCSC(B.m, B.n, copy(B.colptr), copy(B.rowval), A .* B.nzval) +end -Base.:+(x::AbstractArray{T}) where {T<:JuMPTypes} = x -function Base.:-(x::AbstractArray{T}) where T<:JuMPTypes +function Base.:*(A::SparseMatrixCSC, B::JuMPTypes) + return SparseMatrixCSC(A.m, A.n, copy(A.colptr), copy(A.rowval), A.nzval .* B) +end + +function Base.:/(A::SparseMatrixCSC{T}, B::Number) where {T <: JuMPTypes} + return SparseMatrixCSC(A.m, A.n, copy(A.colptr), copy(A.rowval), A.nzval ./ B) +end + +Base.:*(x::AbstractArray{T}) where {T <: JuMPTypes} = x + +Base.:+(x::AbstractArray{T}) where {T <: JuMPTypes} = x + +function Base.:-(x::AbstractArray{T}) where {T <: JuMPTypes} ret = similar(x, typeof(-one(T))) for I in eachindex(ret) ret[I] = -x[I] end ret end -Base.:*(x::AbstractArray{T}) where {T<:JuMPTypes} = x + +if VERSION < v"0.7-" + for (op, opsymbol) in [(+, :+), (-, :-), (*, :*), (/, :/)] + @eval begin + Base.broadcast(::typeof($op), x::Number, y::JuMPTypes) = $opsymbol(x, y) + Base.broadcast(::typeof($op), x::JuMPTypes, y::Number) = $opsymbol(x, y) + end + end +end ############################################################################### -# Add nonlinear function fallbacks for JuMP built-in types +# nonlinear function fallbacks for JuMP built-in types +############################################################################### + const op_hint = "Are you trying to build a nonlinear problem? Make sure you use @NLconstraint/@NLobjective." for (func,_) in Calculus.symbolic_derivatives_1arg(), typ in [:AbstractVariableRef,:GenericAffExpr,:GenericQuadExpr] errstr = "$func is not defined for type $typ. $op_hint" @@ -698,11 +714,13 @@ for (func,_) in Calculus.symbolic_derivatives_1arg(), typ in [:AbstractVariableR end end -Base.:*(::T,::S) where {T<:GenericQuadExpr,S<:Union{AbstractVariableRef,GenericAffExpr,GenericQuadExpr}} = +Base.:*(::T, ::S) where {T <: GenericQuadExpr, S <: Union{AbstractVariableRef, GenericAffExpr, GenericQuadExpr}} = error( "*(::$T,::$S) is not defined. $op_hint") Base.:*(lhs::GenericQuadExpr, rhs::GenericQuadExpr) = error( "*(::GenericQuadExpr,::GenericQuadExpr) is not defined. $op_hint") -Base.:*(::S,::T) where {T<:GenericQuadExpr,S<:Union{AbstractVariableRef,GenericAffExpr,GenericQuadExpr}} = +Base.:*(::S, ::T) where {T <: GenericQuadExpr, + S <: Union{AbstractVariableRef, GenericAffExpr, GenericQuadExpr}} = error( "*(::$S,::$T) is not defined. $op_hint") -Base.:/(::S,::T) where {S<:Union{Number,AbstractVariableRef,GenericAffExpr,GenericQuadExpr},T<:Union{AbstractVariableRef,GenericAffExpr,GenericQuadExpr}} = +Base.:/(::S, ::T) where {S <: Union{Number, AbstractVariableRef, GenericAffExpr, GenericQuadExpr}, + T <: Union{AbstractVariableRef, GenericAffExpr, GenericQuadExpr}} = error( "/(::$S,::$T) is not defined. $op_hint") diff --git a/test/operator.jl b/test/operator.jl index d9486a420bc..1590019282d 100644 --- a/test/operator.jl +++ b/test/operator.jl @@ -360,7 +360,7 @@ function operators_test(ModelType::Type{<:JuMP.AbstractModel}, VariableRefType:: @variable(m, X11) @variable(m, X23) X = sparse([1, 2], [1, 3], [X11, X23], 3, 3) # for testing Variable - @test JuMP.isequal_canonical([X11 0. 0.; 0. 0. X23; 0. 0. 0.], @inferred Matrix(X)) + @test JuMP.isequal_canonical([X11 0. 0.; 0. 0. X23; 0. 0. 0.], @inferred JuMP.densify_with_jump_eltype(X)) @variable(m, Xd[1:3, 1:3]) Y = sparse([1, 2], [1, 3], [2X11, 4X23], 3, 3) # for testing GenericAffExpr Yd = [2X11 0 0 @@ -547,9 +547,14 @@ function operators_test(ModelType::Type{<:JuMP.AbstractModel}, VariableRefType:: @test_throws ErrorException A./y @test_throws ErrorException B./y - @test JuMP.isequal_canonical((2 .* x) ./ 3, Matrix((2 .* y) ./ 3)) - @test JuMP.isequal_canonical(2 .* (x ./ 3), Matrix(2 * (y ./ 3))) - @test JuMP.isequal_canonical((x[1,1],) .* A, Matrix((x[1,1],) .* B)) + # TODO: Refactor to avoid calling the internal JuMP function + # `densify_with_jump_eltype`. + z = JuMP.densify_with_jump_eltype((2 .* y) ./ 3) + @test JuMP.isequal_canonical((2 .* x) ./ 3, z) + z = JuMP.densify_with_jump_eltype(2 * (y ./ 3)) + @test JuMP.isequal_canonical(2 .* (x ./ 3), z) + z = JuMP.densify_with_jump_eltype((x[1,1],) .* B) + @test JuMP.isequal_canonical((x[1,1],) .* A, z) end @testset "Vectorized comparisons" begin