[WIP] Support convex qp

src/CoolPDLP.jl

@@ -16,7 +16,7 @@ using StableRNGs: StableRNG
 using LinearAlgebra: LinearAlgebra, Diagonal, axpby!, diag, dot, mul!, norm
 using Printf: @sprintf
 using Random: randn!
-using SparseArrays: SparseArrays, SparseMatrixCSC, AbstractSparseMatrix, findnz, nnz, nonzeros, nzrange, sparse, sprandn
+using SparseArrays: SparseArrays, SparseMatrixCSC, AbstractSparseMatrix, findnz, nnz, nonzeros, nzrange, sparse, sprandn, spzeros
 
 include("public.jl")
 
@@ -28,7 +28,7 @@ include("public.jl")
     include("utils/linalg.jl")
     include("utils/test.jl")
 
-    include("problems/milp.jl")
+    include("problems/program.jl")
     include("problems/solution.jl")
     include("problems/modify.jl")
 
@@ -52,7 +52,8 @@ include("MOI_wrapper.jl")
 
 export GPUSparseMatrixCOO, GPUSparseMatrixCSR, GPUSparseMatrixELL
 
-export MILP, nbvar, nbvar_int, nbvar_cont, nbcons, nbcons_eq, nbcons_ineq
+export AbstractProgram, LinearProgram, QuadraticProgram
+export nbvar, nbvar_int, nbvar_cont, nbcons, nbcons_eq, nbcons_ineq
 export PrimalDualSolution
 
 export preprocess, initialize, solve, solve!

src/MOI_wrapper.jl

@@ -71,6 +71,7 @@ const SUPPORTED_SET_TYPE{T} = Union{MOI.EqualTo{T}, MOI.GreaterThan{T}, MOI.Less
 MOI.supports_constraint(::Optimizer{T}, ::Type{MOI.VariableIndex}, ::Type{<:SUPPORTED_SET_TYPE{T}}) where {T} = true
 MOI.supports_constraint(::Optimizer{T}, ::Type{MOI.ScalarAffineFunction{T}}, ::Type{<:SUPPORTED_SET_TYPE{T}}) where {T} = true
 MOI.supports(::Optimizer{T}, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}) where {T} = true
+MOI.supports(::Optimizer{T}, ::MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{T}}) where {T} = true
 
 MOI.supports(::Optimizer, ::MOI.ObjectiveSense) = true
 MOI.supports(::Optimizer, ::MOI.Silent) = true
@@ -200,14 +201,34 @@ function MOI.optimize!(dest::Optimizer{T}, fcache::MOI.Utilities.UniversalFallba
 
     c = zeros(T, n)
     obj_constant = zero(T)
-    if cache.objective.scalar_affine !== nothing
+    if !isnothing(cache.objective.scalar_affine)
         for term in cache.objective.scalar_affine.terms
             c[term.variable.value] += term.coefficient
         end
         obj_constant = cache.objective.scalar_affine.constant
     end
+
+    Q = spzeros(T, n, n)
+    if !isnothing(cache.objective.scalar_quadratic)
+        for term in cache.objective.scalar_quadratic.quadratic_terms
+            i = term.variable_1.value
+            j = term.variable_2.value
+            # MOI stores upper-triangular entries of the symmetric Q in ½xᵀQx.
+            # Each ScalarQuadraticTerm(c, xi, xj) with i ≤ j means Q[i,j] = c.
+            Q[i, j] += term.coefficient
+            if i != j
+                Q[j, i] += term.coefficient
+            end
+        end
+        for term in cache.objective.scalar_quadratic.affine_terms
+            c[term.variable.value] += term.coefficient
+        end
+        obj_constant += cache.objective.scalar_quadratic.constant
+    end
+
     if max_sense
         c .*= -one(T)
+        Q .*= -one(T)
     end
 
     dest.sets = cache.constraints.sets
@@ -216,7 +237,11 @@ function MOI.optimize!(dest::Optimizer{T}, fcache::MOI.Utilities.UniversalFallba
     lc = cache.constraints.constants.lower
     uc = cache.constraints.constants.upper
 
-    milp = MILP(; c, lv, uv, A, lc, uc)
+    milp = if nnz(Q) == 0
+        LinearProgram(; c, lv, uv, A, lc, uc)
+    else
+        QuadraticProgram(; c, lv, uv, A, Q, lc, uc)
+    end
 
     algorithm = pop!(dest.options, :algorithm, PDLP)
 
@@ -237,15 +262,26 @@ function MOI.optimize!(dest::Optimizer{T}, fcache::MOI.Utilities.UniversalFallba
 
     dest.x = Array(sol.x)
     dest.y = Array(sol.y)
-    dest.z = proj_multiplier.(c .- milp.At * dest.y, lv, uv)
+
+    grad = if milp isa QuadraticProgram
+        c .+ milp.Q * dest.x
+    else
+        c
+    end
+    dest.z = proj_multiplier.(grad .- milp.At * dest.y, lv, uv)
 
     raw_obj = objective_value(dest.x, milp)
-    raw_dual_obj = (  # lᵀ|y|⁺ - uᵀ|y|⁻ + lᵥᵀ|z|⁺ - uᵥᵀ|z|⁻
+    raw_dual_obj_lp = (  # lᵀ|y|⁺ - uᵀ|y|⁻ + lᵥᵀ|z|⁺ - uᵥᵀ|z|⁻
         sum(safeprod_left.(lc, positive_part.(dest.y)))
             - sum(safeprod_left.(uc, negative_part.(dest.y)))
             + sum(safeprod_left.(lv, positive_part.(dest.z)))
             - sum(safeprod_left.(uv, negative_part.(dest.z)))
     )
+    raw_dual_obj = if milp isa QuadraticProgram
+        raw_dual_obj_lp - dot(dest.x, milp.Q * dest.x) / 2
+    else
+        raw_dual_obj_lp
+    end
     dest.obj_value = (max_sense ? -raw_obj : raw_obj) + obj_constant
     dest.dual_obj_value = (max_sense ? -raw_dual_obj : raw_dual_obj) + obj_constant
     dest.solve_time = stats.time_elapsed

src/algorithms/common.jl

@@ -152,7 +152,7 @@ Apply preconditioning, type conversion and device transfer to `milp_init` and `s
 Return a tuple `(milp, sol)`.
 """
 function preprocess(
-        milp_init_cpu::MILP,
+        milp_init_cpu::AbstractProgram,
         sol_init_cpu::PrimalDualSolution,
         algo::Algorithm,
     )
@@ -178,20 +178,20 @@ function initialize end
 """
     solve(milp, sol, algo)
     solve(milp, algo)
-    
+
 Solve the continuous relaxation of `milp` starting from solution `sol` using the algorithm defined by `algo`.
 
 Return a couple `(sol, stats)` where `sol` is the last solution and `stats` contains convergence information.
 """
 function solve(
-        milp_init_cpu::MILP,
+        milp_init_cpu::AbstractProgram,
         sol_init_cpu::PrimalDualSolution,
         algo::Algorithm
     )
     starting_time = time()
     milp, sol = preprocess(milp_init_cpu, sol_init_cpu, algo)
     state = initialize(milp, sol, algo; starting_time)
-    if nbcons(milp) == 0 && all(iszero, milp.c) # early exit for 0 obj/no cons
+    if milp isa LinearProgram && nbcons(milp) == 0 && all(iszero, milp.c) # early exit for 0 obj/no cons
         @. sol.x = proj_box(zero(eltype(milp.lv)), milp.lv, milp.uv)
         state.stats.termination_status = OPTIMAL
         return get_solution(state, milp), state.stats
@@ -201,7 +201,7 @@ function solve(
 end
 
 function solve(
-        milp_init_cpu::MILP,
+        milp_init_cpu::AbstractProgram,
         algo::Algorithm
     )
     sol_init_cpu = PrimalDualSolution(zero(milp_init_cpu.lv), zero(milp_init_cpu.lc))
@@ -217,7 +217,7 @@ function solve! end
 
 function termination_check!(
         state::AbstractState,
-        milp::MILP,
+        milp::AbstractProgram,
         algo::Algorithm
     )
     (; sol, scratch, stats) = state
@@ -230,6 +230,6 @@ function termination_check!(
     return stats.termination_status !== STILL_RUNNING
 end
 
-function get_solution(state::AbstractState, milp::MILP)
+function get_solution(state::AbstractState, milp::AbstractProgram)
     return unprecondition(state.sol, Preconditioner(milp))
 end

src/algorithms/pdhg.jl

@@ -31,65 +31,67 @@ $(TYPEDFIELDS)
     sol_last::PrimalDualSolution{T, V}
     "step sizes"
     step_sizes::StepSizes{T}
+    "step size parameters"
+    step_size_params::StepSizeParameters{T}
     "scratch space"
     scratch::Scratch{T, V}
     "convergence stats"
     stats::ConvergenceStats{T}
 end
 
 function initialize(
-        milp::MILP{T, V},
+        milp::AbstractProgram{T, V},
         sol::PrimalDualSolution{T, V},
         algo::Algorithm{:PDHG, T};
         starting_time::Float64
     ) where {T, V}
     sol_last = zero(sol)
-    η = fixed_stepsize(milp, algo.step_size)
-    ω = one(η)
-    step_sizes = StepSizes(; η, ω)
-    scratch = Scratch(; x = similar(sol.x), y = similar(sol.y), r = similar(sol.x))
+    η, ω, norm_A, norm_Q = init_stepsize(milp, algo.step_size)
+    step_sizes = StepSizes(; η, ω, norm_A, norm_Q)
+    step_size_params = algo.step_size
+    scratch = Scratch(sol, milp)
     stats = ConvergenceStats(T; starting_time)
-    state = PDHGState(; sol, sol_last, step_sizes, scratch, stats)
+    state = PDHGState(; sol, sol_last, step_sizes, step_size_params, scratch, stats)
     return state
 end
 
 function solve!(
         state::PDHGState,
-        milp::MILP,
+        milp::AbstractProgram,
         algo::Algorithm{:PDHG}
     )
-    prog = ProgressUnknown(desc = "PDHG iterations:", enabled = algo.generic.show_progress)
+    progress = ProgressUnknown(desc = "PDHG iterations:", enabled = algo.generic.show_progress)
     while true
         yield()
         for _ in 1:algo.generic.check_every
             step!(state, milp)
-            next!(prog; showvalues = prog_showvalues(state))
+            next!(progress; showvalues = prog_showvalues(state))
         end
         if termination_check!(state, milp, algo)
             break
         end
     end
-    finish!(prog)
+    finish!(progress)
     return state
 end
 
 function step!(
-        state::PDHGState{T, V},
-        milp::MILP{T, V},
-    ) where {T, V}
-    # switch pointers
+        state::PDHGState,
+        milp::AbstractProgram,
+    )
     state.sol, state.sol_last = state.sol_last, state.sol
-
-    (; sol, sol_last, step_sizes, scratch) = state
+    (; sol, sol_last, step_sizes, step_size_params, scratch) = state
     (; x, y) = sol_last
-    (; η, ω) = step_sizes
-    (; c, lv, uv, A, At, lc, uc) = milp
+    (; lv, uv, A, At, lc, uc) = milp
+
+    η, ω = update_step_size!(step_sizes, milp, step_size_params)
 
     τ, σ = η / ω, η * ω
 
-    # xp = proj_box.(x - τ * (c - At * y), lv, uv)
+    # xp = proj_box.(x - τ * (grad - At * y), lv, uv)
+    g = grad_x(scratch, x, milp)
     At_y = mul!(scratch.x, At, y)
-    @. sol.x = proj_box(x - τ * (c - At_y), lv, uv)
+    @. sol.x = proj_box(x - τ * (g - At_y), lv, uv)
     xdiff = @. scratch.x = 2sol.x - x
 
     # yp = y - σ * A * (2xp - x) - σ * proj_box.(inv(σ) * y - A * (2xp - x), -uc, -lc)

src/algorithms/pdlp.jl

@@ -32,6 +32,8 @@ $(TYPEDFIELDS)
     sol_restart::PrimalDualSolution{T, V}
     "step sizes"
     step_sizes::StepSizes{T}
+    "step size parameters"
+    step_size_params::StepSizeParameters{T}
     "scratch space"
     scratch::Scratch{T, V}
     "iteration counter"
@@ -43,7 +45,7 @@ $(TYPEDFIELDS)
 end
 
 function initialize(
-        milp::MILP{T, V},
+        milp::AbstractProgram{T, V},
         sol::PrimalDualSolution{T, V},
         algo::Algorithm{:PDLP, T};
         starting_time::Float64
@@ -52,59 +54,59 @@ function initialize(
     sol_avg = copy(sol)
     sol_avg_last = zero(sol)
     sol_restart = copy(sol)
-    η = fixed_stepsize(milp, algo.step_size)
-    ω = primal_weight_init(milp, algo.step_size)
-    step_sizes = StepSizes(; η, ω)
-    scratch = Scratch(; x = similar(sol.x), y = similar(sol.y), r = similar(sol.x))
+    η, ω, norm_A, norm_Q = init_stepsize(milp, algo.step_size)
+    step_sizes = StepSizes(; η, ω, norm_A, norm_Q)
+    step_size_params = algo.step_size
+    scratch = Scratch(sol, milp)
     iteration = IterationCounter(0, 0, 0)
     restart_stats = RestartStats(T)
     stats = ConvergenceStats(T; starting_time)
     state = PDLPState(;
         sol, sol_last, sol_avg, sol_avg_last, sol_restart,
-        step_sizes, scratch, iteration, restart_stats, stats
+        step_sizes, step_size_params, scratch, iteration, restart_stats, stats
     )
     return state
 end
 
 function solve!(
         state::PDLPState,
-        milp::MILP,
+        milp::AbstractProgram,
         algo::Algorithm{:PDLP}
     )
-    prog = ProgressUnknown(desc = "PDLP iterations:", enabled = algo.generic.show_progress)
+    progress = ProgressUnknown(desc = "PDLP iterations:", enabled = algo.generic.show_progress)
     while true
         yield()
         for _ in 1:algo.generic.check_every
             step!(state, milp)
-            next!(prog; showvalues = prog_showvalues(state))
+            next!(progress; showvalues = prog_showvalues(state))
         end
         if termination_check!(state, milp, algo)
             break
         elseif restart_check!(state, milp, algo)
-            restart!(state, algo)
+            restart!(state, milp, algo)
         end
     end
-    finish!(prog)
+    finish!(progress)
     return state
 end
 
 function step!(
-        state::PDLPState{T, V},
-        milp::MILP{T, V},
-    ) where {T, V}
-    # switch pointers
+        state::PDLPState,
+        milp::AbstractProgram,
+    )
     state.sol, state.sol_last = state.sol_last, state.sol
-
-    (; sol, sol_last, step_sizes, scratch) = state
+    (; sol, sol_last, step_sizes, step_size_params, scratch) = state
     (; x, y) = sol_last
-    (; η, ω) = step_sizes
-    (; c, lv, uv, A, At, lc, uc) = milp
+    (; lv, uv, A, At, lc, uc) = milp
+
+    η, ω = update_step_size!(step_sizes, milp, step_size_params)
 
     τ, σ = η / ω, η * ω
 
-    # xp = proj_box.(x - τ * (c - At * y), lv, uv)
+    # xp = proj_box.(x - τ * (grad - At * y), lv, uv)
+    g = grad_x(scratch, x, milp)
     At_y = mul!(scratch.x, At, y)
-    @. sol.x = proj_box(x - τ * (c - At_y), lv, uv)
+    @. sol.x = proj_box(x - τ * (g - At_y), lv, uv)
     xdiff = @. scratch.x = 2sol.x - x
 
     # yp = y - σ * A * (2xp - x) - σ * proj_box.(inv(σ) * y - A * (2xp - x), -uc, -lc)
@@ -132,7 +134,7 @@ end
 
 function restart_check!(
         state::PDLPState,
-        milp::MILP,
+        milp::AbstractProgram,
         algo::Algorithm{:PDLP}
     )
     (;
@@ -164,10 +166,10 @@ function restart_check!(
     return should_restart(restart_stats, step_sizes, iteration, algo.restart)
 end
 
-function restart!(state::PDLPState{T}, algo::Algorithm{:PDLP}) where {T}
+function restart!(state::PDLPState{T}, milp::AbstractProgram, algo::Algorithm{:PDLP}) where {T}
     (;
         sol, sol_avg, sol_restart,
-        step_sizes, iteration, scratch, restart_stats,
+        step_sizes, step_size_params, iteration, scratch, restart_stats,
     ) = state
 
     # identify candidate
@@ -179,8 +181,9 @@ function restart!(state::PDLPState{T}, algo::Algorithm{:PDLP}) where {T}
     # update step sizes (must be done before losing previous restart)
     step_sizes.η_sum = zero(T)
     step_sizes.ω = primal_weight_update!(
-        scratch, step_sizes, sol_cand, sol_restart, algo.step_size
+        scratch, step_sizes, sol_cand, sol_restart, step_size_params
     )
+    update_step_size!(step_sizes, milp, step_size_params)
     # update solutions
     copy!(sol, sol_cand)
     zero!(sol_avg)