diff --git a/docs/src/submodules/Utilities/reference.md b/docs/src/submodules/Utilities/reference.md index dc5fc69737..5c868830c8 100644 --- a/docs/src/submodules/Utilities/reference.md +++ b/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 diff --git a/src/Utilities/matrix_of_constraints.jl b/src/Utilities/matrix_of_constraints.jl index 13e9883c07..0468953f4c 100644 --- a/src/Utilities/matrix_of_constraints.jl +++ b/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 ### diff --git a/src/Utilities/sparse_matrix.jl b/src/Utilities/sparse_matrix.jl index 79f8ce7377..88f5f600ac 100644 --- a/src/Utilities/sparse_matrix.jl +++ b/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( diff --git a/test/Utilities/sparse_matrix.jl b/test/Utilities/sparse_matrix.jl index c0b2cad976..e03b398a9e 100644 --- a/test/Utilities/sparse_matrix.jl +++ b/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,