[Utilities] updates to sparse_matrix.jl

docs/src/submodules/Utilities/reference.md

@@ -95,6 +95,7 @@ Utilities.allocate_terms
 Utilities.set_number_of_rows
 Utilities.load_terms
 Utilities.final_touch
+Utilities.extract_function
 ```
 
 ```@docs

src/Utilities/matrix_of_constraints.jl

@@ -151,6 +151,23 @@ No more modification is allowed unless `MOI.empty!` is called.
 """
 function final_touch end
 
+"""
+    extract_function(coefficients, row::Integer, constant::T) where {T}
+
+Return the `MOI.ScalarAffineFunction{T}` function corresponding to row `row` in
+`coefficients`.
+
+    extract_function(
+        coefficients,
+        rows::UnitRange,
+        constants::Vector{T},
+    ) where{T}
+
+Return the `MOI.VectorAffineFunction{T}` function corresponding to rows `rows`
+in `coefficients`.
+"""
+function extract_function end
+
 ###
 ### Interface for the .constants field
 ###

src/Utilities/sparse_matrix.jl

@@ -91,6 +91,7 @@ function add_column(A::MutableSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     push!(A.colptr, zero(Ti))
     return
 end
+
 function add_columns(A::MutableSparseMatrixCSC{Tv,Ti}, n) where {Tv,Ti}
     A.n += n
     for _ in 1:n
@@ -183,15 +184,18 @@ function _first_in_column(
     idx = searchsortedfirst(view(A.rowval, range), row)
     return get(range, idx, last(range) + 1)
 end
+
 function extract_function(
     A::MutableSparseMatrixCSC{T},
     row::Integer,
     constant::T,
 ) where {T}
     func = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{T}[], constant)
-    for col in 1:(A.n)
+    for col in 1:A.n
         idx = _first_in_column(A, row, col)
-        idx <= last(SparseArrays.nzrange(A, col)) || continue
+        if idx > last(SparseArrays.nzrange(A, col))
+            continue
+        end
         r = _shift(A.rowval[idx], A.indexing, OneBasedIndexing())
         if r == row
             push!(
@@ -202,19 +206,26 @@ function extract_function(
     end
     return func
 end
+
 function extract_function(
     A::MutableSparseMatrixCSC{T},
     rows::UnitRange,
     constants::Vector{T},
 ) where {T}
     func = MOI.VectorAffineFunction(MOI.VectorAffineTerm{T}[], constants)
-    isempty(rows) && return func
-    idx = [_first_in_column(A, first(rows), col) for col in 1:(A.n)]
+    if isempty(rows)
+        return func
+    end
+    idx = [_first_in_column(A, first(rows), col) for col in 1:A.n]
     for output_index in eachindex(rows)
-        for col in 1:(A.n)
-            idx[col] <= last(SparseArrays.nzrange(A, col)) || continue
+        for col in 1:A.n
+            if idx[col] > last(SparseArrays.nzrange(A, col))
+                continue
+            end
             row = _shift(A.rowval[idx[col]], A.indexing, OneBasedIndexing())
-            row == rows[output_index] || continue
+            if row != rows[output_index]
+                continue
+            end
             push!(
                 func.terms,
                 MOI.VectorAffineTerm(

test/Utilities/sparse_matrix.jl

@@ -81,6 +81,40 @@ function test_VectorAffine_ZeroBased()
     @test A.colptr == [0, 1, 3, 5]
 end
 
+function test_extract_function()
+    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()
+    MOI.Utilities.add_columns(A, 3)
+    for i in 1:3
+        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)
+    row_1 = MOI.Utilities.extract_function(A, 1, 0.5)
+    @test row_1 ≈ MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5)
+    @test MOI.Utilities.extract_function(A, 1:2, [0.5, 1.2]) ≈ f
+    empty_f = MOI.Utilities.extract_function(A, 1:0, Float64[])
+    @test empty_f ≈
+          MOI.VectorAffineFunction(MOI.VectorAffineTerm{Float64}[], Float64[])
+    return
+end
+
 function test_VectorAffine_OneBased()
     A = MOI.Utilities.MutableSparseMatrixCSC{
         Float64,