diff --git a/src/LowRankOpt.jl b/src/LowRankOpt.jl index e9da0a7..cf56cd8 100644 --- a/src/LowRankOpt.jl +++ b/src/LowRankOpt.jl @@ -17,7 +17,9 @@ include("distance_to_set.jl") include("Test/Test.jl") include("Bridges/Bridges.jl") +include("solution.jl") include("model.jl") +include("buffer.jl") include("schur.jl") include("MOI_wrapper.jl") include("BurerMonteiro.jl") diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl index 049b334..3ac4f57 100644 --- a/src/MOI_wrapper.jl +++ b/src/MOI_wrapper.jl @@ -268,7 +268,9 @@ function _add_constraints( A[lmi_id, k] = _MatrixBuilder{T}(d) end for row in eachindex(func.constants) - _add!(A[lmi_id, 1], row, func.constants[row], set) + if !iszero(func.constants[row]) + _add!(A[lmi_id, 1], row, func.constants[row], set) + end end for term in func.terms scalar = term.scalar_term @@ -335,7 +337,24 @@ function MOI.copy_to( key in SOLVER_OPTIONS && key != "solver" ) dest.solver = dest.options["solver"](dest.model; options...) - return MOI.Utilities.identity_index_map(src) + index_map = MOI.Utilities.identity_index_map(src) + vis_src = MOI.get(src, MOI.ListOfVariableIndices()) + # Just to throw a nice error saying we don't support such attributes + # Filter them out since we already took care of them + no_obj = MOI.Utilities.ModelFilter( + attr -> !( + attr isa MOI.ObjectiveFunction || attr isa MOI.ObjectiveSense + ), + src, + ) + # We error for the rest. In the future, we should also take care of starting values + MOI.Utilities.pass_attributes(dest, no_obj, index_map) + MOI.Utilities.pass_attributes(dest, src, index_map, vis_src) + for (F, S) in constraint_types + cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F,S}()) + MOI.Utilities.pass_attributes(dest, src, index_map, cis_src) + end + return index_map end function MOI.copy_to(dest::Optimizer{T}, src::MOI.ModelLike) where {T} @@ -462,9 +481,6 @@ end struct Solution <: MOI.AbstractModelAttribute end MOI.is_set_by_optimize(::Solution) = true -function MOI.get(solver::SolverCore.AbstractOptimizationSolver, ::Solution) - return VectorizedSolution(solver.stats.solution, solver.model.dim) -end MOI.get(optimizer::Optimizer, attr::Solution) = MOI.get(optimizer.solver, attr) function MOI.get( diff --git a/src/buffer.jl b/src/buffer.jl new file mode 100644 index 0000000..90d5678 --- /dev/null +++ b/src/buffer.jl @@ -0,0 +1,277 @@ +mutable struct BufferedModelForSchur{ + T, + C<:AbstractMatrix{T}, + A<:AbstractMatrix{T}, + JB, + JTB, + SB, +} <: AbstractModel{T} + model::Model{T,C,A} + meta::NLPModels.NLPModelMeta{T,Vector{T}} + jprod_buffer::JB + jtprod_buffer::JTB + schur_buffer::SB +end + +function BufferedModelForSchur(model, datasparsity) + return BufferedModelForSchur( + model, + model.meta, + buffer_for_jprod(model), + buffer_for_jtprod(model), + buffer_for_schur_complement(model, datasparsity), + ) +end + +num_scalars(model::BufferedModelForSchur) = num_scalars(model.model) +num_matrices(model::BufferedModelForSchur) = num_matrices(model.model) +matrix_indices(model::BufferedModelForSchur) = matrix_indices(model.model) +side_dimension(model::BufferedModelForSchur, i) = side_dimension(model.model, i) + +####################### +###### Objective ###### +####################### + +function NLPModels.obj(model::BufferedModelForSchur, x::AbstractVector) + return NLPModels.obj(model.model, x) +end + +function dual_obj(model::BufferedModelForSchur, y::AbstractVector) + return dual_obj(model.model, y) +end + +function NLPModels.grad(model::BufferedModelForSchur, ::Type{ScalarIndex}) + return NLPModels.grad(model.model, ScalarIndex) +end + +function NLPModels.grad(model::BufferedModelForSchur, i::MatrixIndex) + return NLPModels.grad(model.model, i) +end + +######################### +###### Constraints ###### +######################### + +cons_constant(model::BufferedModelForSchur) = cons_constant(model.model) +function NLPModels.jac( + model::BufferedModelForSchur, + j::Integer, + ::Type{ScalarIndex}, +) + return NLPModels.jac(model.model, j, ScalarIndex) +end +function norm_jac(model::BufferedModelForSchur, i::MatrixIndex) + return norm_jac(model.model, i) +end + +####################### +###### J product ###### +####################### + +function jprod!(model, x, v, Jv, ::Type{ScalarIndex}) + return jprod!(model.model, x, v, Jv, ScalarIndex) +end + +function _add_vec!(_, _, _, _, offset, ::FillArrays.Zeros) + return offset +end + +function _add_vec!(I, J, V, j, offset, A::SparseArrays.SparseMatrixCSC) + Ai, Av = SparseArrays.findnz(A[:]) + K = offset .+ eachindex(Ai) + I[K] = Ai + J[K] .= j + V[K] = Av + return offset + length(Ai) +end + +# `SparseMatrixCSC` is stored with an offset by column. +# This means that getting view `view(A, :, I)` can be handles efficently, +# these give `SparseMatrixCSCView` (if `I` is a `UnitRange`) and +# `SparseMatrixCSCColumnSubset` otherwise. +# In `schur.jl`, we therefore get a `SparseMatrixCSCColumnSubset`. +# Since we want to use subsets of constraint indices, we use the columns +# of `A` for constraint indices and the rows of `A` for matrix indices. +function buffer_for_jprod(model::Model{T}, i::MatrixIndex) where {T} + nnz = sum(1:model.meta.ncon; init = 0) do j + return _nnz(model.A[i.value, j]) + end + I = zeros(Int64, nnz) + J = zeros(Int64, nnz) + V = zeros(T, nnz) + offset = 0 + for j in 1:model.meta.ncon + offset = _add_vec!(I, J, V, j, offset, model.A[i.value, j]) + end + A = SparseArrays.sparse( + I, + J, + V, + side_dimension(model, i)^2, + model.meta.ncon, + ) + return A +end + +function buffer_for_jprod(model::Model{T}) where {T} + return SparseArrays.SparseMatrixCSC{T,Int64}[ + buffer_for_jprod(model, i) for i in matrix_indices(model) + ] +end + +_vec(x::AbstractVector) = x +_vec(x::FillArrays.Zeros{T}) where {T} = FillArrays.Zeros{T}(length(x)) +_vec(x::AbstractArray) = UnsafeArrays.uview(x, :) +_vec(x::Base.ReshapedArray) = _vec(parent(x)) + +function _add_jprod!(V, Jv::AbstractArray{T}, A) where {T} + return LinearAlgebra.mul!(Jv, A', _vec(V), true, true) +end + +function add_sub_jprod!( + model::BufferedModelForSchur, + i::MatrixIndex, + V::AbstractMatrix, + Jv::AbstractVector, + I, +) + # `view(cache, I)` would be terribly slow, only the number of elements of `I` matter here + A = model.jprod_buffer[i.value] + return _add_jprod!(V, Jv, view(A, :, I)) +end + +function add_jprod!( + model::BufferedModelForSchur, + V::AbstractMatrix, + Jv::AbstractVector, + i::MatrixIndex, +) + return _add_jprod!(V, Jv, model.jprod_buffer[i.value]) +end + +######################## +###### Jᵀ product ###### +######################## + +function jtprod!( + model::BufferedModelForSchur, + y::AbstractVector, + vJ::AbstractVector, + ::Type{ScalarIndex}, +) + return jtprod!(model.model, y, vJ, ScalarIndex) +end + +function buffer_for_jtprod(model::Model) + if iszero(num_matrices(model)) + return + end + return map(Base.Fix1(buffer_for_jtprod, model), matrix_indices(model)) +end + +function _merge_sparsity( + A::SparseArrays.SparseMatrixCSC, + B::SparseArrays.SparseMatrixCSC, +) + return A + B +end +_merge_sparsity(::FillArrays.Zeros, B::SparseArrays.SparseMatrixCSC) = B +_merge_sparsity(A::SparseArrays.SparseMatrixCSC, ::FillArrays.Zeros) = A +_merge_sparsity(A::FillArrays.Zeros, ::FillArrays.Zeros) = A + +_abs(A::SparseArrays.SparseMatrixCSC) = abs.(A) +_abs(A::FillArrays.Zeros) = A + +function buffer_for_jtprod(model::Model{T}, mat_idx::MatrixIndex) where {T} + if iszero(model.meta.ncon) + d = side_dimension(model, mat_idx) + return FillArrays.Zeros{T}(d, d) + end + # FIXME: at some point, switch to dense + # /!\ If there is only one nonzero matrix and we didn't have `_abs`, + # we would return an alias of that only matrix so that `_abs` has the + # non-obvious role of avoid this as well as avoiding cancellations. + return reduce( + _merge_sparsity, + _abs(model.A[mat_idx.value, j]) for j in 1:model.meta.ncon + ) +end + +function NLPModels.jtprod!( + model::BufferedModelForSchur, + _::AbstractVector, + y::AbstractVector, + vJ::AbstractVector, +) + jtprod!(model, y, vJ[ScalarIndex], ScalarIndex) + for mat_idx in matrix_indices(model) + vJ[mat_idx] .= jtprod!(model, y, mat_idx) + end +end + +_zero!(A::FillArrays.Zeros) = A +_zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) + +# Computes `A .+= B * α` +function _add_mul!(::FillArrays.Zeros, ::FillArrays.Zeros, _) end + +function _add_mul!(A::SparseArrays.SparseMatrixCSC, ::FillArrays.Zeros, _) + return A +end + +function _add_mul!( + A::SparseArrays.SparseMatrixCSC, + B::SparseArrays.SparseMatrixCSC, + α, +) + for col in axes(A, 2) + range_A = SparseArrays.nzrange(A, col) + it_A = iterate(range_A) + for k in SparseArrays.nzrange(B, col) + row_B = SparseArrays.rowvals(B)[k] + while SparseArrays.rowvals(A)[it_A[1]] < row_B + it_A = iterate(range_A, it_A[2]) + end + @assert row_B == SparseArrays.rowvals(A)[it_A[1]] + SparseArrays.nonzeros(A)[it_A[1]] += SparseArrays.nonzeros(B)[k] * α + end + end +end + +function jtprod!(model::Model, y, buffer, mat_idx::MatrixIndex) + _zero!(buffer) + for j in eachindex(y) + _add_mul!(buffer, model.A[mat_idx.value, j], y[j]) + end + return buffer +end + +function jtprod!(model::BufferedModelForSchur, y, mat_idx::MatrixIndex) + return jtprod!(model.model, y, model.jtprod_buffer[mat_idx.value], mat_idx) +end + +function dual_cons!( + model::BufferedModelForSchur, + y::AbstractVector, + res, + ::Type{ScalarIndex}, +) + return dual_cons!(model.model, y, res, ScalarIndex) +end + +# TODO If we rename `dual_cons!` to `unsafe_dual_cons`, +# we can remove the `copy` and remplace `-B` with a mutation + +# Note that we can't use `-` because of https://github.com/JuliaArrays/FillArrays.jl/issues/412 +_sub(A::FillArrays.Zeros, ::FillArrays.Zeros) = A +_sub(A::AbstractArray, ::FillArrays.Zeros) = copy(A) +_sub(::FillArrays.Zeros, B::AbstractArray) = -B +_sub(A::AbstractArray, B::AbstractArray) = A - B + +function dual_cons!( + model::BufferedModelForSchur, + y::AbstractVector, + i::MatrixIndex, +) + return _sub(model.model.C[i.value], jtprod!(model, y, i)) +end diff --git a/src/errors.jl b/src/errors.jl index 8f7349a..15f8ee3 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -4,7 +4,7 @@ Return [the 6 standard DIMACS errors](https://plato.asu.edu/dimacs/node3.html). """ function errors( - model::Model, + model::AbstractModel, x; y = nothing, primal_err = NLPModels.cons(model, x), diff --git a/src/model.jl b/src/model.jl index 6a0eb63..e2a803c 100644 --- a/src/model.jl +++ b/src/model.jl @@ -5,97 +5,32 @@ import MathOptInterface as MOI import NLPModels import UnsafeArrays -struct Dimensions - num_scalars::Int64 - side_dimensions::Vector{Int64} - offsets::Vector{Int64} -end - -num_matrices(d::Dimensions) = length(d.side_dimensions) -Base.length(d::Dimensions) = d.offsets[end] - -struct ScalarIndex - value::Int64 -end +abstract type AbstractModel{T} <: NLPModels.AbstractNLPModel{T,Vector{T}} end +Base.broadcastable(model::AbstractModel) = Ref(model) -struct MatrixIndex - value::Int64 -end -Base.broadcastable(i::MatrixIndex) = Ref(i) - -abstract type AbstractSolution{T} <: AbstractVector{T} end - -function Base.getindex( - s::AbstractSolution, - i::Union{Type{ScalarIndex},MatrixIndex}, +function NLPModels.cons!( + model::AbstractModel, + x::AbstractVector, + cx::AbstractVector, ) - return view(s, i) -end - -struct VectorizedSolution{T} <: AbstractSolution{T} - x::Vector{T} - dim::Dimensions -end - -function LinearAlgebra.dot(x::VectorizedSolution, z::VectorizedSolution) - return LinearAlgebra.dot(x.x, z.x) -end - -num_matrices(x::VectorizedSolution) = num_matrices(x.dim) - -Base.similar(s::VectorizedSolution) = VectorizedSolution(similar(s.x), s.dim) - -Base.size(s::VectorizedSolution) = (length(s.dim),) - -function Base.to_index(s::VectorizedSolution, ::Type{ScalarIndex}) - return Base.OneTo(s.dim.num_scalars) -end - -function Base.to_index(s::VectorizedSolution, mi::MatrixIndex) - i = mi.value - return (1+s.dim.offsets[i]):s.dim.offsets[i+1] -end - -function Base.view(s::VectorizedSolution, ::Type{ScalarIndex}) - return view(s.x, Base.to_index(s, ScalarIndex)) -end - -# `s[i] .= ...` calleds `copyto!(view(s, i), Broadcasted(...))` -function Base.view(s::VectorizedSolution, i::MatrixIndex) - v = view(s.x, Base.to_index(s, i)) - dim = s.dim.side_dimensions[i.value] - X = reshape(v, dim, dim) - return X -end - -Base.setindex!(s::VectorizedSolution, v, i::Integer) = setindex!(s.x, v, i) -Base.getindex(s::VectorizedSolution, i::Integer) = getindex(s.x, i) - -struct ShapedSolution{T,MT<:AbstractMatrix{T}} <: AbstractSolution{T} - scalars::Vector{T} - matrices::Vector{MT} -end - -function Base.size(s::ShapedSolution) - return (length(s.scalars) + sum(length, s.matrices, init = 0),) -end -num_matrices(s::ShapedSolution) = length(s.matrices) - -function LinearAlgebra.norm2(s::ShapedSolution{T}) where {T} - # `LinearAlgebra.generic_norm2` starts by computing the ∞ norm and do a rescaling, we don't do that here - return √(LinearAlgebra.dot(s, s)) + NLPModels.jprod!(model, x, x, cx) + cx .-= cons_constant(model) + return cx end -function LinearAlgebra.dot(a::ShapedSolution{T}, b::ShapedSolution{T}) where {T} - return LinearAlgebra.dot(a.scalars, b.scalars) + - sum(eachindex(a.matrices); init = zero(T)) do i - return LinearAlgebra.dot(a.matrices[i], b.matrices[i]) +function NLPModels.jprod!( + model::AbstractModel, + x::AbstractVector, + v::AbstractVector, + Jv::AbstractVector, +) + jprod!(model, x, v[ScalarIndex], Jv, ScalarIndex) + for i in matrix_indices(model) + add_jprod!(model, v[i], Jv, i) end + return Jv end -Base.view(s::ShapedSolution, ::Type{ScalarIndex}) = s.scalars -Base.view(s::ShapedSolution, i::MatrixIndex) = s.matrices[i.value] - """ Model @@ -127,7 +62,7 @@ The fields of the `struct` as related to the arrays of the above formulation as * The matrix ``A_{i,j}`` is given by `A[i,j]`. """ mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: - NLPModels.AbstractNLPModel{T,Vector{T}} + AbstractModel{T} meta::NLPModels.NLPModelMeta{T,Vector{T}} dim::Dimensions C::Vector{C} @@ -184,20 +119,9 @@ end side_dimension(model::Model, i::MatrixIndex) = model.msizes[i.value] -# Should be only used with `norm` -NLPModels.jac(model::Model, ::Type{ScalarIndex}) = model.C_lin -function NLPModels.jac(model::Model, j::Integer, i::MatrixIndex) - return model.A[i.value, j] -end -function NLPModels.jac(model::Model, j::Integer, ::Type{ScalarIndex}) - return model.C_lin[j, :] -end -function norm_jac(model::Model{T}, i::MatrixIndex) where {T} - if isempty(model.A) - return zero(T) - end - return LinearAlgebra.norm(model.A[i.value, :]) -end +####################### +###### Objective ###### +####################### function NLPModels.obj(model::Model, X::AbstractMatrix, i::MatrixIndex) return LinearAlgebra.dot(model.C[i.value], X) @@ -230,196 +154,65 @@ end dual_obj(model::Model, y::AbstractVector) = LinearAlgebra.dot(model.b, y) -function jtprod(model::Model, ::Type{ScalarIndex}, y::AbstractVector) - return model.C_lin' * y -end - -function buffer_for_jtprod(model::Model) - if iszero(num_matrices(model)) - return - end - return map(Base.Fix1(buffer_for_jtprod, model), matrix_indices(model)) -end - -function buffer_for_jtprod(model::Model, mat_idx::MatrixIndex) - if iszero(model.meta.ncon) - return - end - # FIXME: at some point, switch to dense - return sum(abs.(model.A[mat_idx.value, j]) for j in 1:model.meta.ncon) -end - -# Computes `A .+= B * α` -function _add_mul!( - A::SparseArrays.SparseMatrixCSC, - B::SparseArrays.SparseMatrixCSC, - α, -) - for col in axes(A, 2) - range_A = SparseArrays.nzrange(A, col) - it_A = iterate(range_A) - for k in SparseArrays.nzrange(B, col) - row_B = SparseArrays.rowvals(B)[k] - while SparseArrays.rowvals(A)[it_A[1]] < row_B - it_A = iterate(range_A, it_A[2]) - end - @assert row_B == SparseArrays.rowvals(A)[it_A[1]] - SparseArrays.nonzeros(A)[it_A[1]] += SparseArrays.nonzeros(B)[k] * α - end - end -end - -_zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) - -function jtprod!(buffer, model::Model, mat_idx::MatrixIndex, y) - _zero!(buffer) - for j in eachindex(y) - _add_mul!(buffer, model.A[mat_idx.value, j], y[j]) - end - return buffer -end - -function dual_cons(model::Model, ::Type{ScalarIndex}, y::AbstractVector) - return model.d_lin - jtprod(model, ScalarIndex, y) -end - -function dual_cons!( - buffer, +function jtprod!( model::Model, - mat_idx::MatrixIndex, y::AbstractVector, + vJ::AbstractVector, + ::Type{ScalarIndex}, ) - i = mat_idx.value - return model.C[i] - jtprod!(buffer[i], model, mat_idx, y) + return LinearAlgebra.mul!(vJ, model.C_lin', y) end -function NLPModels.jtprod!( - model::Model, - _::AbstractVector, - y::AbstractVector, - vJ::AbstractVector, - buffer, -) - vJ[ScalarIndex] .= jtprod(model, ScalarIndex, y) - for mat_idx in matrix_indices(model) - i = mat_idx.value - vJ[mat_idx] .= jtprod!(buffer[i], model, mat_idx, y) - end +function dual_cons!(model::Model, y::AbstractVector, res, ::Type{ScalarIndex}) + copyto!(res, model.d_lin) + return LinearAlgebra.mul!(res, model.C_lin', y, -1, true) end NLPModels.grad(model::Model, ::Type{ScalarIndex}) = model.d_lin NLPModels.grad(model::Model, i::MatrixIndex) = model.C[i.value] -cons_constant(model::Model) = model.b - -function NLPModels.cons!( - model::Model, - x::AbstractVector, - cx::AbstractVector, - args::Vararg{Any,N}, -) where {N} - NLPModels.jprod!(model, x, x, cx, args...) - cx .-= model.b - return cx -end - -# `SparseMatrixCSC` is stored with an offset by column. -# This means that getting view `view(A, :, I)` can be handles efficently, -# these give `SparseMatrixCSCView` (if `I` is a `UnitRange`) and -# `SparseMatrixCSCColumnSubset` otherwise. -# In `schur.jl`, we therefore get a `SparseMatrixCSCColumnSubset`. -# Since we want to use subsets of constraint indices, we use the columns -# of `A` for constraint indices and the rows of `A` for matrix indices. -function buffer_for_jprod(model::Model{T}, i::MatrixIndex) where {T} - nnz = sum(1:model.meta.ncon; init = 0) do j - return SparseArrays.nnz(model.A[i.value, j]) - end - I = zeros(Int64, nnz) - J = zeros(Int64, nnz) - V = zeros(T, nnz) - offset = 0 - for j in 1:model.meta.ncon - Ai, Av = SparseArrays.findnz(model.A[i.value, j][:]) - K = offset .+ eachindex(Ai) - I[K] = Ai - J[K] .= j - V[K] = Av - offset += length(Ai) - end - A = SparseArrays.sparse( - I, - J, - V, - side_dimension(model, i)^2, - model.meta.ncon, - ) - return A -end +######################### +###### Constraints ###### +######################### -# We define a new type so that we can define a custom `getindex` -struct JProdBuffer{T} - A::Vector{SparseArrays.SparseMatrixCSC{T,Int64}} -end +cons_constant(model::Model) = model.b -function buffer_for_jprod(model::Model{T}) where {T} - return JProdBuffer([ - buffer_for_jprod(model, i) for i in matrix_indices(model) - ],) +# Should be only used with `norm` +NLPModels.jac(model::Model, ::Type{ScalarIndex}) = model.C_lin +function NLPModels.jac(model::Model, j::Integer, i::MatrixIndex) + return model.A[i.value, j] end - -Base.getindex(buf::JProdBuffer, i::MatrixIndex) = buf.A[i.value] - -_vec(x::AbstractVector) = x -_vec(x::AbstractArray) = UnsafeArrays.uview(x, :) -_vec(x::Base.ReshapedArray) = _vec(parent(x)) - -function _add_jprod!(V, Jv::AbstractArray{T}, A) where {T} - return LinearAlgebra.mul!(Jv, A', _vec(V), true, true) +function NLPModels.jac(model::Model, j::Integer, ::Type{ScalarIndex}) + return model.C_lin[j, :] end - -function add_sub_jprod!( - _::Model, - _::MatrixIndex, - V::AbstractMatrix, - Jv::AbstractVector, - I, - A, -) - # `view(cache, I)` would be terribly slow, only the number of elements of `I` matter here - return _add_jprod!(V, Jv, view(A, :, I)) +function norm_jac(model::Model{T}, i::MatrixIndex) where {T} + if isempty(model.A) + return zero(T) + end + return LinearAlgebra.norm(model.A[i.value, :]) end -function add_jprod!( - ::Model, - ::MatrixIndex, - V::AbstractMatrix, - Jv::AbstractVector, - buffer, -) - return _add_jprod!(V, Jv, buffer) -end +####################### +###### J product ###### +####################### function add_jprod!( model::Model, - i::MatrixIndex, V::AbstractMatrix, Jv::AbstractVector, + i::MatrixIndex, ) for j in 1:model.meta.ncon Jv[j] += LinearAlgebra.dot(model.A[i.value, j], V) end end -function NLPModels.jprod!( +function jprod!( model::Model, _::AbstractVector, v::AbstractVector, Jv::AbstractVector, - args::Vararg{Any,N}, # Optional buffer -) where {N} - LinearAlgebra.mul!(Jv, model.C_lin, v[ScalarIndex]) - for i in matrix_indices(model) - add_jprod!(model, i, v[i], Jv, getindex.(args, i)...) - end - return Jv + ::Type{ScalarIndex}, +) + return LinearAlgebra.mul!(Jv, model.C_lin, v) end diff --git a/src/schur.jl b/src/schur.jl index 8760ca6..68e4990 100644 --- a/src/schur.jl +++ b/src/schur.jl @@ -42,6 +42,9 @@ function _dot( return result end +_nnz(::FillArrays.Zeros) = 0 +_nnz(A::SparseArrays.SparseMatrixCSC) = SparseArrays.nnz(A) + # The `jprod!` buffer is guaranteed to be the first argument of the tuple. # This assumption is used by Loraine. function buffer_for_schur_complement(model::Model{T}, κ) where {T} @@ -51,7 +54,7 @@ function buffer_for_schur_complement(model::Model{T}, κ) where {T} for mat_idx in matrix_indices(model) i = mat_idx.value - nzA = [SparseArrays.nnz(model.A[i, j]) for j in 1:n] + nzA = [_nnz(model.A[i, j]) for j in 1:n] σ[:, i] = sortperm(nzA, rev = true) sorted = nzA[σ[:, i]] @@ -63,48 +66,51 @@ function buffer_for_schur_complement(model::Model{T}, κ) where {T} for dim in model.msizes] WAW = copy.(AW) - return buffer_for_jprod(model), AW, WAW, σ, last_dense + return AW, WAW, σ, last_dense end -function add_schur_complement!(model::Model, W, ::Type{MatrixIndex}, H, buffer) +function add_schur_complement!( + model::BufferedModelForSchur, + W, + ::Type{MatrixIndex}, + H, +) for i in matrix_indices(model) - add_schur_complement!(model, i, W[i], H, buffer) + add_schur_complement!(model, i, W[i], H) end return H end # /!\ W needs to be symmetric function add_schur_complement!( - model::Model, + model::BufferedModelForSchur, mat_idx::MatrixIndex, W::AbstractMatrix{T}, H, - buffer, ) where {T} - jprod_buffer, AW, WAW, σ, last_dense = buffer - buf = jprod_buffer[mat_idx] + AW, WAW, σ, last_dense = model.schur_buffer ilmi = mat_idx.value n = model.meta.ncon for ii in axes(H, 1) i = σ[ii, ilmi] - Ai = model.A[ilmi, i] - if SparseArrays.nnz(Ai) > 0 + Ai = model.model.A[ilmi, i] + if _nnz(Ai) > 0 if ii <= last_dense[ilmi] LinearAlgebra.mul!(AW[ilmi], W, Ai) LinearAlgebra.mul!(WAW[ilmi], AW[ilmi], W) I = view(σ, ii:n, ilmi) - add_sub_jprod!(model, mat_idx, WAW[ilmi], view(H, I, i), I, buf) + add_sub_jprod!(model, mat_idx, WAW[ilmi], view(H, I, i), I) for jj in (ii+1):n j = σ[jj, ilmi] H[i, j] = H[j, i] end else - if SparseArrays.nnz(Ai) > 1 + if _nnz(Ai) > 1 @inbounds for jj in ii:n j = σ[jj, ilmi] - Aj = model.A[ilmi, j] - if !iszero(SparseArrays.nnz(Aj)) + Aj = model.model.A[ilmi, j] + if !iszero(_nnz(Aj)) ttt = _dot(Ai, Aj, W) H[i, j] += ttt if i != j @@ -112,17 +118,17 @@ function add_schur_complement!( end end end - elseif SparseArrays.nnz(Ai) == 1 + elseif _nnz(Ai) == 1 # A is symmetric iiiiAi = jjjiAi = only(SparseArrays.rowvals(Ai)) vvvi = only(SparseArrays.nonzeros(Ai)) @inbounds for jj in ii:n j = σ[jj, ilmi] - Ajjj = model.A[ilmi, j] + Ajjj = model.model.A[ilmi, j] # As we sort the matrices in decreasing `nnz` order, # the rest of matrices is either zero or have only # one entry - if !iszero(SparseArrays.nnz(Ajjj)) + if !iszero(_nnz(Ajjj)) iiijAj = jjjjAj = only(SparseArrays.rowvals(Ajjj)) vvvj = only(SparseArrays.nonzeros(Ajjj)) ttt = @@ -143,18 +149,23 @@ function add_schur_complement!( return H end -function add_schur_complement!(model::Model, w, ::Type{ScalarIndex}, H) - H .+= model.C_lin * SparseArrays.spdiagm(w) * model.C_lin' +function add_schur_complement!( + model::BufferedModelForSchur, + w, + ::Type{ScalarIndex}, + H, +) + H .+= model.model.C_lin * SparseArrays.spdiagm(w) * model.model.C_lin' return H end # [HKS24, (5b)] # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA_jW⟩ -function schur_complement!(model::Model, W::AbstractVector, H, buffer) +function schur_complement!(model::BufferedModelForSchur, W::AbstractVector, H) fill!(H, zero(eltype(H))) if num_matrices(model) > 0 - add_schur_complement!(model, W, MatrixIndex, H, buffer) + add_schur_complement!(model, W, MatrixIndex, H) end if num_scalars(model) > 0 add_schur_complement!(model, W[ScalarIndex], ScalarIndex, H) @@ -165,24 +176,11 @@ end # [HKS24, (5b)] # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA(y)W⟩ -function eval_schur_complement!( - result, - model::Model, - W, - y, - jprod_buffer, - jtprod_buffer, -) +function eval_schur_complement!(model::BufferedModelForSchur, W, y, result) fill!(result, zero(eltype(result))) for i in matrix_indices(model) - add_jprod!( - model, - i, - W[i] * jtprod!(jtprod_buffer[i.value], model, i, y) * W[i], - result, - jprod_buffer[i], - ) + add_jprod!(model, W[i] * jtprod!(model, y, i) * W[i], result, i) end - result .+= model.C_lin * (W[ScalarIndex] .* (model.C_lin' * y)) + result .+= model.model.C_lin * (W[ScalarIndex] .* (model.model.C_lin' * y)) return result end diff --git a/src/solution.jl b/src/solution.jl new file mode 100644 index 0000000..3009403 --- /dev/null +++ b/src/solution.jl @@ -0,0 +1,90 @@ +struct Dimensions + num_scalars::Int64 + side_dimensions::Vector{Int64} + offsets::Vector{Int64} +end + +num_matrices(d::Dimensions) = length(d.side_dimensions) +Base.length(d::Dimensions) = d.offsets[end] + +struct ScalarIndex + value::Int64 +end + +struct MatrixIndex + value::Int64 +end +Base.broadcastable(i::MatrixIndex) = Ref(i) + +abstract type AbstractSolution{T} <: AbstractVector{T} end + +function Base.getindex( + s::AbstractSolution, + i::Union{Type{ScalarIndex},MatrixIndex}, +) + return view(s, i) +end + +struct VectorizedSolution{T} <: AbstractSolution{T} + x::Vector{T} + dim::Dimensions +end + +function LinearAlgebra.dot(x::VectorizedSolution, z::VectorizedSolution) + return LinearAlgebra.dot(x.x, z.x) +end + +num_matrices(x::VectorizedSolution) = num_matrices(x.dim) + +Base.similar(s::VectorizedSolution) = VectorizedSolution(similar(s.x), s.dim) + +Base.size(s::VectorizedSolution) = (length(s.dim),) + +function Base.to_index(s::VectorizedSolution, ::Type{ScalarIndex}) + return Base.OneTo(s.dim.num_scalars) +end + +function Base.to_index(s::VectorizedSolution, mi::MatrixIndex) + i = mi.value + return (1+s.dim.offsets[i]):s.dim.offsets[i+1] +end + +function Base.view(s::VectorizedSolution, ::Type{ScalarIndex}) + return view(s.x, Base.to_index(s, ScalarIndex)) +end + +# `s[i] .= ...` calleds `copyto!(view(s, i), Broadcasted(...))` +function Base.view(s::VectorizedSolution, i::MatrixIndex) + v = view(s.x, Base.to_index(s, i)) + dim = s.dim.side_dimensions[i.value] + X = reshape(v, dim, dim) + return X +end + +Base.setindex!(s::VectorizedSolution, v, i::Integer) = setindex!(s.x, v, i) +Base.getindex(s::VectorizedSolution, i::Integer) = getindex(s.x, i) + +struct ShapedSolution{T,MT<:AbstractMatrix{T}} <: AbstractSolution{T} + scalars::Vector{T} + matrices::Vector{MT} +end + +function Base.size(s::ShapedSolution) + return (length(s.scalars) + sum(length, s.matrices, init = 0),) +end +num_matrices(s::ShapedSolution) = length(s.matrices) + +function LinearAlgebra.norm2(s::ShapedSolution{T}) where {T} + # `LinearAlgebra.generic_norm2` starts by computing the ∞ norm and do a rescaling, we don't do that here + return √(LinearAlgebra.dot(s, s)) +end + +function LinearAlgebra.dot(a::ShapedSolution{T}, b::ShapedSolution{T}) where {T} + return LinearAlgebra.dot(a.scalars, b.scalars) + + sum(eachindex(a.matrices); init = zero(T)) do i + return LinearAlgebra.dot(a.matrices[i], b.matrices[i]) + end +end + +Base.view(s::ShapedSolution, ::Type{ScalarIndex}) = s.scalars +Base.view(s::ShapedSolution, i::MatrixIndex) = s.matrices[i.value] diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl index 458344c..e38f242 100644 --- a/test/BurerMonteiro.jl +++ b/test/BurerMonteiro.jl @@ -182,7 +182,8 @@ end; ) @test model.meta.ncon == 0 @test LRO.norm_jac(model, LRO.MatrixIndex(1)) == 0 - @test isnothing(LRO.buffer_for_jtprod(model, LRO.MatrixIndex(1))) + @test LRO.buffer_for_jtprod(model, LRO.MatrixIndex(1)) isa + LRO.FillArrays.Zeros end; @testset "Fallback" begin diff --git a/test/buffer.jl b/test/buffer.jl new file mode 100644 index 0000000..068b61e --- /dev/null +++ b/test/buffer.jl @@ -0,0 +1,78 @@ +module TestBuffer + +import FillArrays, SparseArrays +using JuMP, Dualization +include("diff_check.jl") + +# Test with zero Ai matrices +function _test_zero_Ai(all_zero::Bool, matrix_in_objective::Bool) + model = Model(dual_optimizer(LRO.Optimizer)) + @variable(model, x[1:2] in MOI.Nonnegatives(2)) + @variable(model, X[1:2, 1:2] in PSDCone()) + @constraint(model, sum(x) == 1) + @constraint(model, 2sum(x) == 2) + if !all_zero + @constraint(model, sum(X) == 2) + end + @constraint(model, x[1] - x[2] == 1) + if matrix_in_objective + @objective(model, Max, x[1] + X[1, 2] - X[1, 1]) + else + @objective(model, Max, x[1]) + end + set_attribute(model, "solver", ConvexSolver) + optimize!(model) + b = _backend(model) + T = Float64 + Z = FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}} + S = SparseArrays.SparseMatrixCSC{T,Int} + @test b.model.C isa Vector{matrix_in_objective ? S : Z} + if all_zero + @test b.model.A isa Matrix{Z} + else + @test b.model.A isa Matrix{Union{Z,S}} + end + @test b.model.A[1] isa Z + @test b.model.A[2] isa Z + if all_zero + @test b.model.A[3] isa Z + else + @test b.model.A[3] isa S + @test b.model.A[4] isa Z + end + buf = LRO.BufferedModelForSchur(b.model, 1) + for A in b.model.A + if all_zero + @test buf.jtprod_buffer[] isa LRO.FillArrays.Zeros + else + @test buf.jtprod_buffer[] !== A + end + end + for κ in 0:5 + schur_test(model, κ) + end +end + +function test_zero_Ai() + @testset "all_zero=$all_zero" for all_zero in [false, true] + @testset "matrix_in_objective=$matrix_in_objective" for matrix_in_objective in + [false, true] + _test_zero_Ai(all_zero, matrix_in_objective) + end + end + return +end + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$name", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end +end + +end + +TestBuffer.runtests() diff --git a/test/diff_check.jl b/test/diff_check.jl index 94d5567..026a462 100644 --- a/test/diff_check.jl +++ b/test/diff_check.jl @@ -4,8 +4,13 @@ # in the LICENSE.md file or at https://opensource.org/licenses/MIT. using Test +using LinearAlgebra +using FillArrays +import SolverCore +using Dualization +import LowRankOpt as LRO -function test_vecprod(f, len, J; tol = 1e-6) +function _test_vecprod(f, len, J; tol = 1e-6) v = ones(len) @test f(v) ≈ J * v rtol = tol atol = tol v = -ones(len) @@ -33,7 +38,7 @@ function jac_check(model, x; kws...) f(x) = NLPModels.cons(model, x) J = FiniteDiff.finite_difference_jacobian(f, x) @testset "jprod" begin - test_vecprod( + _test_vecprod( v -> NLPModels.jprod(model, x, v), model.meta.nvar, J; @@ -41,7 +46,7 @@ function jac_check(model, x; kws...) ) end @testset "jtprod" begin - test_vecprod( + _test_vecprod( v -> NLPModels.jtprod(model, x, v), model.meta.ncon, J'; @@ -56,7 +61,7 @@ function hess_check(model, x; kws...) f(x) = obj_weight * NLPModels.obj(model, x) - dot(y, NLPModels.cons(model, x)) J = FiniteDiff.finite_difference_hessian(f, x) - return test_vecprod( + return _test_vecprod( v -> NLPModels.hprod(model, x, y, v; obj_weight), model.meta.nvar, J; @@ -74,18 +79,117 @@ function _backend(model) return b end -function diff_check(model) - b = _backend(model) - bm = b.solver.model - x = rand(bm.meta.nvar) +function diff_check(model::NLPModels.AbstractNLPModel) + x = rand(model.meta.nvar) @testset "Gradient" begin - grad_check(bm, x) - @test isempty(NLPModelsTest.gradient_check(bm; x)) + grad_check(model, x) + @test isempty(NLPModelsTest.gradient_check(model; x)) end @testset "Jacobian" begin - jac_check(bm, x) + jac_check(model, x) end @testset "Hessian" begin - hess_check(bm, x) + hess_check(model, x) + end +end + +function diff_check(model::JuMP.AbstractModel) + b = _backend(model) + return diff_check(b.solver.model) +end + +struct ConvexSolver{T} <: SolverCore.AbstractOptimizationSolver + model::LRO.Model{T} + stats::SolverCore.GenericExecutionStats{T,Vector{T},Vector{T},Any} +end + +function ConvexSolver(model::LRO.Model) + stats = SolverCore.GenericExecutionStats(model) + return ConvexSolver(model, stats) +end + +function LRO.MOI.get(solver::ConvexSolver, ::LRO.Solution) + return LRO.VectorizedSolution(solver.stats.solution, solver.model.dim) +end + +function SolverCore.solve!(::ConvexSolver, ::LRO.Model) + return +end + +function _alloc_schur_complement(model, i, Wi, H) + if VERSION < v"1.11" + return + end + LRO.add_schur_complement!(model, i, Wi, H) + @test 0 == @allocated LRO.add_schur_complement!(model, i, Wi, H) +end + +function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} + n = model.meta.ncon + y = rand(T, n) + + idx = LRO.matrix_indices(model) + @test NLPModels.obj(model, w) ≈ + dot(NLPModels.grad(model, LRO.ScalarIndex), w[LRO.ScalarIndex]) + + dot(NLPModels.grad.(model, idx), getindex.(Ref(w), idx)) + @test LRO.dual_obj(model, y) ≈ dot(LRO.cons_constant(model), y) + + @test NLPModels.jac(model, 1, LRO.ScalarIndex) == + NLPModels.jac(model.model, 1, LRO.ScalarIndex) + @test LRO.norm_jac.(model, idx) == LRO.norm_jac.(model.model, idx) + @test LRO.side_dimension.(model, idx) == + LRO.side_dimension.(model.model, idx) + if !isempty(idx) + @test LRO.side_dimension.([model], idx[1]) == + [LRO.side_dimension(model, idx[1])] + end + + Jv = similar(y) + vJ = similar(w) + NLPModels.jprod!(model, w, w, Jv) + NLPModels.jtprod!(model, w, y, vJ) + @test dot(Jv, y) ≈ dot(vJ, w) + + H = zeros(n, n) + H = LRO.schur_complement!(model, w, H) + Hy = similar(y) + LRO.eval_schur_complement!(model, w, y, Hy) + @test Hy ≈ H * y + for i in LRO.matrix_indices(model) + Wi = @inferred w[i] + _alloc_schur_complement(model, i, Wi, H) + end + for i in LRO.matrix_indices(model) + @test model.jtprod_buffer[i.value] isa + Union{FillArrays.Zeros,SparseArrays.SparseMatrixCSC} + ret = LRO.jtprod!(model, y, i) + @test ret === model.jtprod_buffer[i.value] + ret = LRO.dual_cons!(model, y, i) + if ret isa SparseArrays.SparseMatrixCSC + @test ret !== model.model.C[i.value] + else + @test ret isa FillArrays.Zeros + end end + dcons = ones(LRO.num_scalars(model)) + LRO.dual_cons!(model, y, dcons, LRO.ScalarIndex) + @test dcons ≈ model.model.d_lin - model.model.C_lin' * y +end + +function schur_test(model::LRO.BufferedModelForSchur{T}) where {T} + w = rand(T, model.meta.nvar) + W = LRO.VectorizedSolution(w, model.model.dim) + for i in LRO.matrix_indices(model) + W[i] .= W[i] .+ W[i]' + end + return schur_test(model, W) +end + +function schur_test(model::LRO.Model, κ) + return schur_test(LRO.BufferedModelForSchur(model, κ)) +end + +function schur_test(model::JuMP.AbstractModel, κ) + b = _backend(model) + return schur_test(b.solver.model, κ) end diff --git a/test/maxcut.jl b/test/maxcut.jl index 09befc5..b11a619 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -3,6 +3,8 @@ # Use of this source code is governed by an MIT-style license that can be found # in the LICENSE.md file or at https://opensource.org/licenses/MIT. +import Percival + include("diff_check.jl") include(joinpath(dirname(@__DIR__), "examples", "maxcut.jl")) @@ -52,7 +54,9 @@ function test_maxcut(; is_dual, sparse, vector) @test lro_model.A isa Matrix{LowRankOpt.Factorization{Float64,F,D}} else lro_model = unsafe_backend(model).model - @test lro_model.C isa Vector{SparseMatrixCSC{T,Int64}} + @test lro_model.C isa Vector{ + FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}}, + } @test lro_model.A isa Matrix{SparseMatrixCSC{T,Int64}} solver = unsafe_backend(model).solver LRO.BurerMonteiro.set_rank!(solver.model, LRO.MatrixIndex(1), 4) @@ -84,66 +88,6 @@ end end end; -struct ConvexSolver{T} <: SolverCore.AbstractOptimizationSolver - model::LRO.Model{T} - stats::SolverCore.GenericExecutionStats{T,Vector{T},Vector{T},Any} -end - -function ConvexSolver(model::LRO.Model) - stats = SolverCore.GenericExecutionStats(model) - return ConvexSolver(model, stats) -end - -function SolverCore.solve!(::ConvexSolver, ::LRO.Model) - return -end - -function _alloc_schur_complement(model, i, Wi, H, schur_buffer) - if VERSION < v"1.11" - return - end - LRO.add_schur_complement!(model, i, Wi, H, schur_buffer) - @test 0 == - @allocated LRO.add_schur_complement!(model, i, Wi, H, schur_buffer) -end - -function schur_test(model::LRO.Model{T}, w, κ) where {T} - schur_buffer = LRO.buffer_for_schur_complement(model, κ) - jtprod_buffer = LRO.buffer_for_jtprod(model) - n = model.meta.ncon - y = rand(T, n) - - Jv = similar(y) - vJ = similar(w) - NLPModels.jprod!(model, w, w, Jv, schur_buffer[1]) - NLPModels.jtprod!(model, w, y, vJ, jtprod_buffer) - @test dot(Jv, y) ≈ dot(vJ, w) - - H = zeros(n, n) - H = LRO.schur_complement!(model, w, H, schur_buffer) - Hy = similar(y) - LRO.eval_schur_complement!(Hy, model, w, y, schur_buffer[1], jtprod_buffer) - @test Hy ≈ H * y - for i in LRO.matrix_indices(model) - Wi = @inferred w[i] - _alloc_schur_complement(model, i, Wi, H, schur_buffer) - end - for i in LRO.matrix_indices(model) - ret = LRO.dual_cons!(jtprod_buffer, model, i, y) - @test ret isa SparseMatrixCSC - end - @test LRO.dual_cons(model, LRO.ScalarIndex, y) isa SparseArrays.SparseVector -end - -function schur_test(model::LRO.Model{T}, κ) where {T} - w = rand(T, model.meta.nvar) - W = LRO.VectorizedSolution(w, model.dim) - for i in LRO.matrix_indices(model) - W[i] .= W[i] .+ W[i]' - end - return schur_test(model, W, κ) -end - @testset "ConvexSolver $T" for T in [Float32, Float64] model = maxcut(T.(weights), LRO.Optimizer{T}) set_attribute(model, "solver", ConvexSolver)