diff --git a/src/BurerMonteiro/BurerMonteiro.jl b/src/BurerMonteiro/BurerMonteiro.jl new file mode 100644 index 0000000..fb99c96 --- /dev/null +++ b/src/BurerMonteiro/BurerMonteiro.jl @@ -0,0 +1,15 @@ +module BurerMonteiro + +import LinearAlgebra +import FillArrays +import SolverCore +import NLPModels +import MathOptInterface as MOI +import NLPModelsJuMP +import LowRankOpt as LRO + +include("solution.jl") +include("model.jl") +include("solver.jl") + +end diff --git a/src/BurerMonteiro.jl b/src/BurerMonteiro/model.jl similarity index 59% rename from src/BurerMonteiro.jl rename to src/BurerMonteiro/model.jl index 9a41b7b..3f24cd0 100644 --- a/src/BurerMonteiro.jl +++ b/src/BurerMonteiro/model.jl @@ -1,40 +1,3 @@ -module BurerMonteiro - -import LinearAlgebra -import FillArrays -import SolverCore -import NLPModels -import MathOptInterface as MOI -import NLPModelsJuMP -import LowRankOpt as LRO - -# `Dimensions{false}` means that nonnegative scalars have a zero lower bound -# `Dimensions{true}` means that nonnegative scalars are the square of a free variable -struct Dimensions{S} - num_scalars::Int64 - side_dimensions::Vector{Int64} - ranks::Vector{Int64} - offsets::Vector{Int64} -end - -function Dimensions{S}(model::LRO.Model, ranks) where {S} - side_dimensions = - [LRO.side_dimension(model, i) for i in LRO.matrix_indices(model)] - num_scalars = LRO.num_scalars(model) - offsets = num_scalars .+ [0; cumsum(side_dimensions .* ranks)] - return Dimensions{S}(num_scalars, side_dimensions, ranks, offsets) -end - -Base.length(d::Dimensions) = d.offsets[end] - -function set_rank!(d::Dimensions, i::LRO.MatrixIndex, rank) - d.ranks[i.value] = rank - for j in (i.value+1):length(d.offsets) - d.offsets[j] = d.offsets[j-1] + d.side_dimensions[j-1] * d.ranks[j-1] - end - return -end - mutable struct Model{S,T,CT,AT} <: NLPModels.AbstractNLPModel{T,Vector{T}} model::LRO.Model{T,CT,AT} dim::Dimensions{S} @@ -82,71 +45,9 @@ function set_rank!(model::Model, i::LRO.MatrixIndex, r) return end -struct Solution{S,T,VT<:AbstractVector{T}} <: AbstractVector{T} - x::VT - dim::Dimensions{S} -end - -struct _OuterProduct{S,T,UT<:AbstractVector{T},VT<:AbstractVector{T}} <: - AbstractVector{T} - x::Solution{S,T,VT} - v::Solution{S,T,UT} -end - -Base.eltype(::Type{<:Union{Solution{S,T},_OuterProduct{S,T}}}) where {S,T} = T -Base.eltype(x::Union{Solution,_OuterProduct}) = eltype(typeof(x)) - -Base.size(s::Solution) = size(s.x) -Base.getindex(s::Solution, i::Integer) = getindex(s.x, i) - -Base.size(s::_OuterProduct) = size(s.x) -function Base.show(io::IO, s::_OuterProduct) - print(io, "_OuterProduct(") - print(io, s.x) - print(io, ", ") - print(io, s.v) - print(io, ")") - return -end - -function LRO.left_factor(s::Solution, ::Type{LRO.ScalarIndex}) - return view(s.x, Base.OneTo(s.dim.num_scalars)) -end - -function Base.getindex(s::Solution{false}, ::Type{LRO.ScalarIndex}) - return LRO.left_factor(s::Solution, LRO.ScalarIndex) -end - -function Base.getindex(s::Solution{true,T}, ::Type{LRO.ScalarIndex}) where {T} - s = LRO.left_factor(s::Solution, LRO.ScalarIndex) - return MOI.Utilities.VectorLazyMap{T}(abs2, s) -end - -function Base.getindex(s::_OuterProduct{false}, ::Type{LRO.ScalarIndex}) - return getindex(s.v, LRO.ScalarIndex) -end - -function Base.getindex(s::_OuterProduct{true}, ::Type{LRO.ScalarIndex}) - # TODO Lazy - return 2 .* LRO.left_factor(s.x, LRO.ScalarIndex) .* - LRO.left_factor(s.v, LRO.ScalarIndex) -end - -function Base.getindex(s::Solution, mi::LRO.MatrixIndex) - i = mi.value - U = reshape( - view(s.x, (1+s.dim.offsets[i]):s.dim.offsets[i+1]), - s.dim.side_dimensions[i], - s.dim.ranks[i], - ) - return LRO.positive_semidefinite_factorization(U) -end - -function Base.getindex(s::_OuterProduct{S,T}, i::LRO.MatrixIndex) where {S,T} - U = s.x[i].factor - V = s.v[i].factor - return LRO.AsymmetricFactorization(U, V, FillArrays.Fill(T(2), size(U, 2))) -end +####################### +###### Objective ###### +####################### function NLPModels.obj(model::Model, x::AbstractVector) return NLPModels.obj(model.model, Solution(x, model.dim)) @@ -196,6 +97,10 @@ function gprod(model::Model, x::AbstractVector, v::AbstractVector) return NLPModels.obj(model.model, _OuterProduct(X, V)) end +######################### +###### Constraints ###### +######################### + function NLPModels.cons!(model::Model, x::AbstractVector, cx::AbstractVector) X = Solution(x, model.dim) # We don't call `cons!` as we don't want to include `-b` since the constraint @@ -204,6 +109,10 @@ function NLPModels.cons!(model::Model, x::AbstractVector, cx::AbstractVector) return NLPModels.jprod!(model.model, X, X, cx) end +####################### +###### J product ###### +####################### + function NLPModels.jprod!( model::Model, x::AbstractVector, @@ -217,6 +126,10 @@ function NLPModels.jprod!( return NLPModels.jprod!(model.model, X, _OuterProduct(X, V), Jv) end +######################## +###### Jáµ€ product ###### +######################## + function jtprod!( model::Model{false}, _, @@ -278,6 +191,10 @@ function NLPModels.jtprod!( return Jtv end +####################### +###### H product ###### +####################### + function NLPModels.hprod!( ::Model{false}, ::AbstractVector, @@ -342,46 +259,3 @@ function NLPModels.hprod!( end return Hv end - -struct Solver{S,T,CT,AT,ST} <: SolverCore.AbstractOptimizationSolver - model::Model{S,T,CT,AT} - solver::ST - stats::SolverCore.GenericExecutionStats{T,Vector{T},Vector{T},Any} -end - -function Solver( - src::LRO.Model; - sub_solver, - ranks, - square_scalars = false, - kws..., -) - model = Model{square_scalars}(src, ranks) - solver = sub_solver(model; kws...) - stats = SolverCore.GenericExecutionStats(model) - return Solver(model, solver, stats) -end - -function SolverCore.solve!( - solver::Solver, - model::NLPModels.AbstractNLPModel; # Same as `solver.model.model` - kws..., -) - return SolverCore.solve!(solver.solver, solver.model, solver.stats; kws...) -end - -function MOI.get(solver::Solver, attr::MOI.SolverName) - return "BurerMonteiro with " * MOI.get(solver.solver, attr) -end - -function MOI.get(solver::Solver, ::LRO.ConvexTerminationStatus) - return NLPModelsJuMP.TERMINATION_STATUS[solver.stats.status] - # TODO if the dual is feasible, we can still claim that we found the optimal - # and turn `LOCALLY_SOLVED` into `OPTIMAL` -end - -function MOI.get(solver::Solver, ::LRO.Solution) - return Solution(solver.stats.solution, solver.model.dim) -end - -end diff --git a/src/BurerMonteiro/solution.jl b/src/BurerMonteiro/solution.jl new file mode 100644 index 0000000..dbc9a4a --- /dev/null +++ b/src/BurerMonteiro/solution.jl @@ -0,0 +1,92 @@ +# `Dimensions{false}` means that nonnegative scalars have a zero lower bound +# `Dimensions{true}` means that nonnegative scalars are the square of a free variable +struct Dimensions{S} + num_scalars::Int64 + side_dimensions::Vector{Int64} + ranks::Vector{Int64} + offsets::Vector{Int64} +end + +function Dimensions{S}(model::LRO.Model, ranks) where {S} + side_dimensions = + [LRO.side_dimension(model, i) for i in LRO.matrix_indices(model)] + num_scalars = LRO.num_scalars(model) + offsets = num_scalars .+ [0; cumsum(side_dimensions .* ranks)] + return Dimensions{S}(num_scalars, side_dimensions, ranks, offsets) +end + +Base.length(d::Dimensions) = d.offsets[end] + +function set_rank!(d::Dimensions, i::LRO.MatrixIndex, rank) + d.ranks[i.value] = rank + for j in (i.value+1):length(d.offsets) + d.offsets[j] = d.offsets[j-1] + d.side_dimensions[j-1] * d.ranks[j-1] + end + return +end + +struct Solution{S,T,VT<:AbstractVector{T}} <: AbstractVector{T} + x::VT + dim::Dimensions{S} +end + +struct _OuterProduct{S,T,UT<:AbstractVector{T},VT<:AbstractVector{T}} <: + AbstractVector{T} + x::Solution{S,T,VT} + v::Solution{S,T,UT} +end + +Base.eltype(::Type{<:Union{Solution{S,T},_OuterProduct{S,T}}}) where {S,T} = T +Base.eltype(x::Union{Solution,_OuterProduct}) = eltype(typeof(x)) + +Base.size(s::Solution) = size(s.x) +Base.getindex(s::Solution, i::Integer) = getindex(s.x, i) + +Base.size(s::_OuterProduct) = size(s.x) +function Base.show(io::IO, s::_OuterProduct) + print(io, "_OuterProduct(") + print(io, s.x) + print(io, ", ") + print(io, s.v) + print(io, ")") + return +end + +function LRO.left_factor(s::Solution, ::Type{LRO.ScalarIndex}) + return view(s.x, Base.OneTo(s.dim.num_scalars)) +end + +function Base.getindex(s::Solution{false}, ::Type{LRO.ScalarIndex}) + return LRO.left_factor(s::Solution, LRO.ScalarIndex) +end + +function Base.getindex(s::Solution{true,T}, ::Type{LRO.ScalarIndex}) where {T} + s = LRO.left_factor(s::Solution, LRO.ScalarIndex) + return MOI.Utilities.VectorLazyMap{T}(abs2, s) +end + +function Base.getindex(s::_OuterProduct{false}, ::Type{LRO.ScalarIndex}) + return getindex(s.v, LRO.ScalarIndex) +end + +function Base.getindex(s::_OuterProduct{true}, ::Type{LRO.ScalarIndex}) + # TODO Lazy + return 2 .* LRO.left_factor(s.x, LRO.ScalarIndex) .* + LRO.left_factor(s.v, LRO.ScalarIndex) +end + +function Base.getindex(s::Solution, mi::LRO.MatrixIndex) + i = mi.value + U = reshape( + view(s.x, (1+s.dim.offsets[i]):s.dim.offsets[i+1]), + s.dim.side_dimensions[i], + s.dim.ranks[i], + ) + return LRO.positive_semidefinite_factorization(U) +end + +function Base.getindex(s::_OuterProduct{S,T}, i::LRO.MatrixIndex) where {S,T} + U = s.x[i].factor + V = s.v[i].factor + return LRO.AsymmetricFactorization(U, V, FillArrays.Fill(T(2), size(U, 2))) +end diff --git a/src/BurerMonteiro/solver.jl b/src/BurerMonteiro/solver.jl new file mode 100644 index 0000000..737f0ba --- /dev/null +++ b/src/BurerMonteiro/solver.jl @@ -0,0 +1,40 @@ +struct Solver{S,T,CT,AT,ST} <: SolverCore.AbstractOptimizationSolver + model::Model{S,T,CT,AT} + solver::ST + stats::SolverCore.GenericExecutionStats{T,Vector{T},Vector{T},Any} +end + +function Solver( + src::LRO.Model; + sub_solver, + ranks, + square_scalars = false, + kws..., +) + model = Model{square_scalars}(src, ranks) + solver = sub_solver(model; kws...) + stats = SolverCore.GenericExecutionStats(model) + return Solver(model, solver, stats) +end + +function SolverCore.solve!( + solver::Solver, + model::NLPModels.AbstractNLPModel; # Same as `solver.model.model` + kws..., +) + return SolverCore.solve!(solver.solver, solver.model, solver.stats; kws...) +end + +function MOI.get(solver::Solver, attr::MOI.SolverName) + return "BurerMonteiro with " * MOI.get(solver.solver, attr) +end + +function MOI.get(solver::Solver, ::LRO.ConvexTerminationStatus) + return NLPModelsJuMP.TERMINATION_STATUS[solver.stats.status] + # TODO if the dual is feasible, we can still claim that we found the optimal + # and turn `LOCALLY_SOLVED` into `OPTIMAL` +end + +function MOI.get(solver::Solver, ::LRO.Solution) + return Solution(solver.stats.solution, solver.model.dim) +end diff --git a/src/LowRankOpt.jl b/src/LowRankOpt.jl index cf56cd8..bfedd42 100644 --- a/src/LowRankOpt.jl +++ b/src/LowRankOpt.jl @@ -22,7 +22,7 @@ include("model.jl") include("buffer.jl") include("schur.jl") include("MOI_wrapper.jl") -include("BurerMonteiro.jl") include("errors.jl") +include("BurerMonteiro/BurerMonteiro.jl") end # module LowRankOpt