[Utilities] simplify MutableSparseMatrix to remove final_touch

src/Utilities/sparse_matrix.jl

@@ -27,8 +27,8 @@ allocation-free conversion of [`MutableSparseMatrixCSC`](@ref) to
 """
 struct OneBasedIndexing <: AbstractIndexing end
 
-_first_index(::ZeroBasedIndexing) = 0
-_first_index(::OneBasedIndexing) = 1
+_first_index(::Type{T}, ::ZeroBasedIndexing) where {T} = T(0)
+_first_index(::Type{T}, ::OneBasedIndexing) where {T} = T(1)
 
 _shift(x, ::T, ::T) where {T<:AbstractIndexing} = x
 _shift(x::Integer, ::ZeroBasedIndexing, ::OneBasedIndexing) = x + 1
@@ -43,12 +43,15 @@ _shift(x::Array{<:Integer}, ::ZeroBasedIndexing, ::OneBasedIndexing) = x .+ 1
         colptr::Vector{Ti}
         rowval::Vector{Ti}
         nzval::Vector{Tv}
+        nz_added::Vector{Ti}
     end
 
 Matrix type loading sparse matrices in the Compressed Sparse Column format.
 The indexing used is `indexing`, see [`AbstractIndexing`](@ref). The other
 fields have the same meaning than for `SparseArrays.SparseMatrixCSC` except
-that the indexing is different unless `indexing` is `OneBasedIndexing`.
+that the indexing is different unless `indexing` is `OneBasedIndexing`. In
+addition, `nz_added` is used to cache the number of non-zero terms that have
+been added to each column due to the incremental nature of `load_terms`.
 
 The matrix is loaded in 5 steps:
 1) `MOI.empty!` is called.
@@ -65,8 +68,11 @@ mutable struct MutableSparseMatrixCSC{Tv,Ti<:Integer,I<:AbstractIndexing}
     colptr::Vector{Ti}
     rowval::Vector{Ti}
     nzval::Vector{Tv}
+    nz_added::Vector{Ti}
     function MutableSparseMatrixCSC{Tv,Ti,I}() where {Tv,Ti<:Integer,I}
-        return new{Tv,Ti,I}(I(), 0, 0, [zero(Ti)], Ti[], Tv[])
+        indexing = I()
+        colptr = [_first_index(Ti, indexing)]
+        return new{Tv,Ti,I}(indexing, 0, 0, colptr, Ti[], Tv[], Ti[])
     end
 end
 
@@ -80,43 +86,39 @@ function MOI.empty!(A::MutableSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     A.m = 0
     A.n = 0
     resize!(A.colptr, 1)
-    A.colptr[1] = zero(Ti)
+    A.colptr[1] = _first_index(Ti, A.indexing)
     empty!(A.rowval)
     empty!(A.nzval)
+    empty!(A.nz_added)
     return
 end
 
 function add_column(A::MutableSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     A.n += 1
-    push!(A.colptr, zero(Ti))
+    push!(A.colptr, _first_index(Ti, A.indexing))
+    push!(A.nz_added, Ti(0))
     return
 end
 
 function add_columns(A::MutableSparseMatrixCSC{Tv,Ti}, n) where {Tv,Ti}
-    A.n += n
     for _ in 1:n
-        push!(A.colptr, zero(Ti))
+        add_column(A)
     end
     return
 end
 
 function set_number_of_rows(A::MutableSparseMatrixCSC, num_rows)
     A.m = num_rows
+    offset = _first_index(Int, A.indexing)
     for i in 3:length(A.colptr)
-        A.colptr[i] += A.colptr[i-1]
+        A.colptr[i] += A.colptr[i-1] - offset
     end
-    resize!(A.rowval, A.colptr[end])
-    resize!(A.nzval, A.colptr[end])
+    resize!(A.rowval, A.colptr[end] - offset)
+    resize!(A.nzval, A.colptr[end] - offset)
     return
 end
 
-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
-    A.colptr[1] = _first_index(A.indexing)
-    return
-end
+final_touch(::MutableSparseMatrixCSC) = nothing
 
 _variable(term::MOI.ScalarAffineTerm) = term.variable
 _variable(term::MOI.VectorAffineTerm) = _variable(term.scalar_term)
@@ -141,11 +143,14 @@ function load_terms(
     func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction},
     offset::Int,
 )
-    new_offset = _shift(offset, OneBasedIndexing(), A.indexing)
+    row_offset = _shift(offset, OneBasedIndexing(), A.indexing)
+    zero_offset = 1 - _first_index(Int, A.indexing)
     for term in func.terms
-        ptr = A.colptr[index_map[_variable(term)].value] += 1
-        A.rowval[ptr] = new_offset + _output_index(term)
+        col = index_map[_variable(term)].value
+        ptr = A.colptr[col] + A.nz_added[col] + zero_offset
+        A.rowval[ptr] = row_offset + _output_index(term)
         A.nzval[ptr] = MOI.coefficient(term)
+        A.nz_added[col] += 1
     end
     return
 end

test/Utilities/sparse_matrix.jl

@@ -79,6 +79,11 @@ function test_VectorAffine_ZeroBased()
     @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]
+    MOI.empty!(A)
+    MOI.Utilities.add_column(A)
+    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
+    @test size(C) == (0, 1)
+    return
 end
 
 function test_extract_function()
@@ -145,6 +150,11 @@ function test_VectorAffine_OneBased()
     @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]
+    MOI.empty!(A)
+    MOI.Utilities.add_column(A)
+    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
+    @test size(C) == (0, 1)
+    return
 end
 
 function test_ScalarAffine_ZeroBased()
@@ -174,6 +184,11 @@ function test_ScalarAffine_ZeroBased()
     @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]
+    MOI.empty!(A)
+    MOI.Utilities.add_column(A)
+    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
+    @test size(C) == (0, 1)
+    return
 end
 
 function test_ScalarAffine_OneBased()
@@ -203,6 +218,11 @@ function test_ScalarAffine_OneBased()
     @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]
+    MOI.empty!(A)
+    MOI.Utilities.add_column(A)
+    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
+    @test size(C) == (0, 1)
+    return
 end
 
 end