diff --git a/src/Utilities/sparse_matrix.jl b/src/Utilities/sparse_matrix.jl index 35905fd40b..91b514df72 100644 --- a/src/Utilities/sparse_matrix.jl +++ b/src/Utilities/sparse_matrix.jl @@ -29,11 +29,11 @@ struct OneBasedIndexing <: AbstractIndexing end _first_index(::ZeroBasedIndexing) = 0 _first_index(::OneBasedIndexing) = 1 -_shift(x, ::ZeroBasedIndexing, ::ZeroBasedIndexing) = x + +_shift(x, ::T, ::T) where {T<:AbstractIndexing} = x _shift(x::Integer, ::ZeroBasedIndexing, ::OneBasedIndexing) = x + 1 -_shift(x::Array{<:Integer}, ::ZeroBasedIndexing, ::OneBasedIndexing) = x .+ 1 _shift(x::Integer, ::OneBasedIndexing, ::ZeroBasedIndexing) = x - 1 -_shift(x, ::OneBasedIndexing, ::OneBasedIndexing) = x +_shift(x::Array{<:Integer}, ::ZeroBasedIndexing, ::OneBasedIndexing) = x .+ 1 """ mutable struct MutableSparseMatrixCSC{Tv,Ti<:Integer,I<:AbstractIndexing} @@ -76,7 +76,8 @@ function MOI.empty!(A::MutableSparseMatrixCSC) resize!(A.colptr, 1) A.colptr[1] = 0 empty!(A.rowval) - return empty!(A.nzval) + empty!(A.nzval) + return end """ @@ -103,7 +104,8 @@ function set_number_of_rows(A::MutableSparseMatrixCSC, num_rows) A.colptr[i] += A.colptr[i-1] end resize!(A.rowval, A.colptr[end]) - return resize!(A.nzval, A.colptr[end]) + resize!(A.nzval, A.colptr[end]) + return end """ @@ -116,61 +118,71 @@ function final_touch(A::MutableSparseMatrixCSC) for i in length(A.colptr):-1:2 A.colptr[i] = _shift(A.colptr[i-1], ZeroBasedIndexing(), A.indexing) end - return A.colptr[1] = _first_index(A.indexing) + A.colptr[1] = _first_index(A.indexing) + return end _variable(term::MOI.ScalarAffineTerm) = term.variable _variable(term::MOI.VectorAffineTerm) = _variable(term.scalar_term) -function _allocate_terms(colptr, index_map, terms) - for term in terms - colptr[index_map[_variable(term)].value+1] += 1 - end -end """ - allocate_terms(A::MutableSparseMatrixCSC, index_map, func) + allocate_terms( + A::MutableSparseMatrixCSC, + index_map, + func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction}, + ) Informs `A` that the terms of the function `func` where the variable indices are mapped with `index_map` will be loaded with [`load_terms`](@ref). The function `func` should be canonicalized, see [`is_canonical`](@ref). """ -function allocate_terms(A::MutableSparseMatrixCSC, index_map, func) - return _allocate_terms(A.colptr, index_map, func.terms) +function allocate_terms( + A::MutableSparseMatrixCSC, + index_map, + func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction}, +) + for term in func.terms + A.colptr[index_map[_variable(term)].value+1] += 1 + end + return end _output_index(::MOI.ScalarAffineTerm) = 1 _output_index(term::MOI.VectorAffineTerm) = term.output_index -_coefficient(term::MOI.ScalarAffineTerm) = term.coefficient -_coefficient(term::MOI.VectorAffineTerm) = _coefficient(term.scalar_term) -function _load_terms(colptr, rowval, nzval, index_map, terms, offset) - for term in terms - ptr = colptr[index_map[_variable(term)].value] += 1 - rowval[ptr] = offset + _output_index(term) - nzval[ptr] = _coefficient(term) - end -end """ - load_terms(A::MutableSparseMatrixCSC, index_map, func, offset) + load_terms( + A::MutableSparseMatrixCSC, + index_map, + func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction}, + offset::Int, + ) Loads the terms of `func` to `A` mapping the variable indices with `index_map`. The `i`th dimension of `func` is loaded at the `(offset + i)`th row of `A`. The function should be allocated first with [`allocate_terms`](@ref). The function `func` should be canonicalized, see [`is_canonical`](@ref). """ -function load_terms(A::MutableSparseMatrixCSC, index_map, func, offset) - return _load_terms( - A.colptr, - A.rowval, - A.nzval, - index_map, - func.terms, - _shift(offset, OneBasedIndexing(), A.indexing), - ) +function load_terms( + A::MutableSparseMatrixCSC, + index_map, + func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction}, + offset::Int, +) + new_offset = _shift(offset, OneBasedIndexing(), A.indexing) + for term in func.terms + ptr = A.colptr[index_map[_variable(term)].value] += 1 + A.rowval[ptr] = new_offset + _output_index(term) + A.nzval[ptr] = MOI.coefficient(term) + end + return end """ - Base.convert(::Type{SparseMatrixCSC{Tv, Ti}}, A::MutableSparseMatrixCSC{Tv, Ti, I}) where {Tv, Ti, I} + Base.convert( + ::Type{SparseMatrixCSC{Tv,Ti}}, + A::MutableSparseMatrixCSC{Tv,Ti,I}, + ) where {Tv,Ti,I} Converts `A` to a `SparseMatrixCSC`. Note that the field `A.nzval` is **not copied** so if `A` is modified after the call of this function, it can still diff --git a/test/Utilities/sparse_matrix.jl b/test/Utilities/sparse_matrix.jl new file mode 100644 index 0000000000..c0b2cad976 --- /dev/null +++ b/test/Utilities/sparse_matrix.jl @@ -0,0 +1,176 @@ +module TestSparseMatrix + +import SparseArrays +using Test + +import MathOptInterface +const MOI = MathOptInterface + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$(name)", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end + return +end + +""" + test_empty_size() + +Julia 1.7 is going to outlaw constructing empty sparse matrices. Test that we +can still do so. +""" +function test_empty_size() + A = MOI.Utilities.MutableSparseMatrixCSC{ + Float64, + Int, + MOI.Utilities.ZeroBasedIndexing, + }() + MOI.empty!(A) + @test A.rowval == Int[] + @test A.nzval == Float64[] + @test A.colptr == Int[0] +end + +function test_empty_large() + A = MOI.Utilities.MutableSparseMatrixCSC{ + Float64, + Int, + MOI.Utilities.ZeroBasedIndexing, + }() + MOI.empty!(A) + MOI.Utilities.add_column(A) + MOI.Utilities.add_column(A) + MOI.Utilities.set_number_of_rows(A, 5) + B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A) + @test size(B) == (5, 2) +end + +function test_VectorAffine_ZeroBased() + A = MOI.Utilities.MutableSparseMatrixCSC{ + Float64, + Int, + MOI.Utilities.ZeroBasedIndexing, + }() + MOI.empty!(A) + x = MOI.VariableIndex.(1:3) + f = MOI.VectorAffineFunction( + vcat( + MOI.VectorAffineTerm.(1, MOI.ScalarAffineTerm.(1.0, x)), + MOI.VectorAffineTerm.(2, MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3])), + ), + [0.5, 1.2], + ) + + index_map = MOI.Utilities.IndexMap() + for i in 1:3 + MOI.Utilities.add_column(A) + index_map[x[i]] = x[i] + end + MOI.Utilities.allocate_terms(A, index_map, f) + MOI.Utilities.set_number_of_rows(A, 2) + MOI.Utilities.load_terms(A, index_map, f, 0) + MOI.Utilities.final_touch(A) + B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A) + @test B == [1.0 1.0 1.0; 0.0 2.0 3.0] + @test A.rowval == [0, 0, 1, 0, 1] + @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0] + @test A.colptr == [0, 1, 3, 5] +end + +function test_VectorAffine_OneBased() + A = MOI.Utilities.MutableSparseMatrixCSC{ + Float64, + Int, + MOI.Utilities.OneBasedIndexing, + }() + MOI.empty!(A) + x = MOI.VariableIndex.(1:3) + f = MOI.VectorAffineFunction( + vcat( + MOI.VectorAffineTerm.(1, MOI.ScalarAffineTerm.(1.0, x)), + MOI.VectorAffineTerm.(2, MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3])), + ), + [0.5, 1.2], + ) + + index_map = MOI.Utilities.IndexMap() + for i in 1:3 + MOI.Utilities.add_column(A) + index_map[x[i]] = x[i] + end + MOI.Utilities.allocate_terms(A, index_map, f) + MOI.Utilities.set_number_of_rows(A, 2) + MOI.Utilities.load_terms(A, index_map, f, 0) + MOI.Utilities.final_touch(A) + B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A) + @test B == [1.0 1.0 1.0; 0.0 2.0 3.0] + @test A.rowval == [1, 1, 2, 1, 2] + @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0] + @test A.colptr == [1, 2, 4, 6] +end + +function test_ScalarAffine_ZeroBased() + A = MOI.Utilities.MutableSparseMatrixCSC{ + Float64, + Int, + MOI.Utilities.ZeroBasedIndexing, + }() + MOI.empty!(A) + x = MOI.VariableIndex.(1:3) + f1 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5) + f2 = + MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3]), 1.2) + index_map = MOI.Utilities.IndexMap() + for i in 1:3 + MOI.Utilities.add_column(A) + index_map[x[i]] = x[i] + end + MOI.Utilities.allocate_terms(A, index_map, f1) + MOI.Utilities.allocate_terms(A, index_map, f2) + MOI.Utilities.set_number_of_rows(A, 2) + MOI.Utilities.load_terms(A, index_map, f1, 0) + MOI.Utilities.load_terms(A, index_map, f2, 1) + MOI.Utilities.final_touch(A) + B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A) + @test B == [1.0 1.0 1.0; 0.0 2.0 3.0] + @test A.rowval == [0, 0, 1, 0, 1] + @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0] + @test A.colptr == [0, 1, 3, 5] +end + +function test_ScalarAffine_OneBased() + A = MOI.Utilities.MutableSparseMatrixCSC{ + Float64, + Int, + MOI.Utilities.OneBasedIndexing, + }() + MOI.empty!(A) + x = MOI.VariableIndex.(1:3) + f1 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5) + f2 = + MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3]), 1.2) + index_map = MOI.Utilities.IndexMap() + for i in 1:3 + MOI.Utilities.add_column(A) + index_map[x[i]] = x[i] + end + MOI.Utilities.allocate_terms(A, index_map, f1) + MOI.Utilities.allocate_terms(A, index_map, f2) + MOI.Utilities.set_number_of_rows(A, 2) + MOI.Utilities.load_terms(A, index_map, f1, 0) + MOI.Utilities.load_terms(A, index_map, f2, 1) + MOI.Utilities.final_touch(A) + B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A) + @test B == [1.0 1.0 1.0; 0.0 2.0 3.0] + @test A.rowval == [1, 1, 2, 1, 2] + @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0] + @test A.colptr == [1, 2, 4, 6] +end + +end + +TestSparseMatrix.runtests()