Skip to content

Refactor BurerMonteiro into subfolder #54

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
blegat merged 1 commit into from
Jul 31, 2025
Merged

Refactor BurerMonteiro into subfolder #54

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions src/BurerMonteiro/BurerMonteiro.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
164 changes: 19 additions & 145 deletions src/BurerMonteiro.jl → src/BurerMonteiro/model.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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}
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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},
_,
Expand Down Expand Up @@ -278,6 +191,10 @@ function NLPModels.jtprod!(
return Jtv
end

#######################
###### H product ######
#######################

function NLPModels.hprod!(
::Model{false},
::AbstractVector,
Expand Down Expand Up @@ -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
92 changes: 92 additions & 0 deletions src/BurerMonteiro/solution.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
40 changes: 40 additions & 0 deletions src/BurerMonteiro/solver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion src/LowRankOpt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading