diff --git a/bench/Project.toml b/bench/Project.toml index eb3791fa3..8ad3c28b6 100644 --- a/bench/Project.toml +++ b/bench/Project.toml @@ -1,6 +1,15 @@ [deps] +Dualization = "191a621a-6537-11e9-281d-650236a99e60" DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" +LowRankOpt = "607ca3ad-272e-43c8-bcbe-fc71b56c935c" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +MultivariateBases = "be282fd4-ad43-11e9-1d11-8bd9d7e43378" MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3" MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0" +NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" +SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" +StarAlgebras = "0c0c59c1-dc5f-42e9-9a8b-b5dc384a6cd1" SumOfSquares = "4b9e565b-77fc-50a5-a571-1244f986bda1" diff --git a/bench/bm_madnlp_bench.jl b/bench/bm_madnlp_bench.jl new file mode 100644 index 000000000..ebefaf637 --- /dev/null +++ b/bench/bm_madnlp_bench.jl @@ -0,0 +1,247 @@ +# Required branches +# +# SumOfSquares → branch `bl/fft` (LowRankOpt + FFT wiring) +# LowRankOpt → branch `bl/sampling` (incl. `hprod!` sign fix) +# MultivariateBases → branch `bl/trigpolys` (TrigEvalMatrix + batched mul!) +# +# Run from this directory: +# +# julia --project=. bm_madnlp_bench.jl + +using SumOfSquares +import DynamicPolynomials +import MultivariateBases as MB +import LowRankOpt as LRO +import MadNLP +import Krylov +import NLPModels +import Dualization +import LinearAlgebra +import MathOptInterface as MOI +import Random +import StarAlgebras as SA + +include(joinpath(@__DIR__, "bm_madnlp_kkt.jl")) + +DynamicPolynomials.@polyvar x + +# Custom sub_solver factory that builds a MadNLPSolver with our BMKKTSystem. +function MadNLPMinresSolver(nlp::NLPModels.AbstractNLPModel; qlp::Bool = true, kws...) + # Deterministic init so every run sees the same KKT trace. + Random.seed!(0) + # Override the random `meta.x0` with a feasibility-friendlier starting + # point. Starting from `rand(n)/n` (small) makes the Hessian near-singular + # at iter 0, which gives Newton steps that overshoot the optimum and the + # IPM oscillates. Use `rand(n)` (BMSOSAL's choice) for a non-degenerate + # initial Hessian. + n = length(nlp.meta.x0) + nlp.meta.x0 .= rand(eltype(nlp.meta.x0), n) + # Quick check whether the BM objective is identically zero (which would + # make the dual-Newton step purely feasibility-driven). + f0 = NLPModels.obj(nlp, nlp.meta.x0) + g0 = NLPModels.grad(nlp, nlp.meta.x0) + @info "BM model at x0" obj_x0=f0 grad_norm=LinearAlgebra.norm(g0) nvar=n ncon=length(nlp.meta.y0) + return MadNLP.MadNLPSolver( + nlp; + callback = MadNLP.DenseCallback, + kkt_system = BMKKTSystem{qlp}, + linear_solver = MadNLP.LapackCPUSolver, # unused — Krylov takes over + nlp_scaling = false, # `jac_dense!` not implemented + # We can't honestly report KKT inertia from a matrix-free MINRES + # solve, so use MadNLP's inertia-free correction (Curtis–Schenk–Wächter + # heuristic). It just needs two extra `solve_kkt!`s per Newton step + # to validate that the computed direction is a descent step. + # Lanczos probe in `factorize_wrapper!` reports honest inertia, so + # MadNLP can drive `del_w` on its own — `InertiaBased` is now the + # right choice (no more `InertiaFree` curvature heuristic). + inertia_correction_method = MadNLP.InertiaBased, + # Stabilisation : la `curv_test` par défaut tolère ζ=0 (curvature ≥ 0), + # ce qui laisse passer des directions presque-singulières → pas de + # Newton qui dépassent l'optimum. Demander une curvature strictement + # positive et un régulariseur primal de base force un comportement + # type «Levenberg-Marquardt léger». + inertia_free_tol = 1e-8, + default_primal_regularization = 1e-8, + default_dual_regularization = 1e-8, + # Default least-squares dual init solves K·[d;y] = [-g;-c] once, + # which on our singular KKT can throw `y` to wild values. Starting + # from `y = 0` lets MadNLP build up the duals iteration by iteration. + dual_initialization_method = MadNLP.DualInitializeSetZero, + print_level = MadNLP.INFO, + max_iter = 500, + ) +end + +inner = optimizer_with_attributes( + LRO.Optimizer, + "solver" => LRO.BurerMonteiro.Solver, + "sub_solver" => MadNLPMinresSolver, + "ranks" => [4], + "square_scalars" => true, +) +bmlbfgs = Dualization.dual_optimizer(inner; assume_min_if_feasibility = true) + +function with_lro_bridges!(model) + backend = JuMP.backend(model) + SumOfSquares.Bridges.add_all_bridges(backend.optimizer, Float64) + MOI.Bridges.remove_bridge( + backend.optimizer, + SumOfSquares.Bridges.Constraint.ImageBridge{Float64}, + ) +end + +p = x^4 - 4x^3 - 2x^2 + 12x + 3 + +function smoke(feas::Bool) + println("\n== BMKKTSystem smoke: ", feas ? "feasibility (γ=-6)" : "max γ", " ==") + model = Model(bmlbfgs) + set_silent(model) + γ = -6 + if !feas + @variable(model, γ) + @objective(model, Max, γ) + end + with_lro_bridges!(model) + @constraint(model, p - γ in SOSCone(), zero_basis = MB.BoxSampling([-1.0], [1.0])) + optimize!(model) + println("primal_status = ", primal_status(model)) + if !feas + println("value(γ) = ", value(γ), " (expected ≈ -6)") + end +end + +smoke(true) +smoke(false) + +# Cross-check : same problem, but with Percival as the sub_solver. If Percival +# *also* diverges on `max γ`, the bug is upstream (Dualization / bridges), +# not in our MadNLP integration. +println("\n== Percival cross-check (max γ) ==") +import Percival +percival_inner = optimizer_with_attributes( + LRO.Optimizer, + "solver" => LRO.BurerMonteiro.Solver, + "sub_solver" => Percival.PercivalSolver, + "ranks" => [4], + "square_scalars" => true, +) +percival_bmlbfgs = Dualization.dual_optimizer(percival_inner; assume_min_if_feasibility = true) +let + model = Model(percival_bmlbfgs) + set_silent(model) + @variable(model, γ) + @objective(model, Max, γ) + with_lro_bridges!(model) + @constraint(model, p - γ in SOSCone(), zero_basis = MB.BoxSampling([-1.0], [1.0])) + optimize!(model) + println("primal_status = ", primal_status(model)) + println("value(γ) = ", value(γ), " (expected ≈ -6)") +end + +# Bump the BM rank to scale up the primal-variable count without changing +# the optimum (γ* = -6) or the constraint structure. Smoke uses `rank=4` +# (`n + m = 19`); we sweep up to `rank = 32` (`n + m ≈ 103`). Tests the +# dense `n+m × n+m` Bunch-Kaufman inertia probe at a non-trivial size. +# (Polynomial-degree scaling hits an unrelated upstream `BoundsError` in +# `MultivariateBases.eval_basis!` for `degree > 4`; bumping rank exercises +# the same `n + m` growth without touching the SOS-bridge path.) +function smoke_box_with_rank(rank::Int) + inner_r = optimizer_with_attributes( + LRO.Optimizer, + "solver" => LRO.BurerMonteiro.Solver, + "sub_solver" => MadNLPMinresSolver, + "ranks" => [rank], + "square_scalars" => true, + ) + opt = Dualization.dual_optimizer(inner_r; assume_min_if_feasibility = true) + model = Model(opt) + set_silent(model) + @variable(model, γ) + @objective(model, Max, γ) + with_lro_bridges!(model) + @constraint(model, p - γ in SOSCone(), zero_basis = MB.BoxSampling([-1.0], [1.0])) + t = @elapsed optimize!(model) + println("rank = ", rank, " value(γ) = ", round(value(γ), digits = 8), + " (expected ≈ -6) time = ", round(t, digits = 2), " s") + return t +end +println("\n== Scale-up via BM rank (smoke poly, expected γ = -6) ==") +for r in (4, 8, 16, 32) + smoke_box_with_rank(r) +end + +# Convert the smoke polynomial `3 + 12x − 2x² − 4x³ + x⁴` to a cos-only +# trig polynomial via Chebyshev as intermediate. Under `x = cos(θ)` we have +# `Tₖ(cos θ) = cos(kθ)`, so the Chebyshev expansion immediately gives the +# Fourier-cosine coefficients of `p(cos θ)`. Sampling `θ ∈ [−π, π]` covers +# `x ∈ [−1, 1]`, so the smoke optimum `γ* = −6` (attained at `x = −1`, i.e. +# `θ = π`) is preserved. +# Reduce to a degree-2 monomial polynomial: `x² − x − 1`. Min on `ℝ` +# is at `x = 1/2`, with `p(1/2) = −5/4 = −1.25`. Its `Monomial → Chebyshev` +# conversion gives `−0.5 T₀ − T₁ + 0.5 T₂`, so the cos-only trig form is +# `p_trig(θ) = −0.5 − cos(θ) + 0.5·cos(2θ)`, with the same minimum `−1.25` +# attained at `θ = π/3` (where `cos(θ) = 1/2`). Sticks to `monomials(x, 0:4)` +# → 5 trig basis functions (odd, dodges the upstream `isodd(n_coef)` assertion). +function _quad_as_cos_only_trig() + mon_coeffs = Float64[-1, -1, 1] # constant, x, x² + mon_sub = MB.SubBasis{MB.Monomial}(DynamicPolynomials.monomials(x, 0:2)) + cheb_full = MB.FullBasis{MB.Chebyshev}([x]) + cheb_sparse = SA.coeffs(mon_coeffs, mon_sub, cheb_full) + cheb_vals = collect(SA.values(cheb_sparse)) + @assert length(cheb_vals) == 3 "expected 3 chebyshev coefficients, got $(length(cheb_vals))" + trig_coeffs = zeros(5) # `monomials(x, 0:4)` → 5 trig basis fns + trig_coeffs[1] = cheb_vals[1] # constant + trig_coeffs[2] = cheb_vals[2] # cos(θ) + trig_coeffs[4] = cheb_vals[3] # cos(2θ) + trig_basis = MB.SubBasis{MB.Trigonometric}(DynamicPolynomials.monomials(x, 0:4)) + return MB.algebra_element(trig_coeffs, trig_basis) +end +function smoke_box_trig(solver, p_trig) + model = Model(solver) + set_silent(model) + @variable(model, γ) + @objective(model, Max, γ) + with_lro_bridges!(model) + # MultivariateBases' `Trigonometric` recurrence interprets the sample + # value as `cos(θ)` directly (see `recurrence_eval` in trigonometric.jl), + # so `BoxSampling([-1, 1])` covers the full `θ ∈ [-π, π]` and the + # Chebyshev→cos-only conversion exactly reproduces the smoke polynomial. + @constraint(model, p_trig - γ in SOSCone(), zero_basis = MB.BoxSampling([-1.0], [1.0])) + t = @elapsed optimize!(model) + println(" value(γ) = ", round(value(γ), digits = 6), + " time = ", round(t, digits = 2), " s") + return value(γ) +end +println("\n== Trigonometric quad (x²−x−1 via Chebyshev → cos-only trig) ==") +p_trig = _quad_as_cos_only_trig() +println("MadNLP+MINRES-QLP:") +γ_mad = smoke_box_trig(bmlbfgs, p_trig) +println("Percival:") +γ_perc = smoke_box_trig(percival_bmlbfgs, p_trig) +println("agreement: |Δγ| = ", round(abs(γ_mad - γ_perc), digits = 6), + " expected ≈ -1.25") + +# Should give +# == BMKKTSystem smoke: max γ s.t. p − γ ∈ SOS == +# ┌ Info: BM model at x0 +# │ obj_x0 = 0.00140682890105754 +# │ grad_norm = 0.12033747127060285 +# │ nvar = 14 +# └ ncon = 5 +# [ Info: Custom SolverCore.solve! → regular! only (skipping restoration) +# Number of nonzeros in constraint Jacobian............: 70 +# Number of nonzeros in Lagrangian Hessian.............: 105 +# +# Total number of variables............................: 14 +# variables with only lower bounds: 0 +# variables with lower and upper bounds: 0 +# variables with only upper bounds: 0 +# Total number of equality constraints.................: 5 +# Total number of inequality constraints...............: 0 +# inequality constraints with only lower bounds: 0 +# inequality constraints with lower and upper bounds: 0 +# inequality constraints with only upper bounds: 0 +# +# iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls +# 0 1.4068289e-03 9.97e+00 1.00e-01 0.00e+00 -1.0 - 0.00e+00 10 0 +# primal_status = NO_SOLUTION diff --git a/bench/bm_madnlp_kkt.jl b/bench/bm_madnlp_kkt.jl new file mode 100644 index 000000000..9f4cec2c3 --- /dev/null +++ b/bench/bm_madnlp_kkt.jl @@ -0,0 +1,482 @@ +# BMKKTSystem — matrix-free MadNLP KKT system for LowRankOpt.BurerMonteiro.Model +# Routes through `BurerMonteiro.Model`'s `NLPModels.hprod!`/`jprod!`/`jtprod!` +# (so the Stage-3 batched-FFT path stays live) and solves each Newton step +# with `Krylov.minres!` (no preconditioner). +# +# Closely modeled after `CompressedSensingIPM/src/fft_kkt.jl`. + +import MadNLP +import Krylov +import NLPModels +import LinearAlgebra +import LowRankOpt as LRO + +mutable struct BMKKTSystem{QLP,T,VT,NLP,LS} <: + MadNLP.AbstractReducedKKTSystem{T,VT,Matrix{T},MadNLP.ExactHessian{T,VT}} + nlp::NLP + n::Int + m::Int + # MadNLP-standard diagonals it updates between iterations + reg::VT + pr_diag::VT + du_diag::VT + l_diag::VT + u_diag::VT + l_lower::VT + u_lower::VT + ind_lb::Vector{Int} + ind_ub::Vector{Int} + # Current Lagrangian multipliers — stashed at `eval_lag_hess_wrapper!` time + # and read by `mul!` (the BM Hessian depends on `y` but not on `x`). + current_y::VT + # Current primal iterate — stashed at `eval_jac_wrapper!` time. The BM + # constraint c(x) is *quadratic* in `x` (since x is the rank factor / + # square-scalar pre-image), so its Jacobian depends on the linearization + # point and we must pass the actual IPM iterate to `jprod!`/`jtprod!`, + # not the Krylov direction. + current_x::VT + # Krylov state + linear_solver::LS + krylov_iterations::Vector{Int} + krylov_residuals::Vector{Float64} + # Buffers reused inside `mul!` + hv_buf::VT + jv_buf::VT + jtv_buf::VT + # Diagonal preconditioner for Krylov. Length `n+m`. Refreshed before + # each `krylov_solve!` from `reg + pr_diag + |diag(H)|` on the primal + # block and `max(|du_diag|, 1)` on the dual block. `diag(H)` is + # extracted by probing `hprod!(e_i)`; cheap for our smoke-test sizes + # but `O(n)` FFTs per IPM iter for the larger benchmark. + precond_diag::VT + precond_probe::VT # buffer for the probe vector + # Inertia estimate produced by a short Lanczos probe of `K` before each + # solve; `MadNLP.inertia(workspace)` reads this so `InertiaBased`'s + # `is_inertia_correct(n_pos, n_zero, n_neg)` can fire and trigger a + # `del_w` bump when the operator is indefinite on the augmented system. + inertia_pos::Base.RefValue{Int} + inertia_zero::Base.RefValue{Int} + inertia_neg::Base.RefValue{Int} + # Lanczos probe buffers (length `n + m`). + lanczos_v_prev::VT + lanczos_v_curr::VT + lanczos_w::VT +end +# Wrapper so Krylov treats `precond_diag` as a left preconditioner via +# `LinearAlgebra.ldiv!` (we set `ldiv = true` in `solve_kkt!`). +struct BMJacobiPrecond{VT} <: AbstractMatrix{Float64} + diag::VT +end +Base.size(M::BMJacobiPrecond) = (length(M.diag), length(M.diag)) +Base.eltype(::BMJacobiPrecond{VT}) where {VT} = eltype(VT) +LinearAlgebra.ldiv!(y::AbstractVector, M::BMJacobiPrecond, x::AbstractVector) = + (y .= x ./ M.diag; return y) +LinearAlgebra.ldiv!(M::BMJacobiPrecond, x::AbstractVector) = (x ./= M.diag; return x) + +# `qlp=true` → `Krylov.MinresQlpWorkspace` (robust on singular K, default). +# `qlp=false` → `Krylov.MinresWorkspace` (cheaper per iter, but falls back to +# the min-norm least-squares solution when K is rank-deficient — yields +# bogus Newton directions there). +function BMKKTSystem(nlp; qlp::Bool = true, T = Float64, VT = Vector{T}) + n = NLPModels.get_nvar(nlp) + m = NLPModels.get_ncon(nlp) + workspace = qlp ? + Krylov.MinresQlpWorkspace(n + m, n + m, VT) : + Krylov.MinresWorkspace(n + m, n + m, VT) + return BMKKTSystem{qlp,T,VT,typeof(nlp),typeof(workspace)}( + nlp, n, m, + VT(undef, n), # reg + VT(undef, n), # pr_diag + VT(undef, m), # du_diag + VT(undef, 0), # l_diag (no lb) + VT(undef, 0), # u_diag (no ub) + VT(undef, 0), # l_lower + VT(undef, 0), # u_lower + Int[], # ind_lb + Int[], # ind_ub + zeros(T, m), # current_y + zeros(T, n), # current_x + workspace, + Int[], + Float64[], + VT(undef, n), # hv_buf + VT(undef, m), # jv_buf + VT(undef, n), # jtv_buf + ones(T, n + m), # precond_diag + zeros(T, n), # precond_probe + Ref(n), # inertia_pos + Ref(0), # inertia_zero + Ref(m), # inertia_neg + zeros(T, n + m), # lanczos_v_prev + zeros(T, n + m), # lanczos_v_curr + zeros(T, n + m), # lanczos_w + ) +end + +function MadNLP.create_kkt_system( + ::Type{BMKKTSystem{QLP}}, + cb::MadNLP.AbstractCallback{T,VT}, + linear_solver::Type; + opt_linear_solver = MadNLP.default_options(linear_solver), + hessian_approximation = MadNLP.ExactHessian, + qn_options = MadNLP.QuasiNewtonOptions(), +) where {QLP,T,VT} + # Only supported with `square_scalars=true` (no bounds, equality + # constraints only). Verify here so misconfiguration surfaces early. + @assert isempty(cb.ind_ineq) "BMKKTSystem assumes equality-only constraints" + @assert isempty(cb.ind_lb) "BMKKTSystem assumes no lower-bound variables (square_scalars=true)" + @assert isempty(cb.ind_ub) "BMKKTSystem assumes no upper-bound variables (square_scalars=true)" + return BMKKTSystem(cb.nlp; qlp = QLP, T = T, VT = VT) +end +# Convenience dispatch — `kkt_system = BMKKTSystem` (no type param) ↔ QLP=true. +MadNLP.create_kkt_system(::Type{BMKKTSystem}, cb, ls; kws...) = + MadNLP.create_kkt_system(BMKKTSystem{true}, cb, ls; kws...) + +MadNLP.num_variables(kkt::BMKKTSystem) = kkt.n +MadNLP.get_hessian(::BMKKTSystem) = nothing +MadNLP.get_jacobian(::BMKKTSystem) = nothing + +# Krylov workspaces themselves carry no inertia. The real estimate lives on +# `BMKKTSystem` (refreshed by `_inertia_probe!` before each `solve_kkt!`); +# the workspace shim below stays a no-op so `is_inertia_correct(::BMKKTSystem, +# ...)` reads the kkt-level refs. +const _KrylovWS = Union{Krylov.MinresWorkspace,Krylov.MinresQlpWorkspace} +MadNLP.is_inertia(::_KrylovWS) = true +MadNLP.inertia(::_KrylovWS) = (0, 0, 0) +MadNLP.introduce(::Krylov.MinresQlpWorkspace) = "Krylov.MINRES-QLP" +MadNLP.introduce(::Krylov.MinresWorkspace) = "Krylov.MINRES" +MadNLP.improve!(::_KrylovWS) = true +MadNLP.factorize!(::_KrylovWS) = nothing +# MadNLP's `InertiaBased` corrector calls `inertia(kkt.linear_solver)` to +# get the inertia tuple, *then* hands it to `is_inertia_correct(kkt, ...)`. +# Our `Krylov.MinresQlpWorkspace` shim is forced to return `(0,0,0)` since +# the workspace has no view of `kkt`. So we ignore the passed args and read +# the inertia stashed on `kkt` by `_inertia_probe!`. Same for +# `should_regularize_dual`. +MadNLP.is_inertia_correct(kkt::BMKKTSystem, _, _, _) = + kkt.inertia_pos[] == kkt.n && + kkt.inertia_zero[] == 0 && + kkt.inertia_neg[] == kkt.m +MadNLP.should_regularize_dual(kkt::BMKKTSystem, _, _, _) = kkt.inertia_zero[] != 0 + +Base.eltype(::BMKKTSystem{QLP,T}) where {QLP,T} = T +Base.size(kkt::BMKKTSystem) = (kkt.n + kkt.m, kkt.n + kkt.m) +Base.size(kkt::BMKKTSystem, ::Int) = kkt.n + kkt.m + +function MadNLP.initialize!(kkt::BMKKTSystem{QLP,T}) where {QLP,T} + fill!(kkt.reg, one(T)) + fill!(kkt.pr_diag, one(T)) + fill!(kkt.du_diag, zero(T)) + fill!(kkt.current_y, zero(T)) + fill!(kkt.current_x, zero(T)) + return +end + +# Never assemble Jacobian or Hessian — but DO stash the current iterate so +# our matrix-free `mul!` linearizes the (nonlinear) BM constraint at it. +function MadNLP.eval_jac_wrapper!( + ::MadNLP.MadNLPSolver, kkt::BMKKTSystem, x::MadNLP.PrimalVector, +) + copyto!(kkt.current_x, MadNLP.full(x)) + return +end +function MadNLP.eval_lag_hess_wrapper!( + ::MadNLP.MadNLPSolver, + kkt::BMKKTSystem, + ::MadNLP.PrimalVector, + l::AbstractVector; + is_resto = false, +) + # MadNLP convention: `L(x, y) = obj_weight·f(x) + y'·c(x)` (dual feasibility + # `∇f + Jᵀy = 0`, see `IPM/kernels.jl:247`). + # NLPModels convention: `L(x, y) = obj_weight·f(x) − y'·c(x)`, so its + # `hprod!` returns `(obj_weight·∇²f − Σ y_i ∇²c_i)·v`. + # To get MadNLP's Hessian `∇²f + Σ y_i ∇²c_i` we must pass `−l`. + kkt.current_y .= .-l + return +end + +MadNLP.compress_jacobian!(::BMKKTSystem) = nothing +MadNLP.compress_hessian!(::BMKKTSystem) = nothing +MadNLP.build_kkt!(::BMKKTSystem) = nothing +function MadNLP.factorize_wrapper!( + solver::MadNLP.MadNLPSolver{T,VT,IT,KKT}, +) where {T,VT,IT,KKT<:BMKKTSystem} + MadNLP.build_kkt!(solver.kkt) + # Lanczos inertia probe runs against the regularization-augmented matvec, + # so MadNLP's subsequent `inertia(kkt)` / `is_inertia_correct(kkt, ...)` + # check honestly reflects the current `reg`/`pr_diag`/`du_diag` state. + _inertia_probe!(solver.kkt, solver.kkt.n + solver.kkt.m) + return true +end + +# Augmented KKT matvec on the `[Δx; Δy]` layout: +# yp = β·yp + α·(H·xp + Aᵀ·xd + pr_diag·xp) +# yd = β·yd + α·(A·xp − du_diag·xd) +function _kkt_apply!( + yp::AbstractVector, yd::AbstractVector, + kkt::BMKKTSystem, + xp::AbstractVector, xd::AbstractVector, + alpha::Number, beta::Number, +) + # Augmented KKT operator (no bounds, no slacks; reg + du_diag added by + # the trailing `_kktmul!`, which mirrors MadNLP's standard sparse path): + # yp = β·yp + α·(H·xp + Aᵀ·xd) + # yd = β·yd + α·(A·xp) + # Linearize at the *current IPM iterate* (`kkt.current_x`), not at the + # Krylov direction `xp` — the BM Jacobian is x-dependent (quadratic + # constraint). + NLPModels.hprod!( + kkt.nlp, kkt.current_x, kkt.current_y, xp, kkt.hv_buf; + obj_weight = one(eltype(yp)), + ) + NLPModels.jtprod!(kkt.nlp, kkt.current_x, xd, kkt.jtv_buf) + yp .= beta .* yp .+ alpha .* (kkt.hv_buf .+ kkt.jtv_buf) + + NLPModels.jprod!(kkt.nlp, kkt.current_x, xp, kkt.jv_buf) + yd .= beta .* yd .+ alpha .* kkt.jv_buf + return +end + +# Variant called by MadNLP on `AbstractKKTVector` (extra `_kktmul!` at the +# end handles the bound-mult blocks; empty in our case but kept for safety). +function MadNLP.mul!( + y::MadNLP.AbstractKKTVector, + kkt::BMKKTSystem, + x::MadNLP.AbstractKKTVector, + alpha::Number, beta::Number, +) + n, m = kkt.n, kkt.m + _x = MadNLP.full(x) + _y = MadNLP.full(y) + _kkt_apply!( + view(_y, 1:n), view(_y, (n+1):(n+m)), + kkt, + view(_x, 1:n), view(_x, (n+1):(n+m)), + alpha, beta, + ) + MadNLP._kktmul!( + y, x, + kkt.reg, kkt.du_diag, kkt.l_lower, kkt.u_lower, kkt.l_diag, kkt.u_diag, + alpha, beta, + ) + return y +end + +# Variant called by `Krylov.kmul!` on plain `Vector{T}` (length `n + m`). +# Must mirror the `AbstractKKTVector` variant — including the `reg`/`du_diag` +# contributions that the sparse path applies via `_kktmul!`. Without these, +# MadNLP's `del_w` regularization never reaches Krylov, the curvature test +# can never be satisfied, and the IPM bails into restoration. +function LinearAlgebra.mul!( + y::AbstractVector, + kkt::BMKKTSystem, + x::AbstractVector, + alpha::Number, beta::Number, +) + n, m = kkt.n, kkt.m + yp = view(y, 1:n); yd = view(y, (n+1):(n+m)) + xp = view(x, 1:n); xd = view(x, (n+1):(n+m)) + _kkt_apply!(yp, yd, kkt, xp, xd, alpha, beta) + # Match `_kktmul!` from the `AbstractKKTVector` path: add `reg` to the + # primal block and `du_diag` to the dual block. These reflect MadNLP's + # cumulative `del_w`/`del_c` Hessian/Jacobian perturbations. + yp .+= alpha .* kkt.reg .* xp + yd .+= alpha .* kkt.du_diag .* xd + return y +end + +LinearAlgebra.mul!(y::AbstractVector, kkt::BMKKTSystem, x::AbstractVector) = + LinearAlgebra.mul!(y, kkt, x, true, false) + +# Hessian-block matvec used by `InertiaFree`'s curvature test (`curv_test` +# in `MadNLP/src/IPM/solver.jl:785`). Computes `wx = (H + pr_diag) · t` +# where `t`, `wx` ∈ ℝⁿ are *primal only* (see `build_inertia_corrector`). +function MadNLP.mul_hess_blk!( + wx::AbstractVector, kkt::BMKKTSystem, t::AbstractVector, +) + NLPModels.hprod!( + kkt.nlp, kkt.current_x, kkt.current_y, t, kkt.hv_buf; + obj_weight = one(eltype(wx)), + ) + copyto!(wx, kkt.hv_buf) + wx .+= t .* kkt.pr_diag + return wx +end + +# `Aᵀ x`, linearized at the current IPM iterate (stashed in `kkt.current_x` +# by our `eval_jac_wrapper!`). Until that fix, this passed `y` as the +# evaluation point of the Jacobian — fine if `c` is linear, but `c(X) = +# jprod(model, X, X)` for our BM model is *quadratic* in `X`, so J depends +# on the linearization point and the wrong one produced bogus duals once +# `y` had non-trivial magnitude (feasibility happens to keep `y ≈ 0` so the +# bug stayed dormant; `max γ` revealed it). +function MadNLP.jtprod!(y::AbstractVector, kkt::BMKKTSystem, x::AbstractVector) + NLPModels.jtprod!(kkt.nlp, kkt.current_x, x, y) + return y +end + +import SolverCore + +# `BurerMonteiro.Solver`'s `solve!` calls +# `SolverCore.solve!(madnlp_solver, bm_model, stats; kws...)`. MadNLP's own +# entrypoints are `solve!(solver)` or `solve!(nlp, solver, stats)`; the +# (solver, nlp, stats) ordering doesn't exist. Bridge here. +# +# CompressedSensingIPM-style: only run `MadNLP.regular!`, never `robust!`. +# That bypasses restoration entirely — which is fine for us because +# restoration would call `eval_lag_hess_wrapper!`/`eval_jac_wrapper!` (which +# we stub to no-ops) and bail anyway. +# MadNLP `Status` → SolverCore status symbol (consumed by NLPModelsJuMP's +# `TERMINATION_STATUS`; `:unknown` → `MOI.OPTIMIZE_NOT_CALLED`, which makes +# JuMP throw `OptimizeNotCalled` on `value(...)` — avoid). +function _madnlp_to_solvercore_status(s::MadNLP.Status) + s == MadNLP.SOLVE_SUCCEEDED && return :first_order + s == MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL && return :acceptable + s == MadNLP.SEARCH_DIRECTION_BECOMES_TOO_SMALL && return :small_step + s == MadNLP.DIVERGING_ITERATES && return :unbounded + s == MadNLP.INFEASIBLE_PROBLEM_DETECTED && return :infeasible + s == MadNLP.MAXIMUM_ITERATIONS_EXCEEDED && return :max_iter + s == MadNLP.MAXIMUM_WALLTIME_EXCEEDED && return :max_time + s == MadNLP.USER_REQUESTED_STOP && return :user + s == MadNLP.ERROR_IN_STEP_COMPUTATION && return :neg_pred + # In-progress states leaking out (e.g. we returned mid-iteration after the + # line search failed) — treat as a stalled / slow-progress termination so + # JuMP still sees a result and can read the iterate. + s in (MadNLP.INITIAL, MadNLP.REGULAR, MadNLP.RESTORE, MadNLP.ROBUST, + MadNLP.LINESEARCH_SUCCEEDED) && return :stalled + return :exception +end + +function SolverCore.solve!( + solver::MadNLP.MadNLPSolver, + nlp::NLPModels.AbstractNLPModel, + stats::SolverCore.GenericExecutionStats; + kws..., +) + @info "Custom SolverCore.solve! → regular! only (skipping restoration)" + MadNLP.print_init(solver) + MadNLP.initialize!(solver) + try + MadNLP.regular!(solver) + catch err + @info "MadNLP.regular! threw — accepting current state" err + end + res = MadNLP.MadNLPExecutionStats(solver) + stats.solution .= res.solution + stats.objective = res.objective + stats.iter = res.iter + SolverCore.set_status!(stats, _madnlp_to_solvercore_status(res.status)) + stats.dual_feas = res.dual_feas + stats.primal_feas = res.primal_feas + stats.solver_specific[:madnlp] = res + return stats +end + +# Dense reconstruction of `K` via `n+m` standard-basis matvecs, then +# `eigvals`. Cheap for the smoke (`n+m = 19`); the trig benchmark +# (`n+m ≈ 1000+`) will need a Lanczos with full reorthogonalization +# instead. Plain three-term Lanczos on indefinite KKT loses orthogonality +# after a few iters and produces spurious eigenvalues, which made +# `is_inertia_correct` permanently false. +function _inertia_probe!(kkt::BMKKTSystem{QLP,T}, k_max::Int) where {QLP,T} + nm = kkt.n + kkt.m + Kmat = zeros(T, nm, nm) + e = kkt.lanczos_v_curr + w = kkt.lanczos_w + fill!(e, zero(T)) + for j in 1:nm + e[j] = one(T) + LinearAlgebra.mul!(w, kkt, e) + @views Kmat[:, j] .= w + e[j] = zero(T) + end + # Symmetrize defensively (drops asymmetry from accumulated FP error). + Kmat .= (Kmat .+ Kmat') ./ 2 + # Bunch-Kaufman on the symmetric `Kmat` — the signs of `D`'s diagonal + # give exact inertia by Sylvester's law, without an eigenvalue tolerance + # to tune. (The 2×2 blocks have one positive and one negative eigenvalue.) + bk = LinearAlgebra.bunchkaufman(LinearAlgebra.Symmetric(Kmat); check = false) + D = bk.D # `Tridiagonal` view on the block-diagonal D + n_pos = 0; n_neg = 0; n_zero = 0 + i = 1 + while i ≤ nm + if i < nm && D[i + 1, i] != zero(T) + # 2×2 block: one positive, one negative eigenvalue + n_pos += 1; n_neg += 1 + i += 2 + else + d = D[i, i] + if d > zero(T); n_pos += 1 + elseif d < zero(T); n_neg += 1 + else; n_zero += 1 + end + i += 1 + end + end + kkt.inertia_pos[] = n_pos + kkt.inertia_zero[] = n_zero + kkt.inertia_neg[] = n_neg + return +end + +# Hook MadNLP reads after `factorize_wrapper!`. We computed the inertia +# during `factorize!` (overridden below) so just return the cached values. +function MadNLP.inertia(kkt::BMKKTSystem) + return (kkt.inertia_pos[], kkt.inertia_zero[], kkt.inertia_neg[]) +end +MadNLP.is_inertia(::BMKKTSystem) = true + +# Probe `diag(H)` via `n` Hessian-vector products on standard basis vectors. +# Cheap for our smoke-test (`n = 14`); for the trigonometric `d = 100` +# benchmark this is `n ≈ 800` FFTs per IPM iter, still negligible compared +# to the Krylov inner loop. +function _refresh_preconditioner!(kkt::BMKKTSystem{QLP,T}) where {QLP,T} + n, m = kkt.n, kkt.m + e = kkt.precond_probe + fill!(e, zero(T)) + for i in 1:n + e[i] = one(T) + NLPModels.hprod!( + kkt.nlp, kkt.current_x, kkt.current_y, e, kkt.hv_buf; + obj_weight = one(T), + ) + # `|diag(H)| + reg + pr_diag`, guarded away from zero. Taking abs + # keeps the preconditioner SPD even where H has indefinite diagonal + # entries — for MINRES-QLP a SPD preconditioner is required. + kkt.precond_diag[i] = max(abs(kkt.hv_buf[i]) + kkt.reg[i] + kkt.pr_diag[i], 1e-12) + e[i] = zero(T) + end + for j in 1:m + kkt.precond_diag[n + j] = max(abs(kkt.du_diag[j]), one(T)) + end + return +end + +function MadNLP.solve_kkt!(kkt::BMKKTSystem, w::MadNLP.AbstractKKTVector) + MadNLP.reduce_rhs!(kkt, w) + # The `[primal; dual]` block — what MadNLP's standard + # `solve!(::AbstractReducedKKTSystem, w)` hands to its linear solver. + b = MadNLP.primal_dual(w) + rhs_norm = LinearAlgebra.norm(b) + rhs_copy = copy(b) + _refresh_preconditioner!(kkt) + M = BMJacobiPrecond(kkt.precond_diag) + Krylov.krylov_solve!( + kkt.linear_solver, kkt, b; + M = M, ldiv = true, + atol = 1e-10, rtol = 1e-8, itmax = 10 * (kkt.n + kkt.m), + verbose = 0, + ) + x = Krylov.solution(kkt.linear_solver) + # Compute true residual ‖K·x − b‖ against the operator we just gave Krylov + Kx = similar(x) + LinearAlgebra.mul!(Kx, kkt, x, true, false) # routes through `_kkt_apply!` + true_res = LinearAlgebra.norm(Kx .- rhs_copy) + nit = Krylov.iteration_count(kkt.linear_solver) + copyto!(b, x) + push!(kkt.krylov_iterations, nit) + push!(kkt.krylov_residuals, Krylov.elapsed_time(kkt.linear_solver)) + MadNLP.finish_aug_solve!(kkt, w) + return w +end diff --git a/docs/Project.toml b/docs/Project.toml index 5f70b20a2..ace27f482 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -5,13 +5,13 @@ Clarabel = "61c947e1-3e6d-4ee4-985a-eec8c727bd6e" ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4" Cyclotomics = "da8f5974-afbb-4dc8-91d8-516d5257c83b" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" -OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244" Dualization = "191a621a-6537-11e9-281d-650236a99e60" DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07" GroupsCore = "d5909c97-4eac-4ecc-a3dc-fdd0858a4120" HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327" +Hypatia = "b99e6be6-89ff-11e8-14f8-45c827f4f8f2" ImplicitPlots = "55ecb840-b828-11e9-1645-43f4a9f9ace7" Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" KnuthBendix = "c2604015-7b3d-4a30-8a26-9074551ec60a" @@ -23,11 +23,13 @@ MultivariateBases = "be282fd4-ad43-11e9-1d11-8bd9d7e43378" MultivariateMoments = "f4abf1af-0426-5881-a0da-e2f168889b5e" MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3" MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0" +OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" PermutationGroups = "8bc5a954-2dfc-11e9-10e6-cd969bffa420" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PolyJuMP = "ddf597a6-d67e-5340-b84c-e37d84115374" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" SCS = "c946c3f1-0d1f-5ce8-9dea-7daa1f7e2d13" +SDPLR = "56161740-ea4e-4253-9d15-43c62ff94d95" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" StarAlgebras = "0c0c59c1-dc5f-42e9-9a8b-b5dc384a6cd1" SumOfSquares = "4b9e565b-77fc-50a5-a571-1244f986bda1" diff --git a/docs/src/tutorials/Getting started/sampling.jl b/docs/src/tutorials/Getting started/sampling.jl new file mode 100644 index 000000000..de99d80df --- /dev/null +++ b/docs/src/tutorials/Getting started/sampling.jl @@ -0,0 +1,90 @@ +# # Sampling basis + +#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Getting started/sampling.ipynb) +#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Getting started/sampling.ipynb) +# **Contributed by**: Benoît Legat + +using Test #src +using DynamicPolynomials +using SumOfSquares +import MultivariateBases as MB + +# In this tutorial, we show how to use a different polynomial basis +# for enforcing the equality between the polynomial and its Sum-of-Squares decomposition. + +@polyvar x +p = x^4 - 4x^3 - 2x^2 + 12x + 3 + +# We want to find the minimum of the above polynomial (which is -6). + +model = Model() +set_silent(model) +γ = -6 +@variable(model, γ) +@objective(model, Max, γ) +@constraint(model, p - γ in SOSCone(), zero_basis = BoxSampling([-1], [1])) +set_optimizer() +optimize!(model) + +function test(solver, feas::Bool) + model = Model(solver) + set_silent(model) + if feas + γ = -6 + else + @variable(model, γ) + @objective(model, Max, γ) + end + @constraint(model, p - γ in SOSCone(), zero_basis = BoxSampling([-1], [1])) + optimize!(model) + @test primal_status == MOI.FEASIBLE_POINT + if !feas + @test value(γ) ≈ -6 rtol=1e-4 + end +end + + +import SDPLR + +import BMSOS + +import LowRankOpt as LRO +import Percival +import Dualization + +sdplr = optimizer_with_attributes(SDPLR.Optimizer, "maxrank" => (m, n) -> 4) +bmsos = BMSOS.Optimizer +bmlbfgs = Dualization.dual_optimizer( + optimizer_with_attributes( + LRO.Optimizer, + "solver" => LRO.BurerMonteiro.Solver, + "sub_solver" => Percival.PercivalSolver, + "ranks" => [4], + "square_scalars" => true, + ); + assume_min_if_feasibility = true, +) +test(sdplr) +test(bmsos, true) +test(bmlbfgs, true) + +function bench_rand(solver, d, B) + model = Model(solver) + set_silent(model) + p = MB.algebra_element(rand(2d+1), MB.SubBasis{B}(monomials(x, 0:2d))) + @constraint(model, p in SOSCone(), zero_basis = BoxSampling([-1], [1])) + optimize!(model) + return solve_time(model) +end + +import SCS +scs = SCS.Optimizer +bench_rand(scs, 100, MultivariateBases.Trigonometric) + +import Hypatia +hypatia = Hypatia.Optimizer +test_rand(hypatia, 100, MultivariateBases.Trigonometric) + +test_rand(sdplr, 100, MultivariateBases.Trigonometric) +test_rand(bmsos, 100, MultivariateBases.Trigonometric) +test_rand(bmlbfgs, 100, MultivariateBases.Trigonometric) diff --git a/src/Bridges/Variable/lowrank.jl b/src/Bridges/Variable/lowrank.jl index c55b85dcd..e0fedbc17 100644 --- a/src/Bridges/Variable/lowrank.jl +++ b/src/Bridges/Variable/lowrank.jl @@ -1,17 +1,31 @@ -struct LowRankBridge{T,M} <: MOI.Bridges.Variable.AbstractBridge +struct LowRankBridge{T,M,B,G,W} <: MOI.Bridges.Variable.AbstractBridge affine::Vector{MOI.ScalarAffineFunction{T}} variables::Vector{Vector{MOI.VariableIndex}} constraints::Vector{MOI.ConstraintIndex{MOI.VectorOfVariables}} - set::SOS.WeightedSOSCone{M} + set::SOS.WeightedSOSCone{M,B,G,W} end import LinearAlgebra +# The matrix returned by `MB.transformation_to(gram_basis, target_basis)`. +# `Base.promote_op` resolves this at the type level so we can declare a +# concrete `LRO.Factorization{T, SubArray{...}, Array{T,0}}` in +# `added_constrained_variable_types` — needed because PolyJuMP's +# `bridgeable` rejects `UnionAll` set types. +_transformation_type(::Type{G}, ::Type{B}) where {G,B} = + Base.promote_op(MB.transformation_to, G, B) + +# Row-view of the transformation matrix, i.e. the type of `view(U, j, :)`. +# Used as the `Factorization.factor` type so the parent `U` survives down +# to downstream batched-FFT consumers (e.g. `LowRankOpt.BurerMonteiro`). +_row_view_type(::Type{MT}) where {MT} = + Base.promote_op(view, MT, Int, Colon) + function MOI.Bridges.Variable.bridge_constrained_variable( - ::Type{LowRankBridge{T,M}}, + ::Type{LowRankBridge{T,M,B,G,W}}, model::MOI.ModelLike, - set::SOS.WeightedSOSCone{M}, -) where {T,M} + set::SOS.WeightedSOSCone{M,B,G,W}, +) where {T,M,B,G,W} variables = Vector{Vector{MOI.VariableIndex}}(undef, length(set.gram_bases)) constraints = Vector{MOI.ConstraintIndex{MOI.VectorOfVariables}}( undef, @@ -20,19 +34,26 @@ function MOI.Bridges.Variable.bridge_constrained_variable( for i in eachindex(set.gram_bases) U = MB.transformation_to(set.gram_bases[i], set.basis) weights = SA.coeffs(set.weights[i], set.basis) + # `view(U, j, :)` preserves the parent matrix `U` so that downstream + # `LowRankOpt.BurerMonteiro` can recognise the row-sharing pattern and + # use a single `mul!(buf, U, X.factor)` (one batched FFT per row of `Y` + # when `U::MB.TrigEvalMatrix`) instead of `n` separate inner products. variables[i], constraints[i] = MOI.add_constrained_variables( model, LRO.SetDotProducts{LRO.WITHOUT_SET}( SOS.matrix_cone(M, length(set.gram_bases[i])), [ LRO.TriangleVectorization( - LRO.Factorization(U[j, :], reshape(T[weights[j]], ())), + LRO.Factorization( + view(U, j, :), + reshape(T[weights[j]], ()), + ), ) for j in eachindex(set.basis) ], ), ) end - return LowRankBridge{T,M}( + return LowRankBridge{T,M,B,G,W}( [ MOI.ScalarAffineFunction( [ @@ -59,19 +80,14 @@ function MOI.Bridges.Variable.supports_constrained_variable( end function MOI.Bridges.added_constrained_variable_types( - ::Type{LowRankBridge{T,M}}, -) where {T,M} + ::Type{LowRankBridge{T,M,B,G,W}}, +) where {T,M,B,G,W} + MT = _transformation_type(G, B) + FT = _row_view_type(MT) + TVT = LRO.TriangleVectorization{T,LRO.Factorization{T,FT,Array{T,0}}} return Tuple{Type}[ - ( - LRO.SetDotProducts{ - LRO.WITHOUT_SET, - S[1], - LRO.TriangleVectorization{ - T, - LRO.Factorization{T,Vector{T},Array{T,0}}, - }, - }, - ) for S in SOS.Bridges.Constraint.constrained_variable_types(M) if + (LRO.SetDotProducts{LRO.WITHOUT_SET,S[1],TVT,Vector{TVT}},) + for S in SOS.Bridges.Constraint.constrained_variable_types(M) if S[1] == MOI.PositiveSemidefiniteConeTriangle # FIXME hack ] end @@ -82,9 +98,9 @@ end function MOI.Bridges.Variable.concrete_bridge_type( ::Type{<:LowRankBridge{T}}, - ::Type{<:SOS.WeightedSOSCone{M}}, -) where {T,M} - return LowRankBridge{T,M} + ::Type{<:SOS.WeightedSOSCone{M,B,G,W}}, +) where {T,M,B,G,W} + return LowRankBridge{T,M,B,G,W} end # Attributes, Bridge acting as a model diff --git a/src/variables.jl b/src/variables.jl index 889c6a672..f1f20695d 100644 --- a/src/variables.jl +++ b/src/variables.jl @@ -50,11 +50,11 @@ function PolyJuMP.bridges( S, LRO.TriangleVectorization{ T, - LRO.Factorization{T,Vector{T},Array{T,0}}, + LRO.Factorization{T,F,Array{T,0}}, }, }, }, -) where {S,T} +) where {S,T,F<:AbstractVector{T}} return Tuple{Type,Type}[ (LRO.Bridges.Variable.ToPositiveBridge, T), (LRO.Bridges.Variable.AppendSetBridge, T), @@ -67,11 +67,11 @@ function PolyJuMP.bridges( S, LRO.TriangleVectorization{ T, - LRO.Factorization{T,Vector{T},LRO.One{T}}, + LRO.Factorization{T,F,LRO.One{T}}, }, }, }, -) where {S,T} +) where {S,T,F<:AbstractVector{T}} return Tuple{Type,Type}[(LRO.Bridges.Variable.AppendSetBridge, T)] end function PolyJuMP.bridges( diff --git a/test/Bridges/Variable/lowrank_param.jl b/test/Bridges/Variable/lowrank_param.jl new file mode 100644 index 000000000..7bb79dd51 --- /dev/null +++ b/test/Bridges/Variable/lowrank_param.jl @@ -0,0 +1,376 @@ +# Copyright (c) 2026: Benoît Legat and contributors +# +# 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. +# +# Self-contained tests for the `Bridges.Variable.LowRankBridge` parameterisation +# rewrite (the change that lets `MultivariateBases.TrigEvalMatrix` survive +# intact to `LowRankOpt.BurerMonteiro`). +# +# What this file pins down: +# +# 1. `LowRankBridge` now carries the basis / gram-basis / weight type +# parameters `{T,M,B,G,W}` (was `{T,M}`). +# 2. `_transformation_type(G, B)` / `_row_view_type(MT)` resolve at the +# type level (via `Base.promote_op`) so `added_constrained_variable_types` +# declares the concrete SubArray factor type that +# `bridge_constrained_variable` actually produces. +# 3. The bridge stores `view(U, j, :)` (not `collect(...)`) inside each +# `LRO.Factorization`, so `parent` on the factor returns the underlying +# matrix `U` — the contract that the (future) batched-FFT path in +# BurerMonteiro relies on. +# 4. For `gram_basis = MultivariateBases.Trigonometric`, the same path +# resolves `U::MultivariateBases.TrigEvalMatrix` and the factor type is +# a `SubArray` of that matrix. + +module TestVariableLowRankParam + +using Test +using DynamicPolynomials +using JuMP +using SumOfSquares +import MultivariateBases as MB +import LowRankOpt as LRO +import StarAlgebras as SA +import MathOptInterface as MOI + +# Build a `WeightedSOSCone{M,B,G,W}` of the shape `LowRankBridge` consumes, +# parameterised by the gram-basis type so we can exercise both the dense +# (`Monomial`) and FFT (`Trigonometric`) paths. +function _weighted_sos_cone(::Type{B}; var = :t, gram_degree = 2, n_pts = 5) where {B} + @polyvar t + pts = [Float64[k / (n_pts - 1)] for k in 0:(n_pts-1)] + lag = MB.LagrangeBasis((t,), pts) + gram = MB.SubBasis{B}(monomials(t, 0:gram_degree)) + weight = MB.algebra_element(1.0 * t^0) + return SumOfSquares.WeightedSOSCone{MOI.PositiveSemidefiniteConeTriangle}( + lag, + [gram], + [weight], + ), lag, gram, weight +end + +# 1 ───────────────────────────────────────────────────────────────────────── +# Bridge type carries 5 parameters. +# ───────────────────────────────────────────────────────────────────────── + +function test_bridge_struct_has_five_type_parameters() + BridgeT = SumOfSquares.Bridges.Variable.LowRankBridge + @test BridgeT isa UnionAll + # `LowRankBridge{T}` should still be a UnionAll over the remaining 4 + # parameters — the existing bridge-graph registration uses that syntax. + @test BridgeT{Float64} isa UnionAll + # Concretely instantiate to make sure 5 type parameters resolve. + set, _, gram, weight = _weighted_sos_cone(MB.Monomial) + SetT = typeof(set) + Mp = MOI.PositiveSemidefiniteConeTriangle + Bp = typeof(set.basis) + Gp = eltype(set.gram_bases) + Wp = eltype(set.weights) + concrete = BridgeT{Float64,Mp,Bp,Gp,Wp} + @test concrete <: BridgeT + @test isconcretetype(concrete) || concrete isa DataType +end + +# 2 ───────────────────────────────────────────────────────────────────────── +# `concrete_bridge_type` round-trips: given a `WeightedSOSCone{M,B,G,W}` +# it returns the 5-param `LowRankBridge{T,M,B,G,W}`. +# ───────────────────────────────────────────────────────────────────────── + +function test_concrete_bridge_type_monomial_gram() + set, _, _, _ = _weighted_sos_cone(MB.Monomial) + bridge_t = MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ) + @test bridge_t isa Type + @test bridge_t <: SumOfSquares.Bridges.Variable.LowRankBridge{Float64} + @test length(bridge_t.parameters) == 5 + @test bridge_t.parameters[1] === Float64 + @test bridge_t.parameters[2] === MOI.PositiveSemidefiniteConeTriangle + @test bridge_t.parameters[3] === typeof(set.basis) + @test bridge_t.parameters[4] === eltype(set.gram_bases) + @test bridge_t.parameters[5] === eltype(set.weights) +end + +function test_concrete_bridge_type_trigonometric_gram() + set, _, _, _ = _weighted_sos_cone(MB.Trigonometric; gram_degree = 4, n_pts = 9) + bridge_t = MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ) + @test bridge_t <: SumOfSquares.Bridges.Variable.LowRankBridge{Float64} + @test bridge_t.parameters[4] === eltype(set.gram_bases) +end + +# 3 ───────────────────────────────────────────────────────────────────────── +# `_transformation_type(G, B)` is a compile-time map from +# `(gram_basis_type, target_basis_type)` to the matrix type that +# `MB.transformation_to` produces. This is what +# `added_constrained_variable_types` consults. +# ───────────────────────────────────────────────────────────────────────── + +function test_transformation_type_monomial_returns_matrix() + _, lag, gram, _ = _weighted_sos_cone(MB.Monomial) + MT = SumOfSquares.Bridges.Variable._transformation_type(typeof(gram), typeof(lag)) + @test MT !== Any + @test MT <: AbstractMatrix{Float64} + @test MT === typeof(MB.transformation_to(gram, lag)) +end + +function test_transformation_type_trigonometric_returns_trig_eval() + _, lag, gram, _ = _weighted_sos_cone(MB.Trigonometric; gram_degree = 4, n_pts = 9) + MT = SumOfSquares.Bridges.Variable._transformation_type(typeof(gram), typeof(lag)) + @test MT <: MB.TrigEvalMatrix{Float64} + @test MT === typeof(MB.transformation_to(gram, lag)) +end + +function test_row_view_type_matches_view_at_runtime() + _, lag, gram, _ = _weighted_sos_cone(MB.Monomial) + MT = SumOfSquares.Bridges.Variable._transformation_type(typeof(gram), typeof(lag)) + FT = SumOfSquares.Bridges.Variable._row_view_type(MT) + U = MB.transformation_to(gram, lag) + @test FT === typeof(view(U, 1, :)) +end + +function test_row_view_type_for_trig_eval() + _, lag, gram, _ = _weighted_sos_cone(MB.Trigonometric; gram_degree = 4, n_pts = 9) + MT = SumOfSquares.Bridges.Variable._transformation_type(typeof(gram), typeof(lag)) + FT = SumOfSquares.Bridges.Variable._row_view_type(MT) + U = MB.transformation_to(gram, lag) + @test FT === typeof(view(U, 1, :)) + # The factor is a `SubArray` whose parent IS the `TrigEvalMatrix` — the + # invariant downstream batched-FFT consumers will rely on. + v = view(U, 2, :) + @test parent(v) === U + @test parent(v) isa MB.TrigEvalMatrix +end + +# 4 ───────────────────────────────────────────────────────────────────────── +# `added_constrained_variable_types` declares the *same* concrete +# `SetDotProducts` type that `bridge_constrained_variable` actually adds. +# (This is the contract MOI's bridge graph relies on; the older `Vector{T}` +# declaration here would silently mismatch the new `SubArray` factor type +# and PolyJuMP would refuse to bridge.) +# ───────────────────────────────────────────────────────────────────────── + +function _added_constraint_set_type(BridgeT) + types = MOI.Bridges.added_constrained_variable_types(BridgeT) + @assert length(types) == 1 + return types[1][1] +end + +function test_added_constrained_variable_types_monomial() + set, _, _, _ = _weighted_sos_cone(MB.Monomial) + BridgeT = MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ) + declared = _added_constraint_set_type(BridgeT) + @test declared <: LRO.SetDotProducts{LRO.WITHOUT_SET} + # 4 type parameters (W, S, V, Vs) — the relaxed `SetDotProducts`. + @test length(declared.parameters) == 4 + # Factor type matches `view(::Matrix{Float64}, ::Int, ::Colon)`. + V = declared.parameters[3] + @test V <: LRO.TriangleVectorization{Float64} + Fact = V.parameters[2] + @test Fact <: LRO.Factorization{Float64} + F = Fact.parameters[2] + @test F === typeof(view(Matrix{Float64}(undef, 1, 1), 1, :)) +end + +function test_added_constrained_variable_types_trigonometric() + set, lag, gram, _ = _weighted_sos_cone( + MB.Trigonometric; + gram_degree = 4, + n_pts = 9, + ) + BridgeT = MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ) + declared = _added_constraint_set_type(BridgeT) + @test declared <: LRO.SetDotProducts{LRO.WITHOUT_SET} + V = declared.parameters[3] + Fact = V.parameters[2] + F = Fact.parameters[2] + # Factor type must be a `SubArray` *of* the `TrigEvalMatrix`, not of a + # plain `Matrix{Float64}` — this is what preserves the FFT path through + # the bridge chain. + U = MB.transformation_to(gram, lag) + @test F === typeof(view(U, 1, :)) + @test F <: SubArray + @test parent(view(U, 1, :)) isa MB.TrigEvalMatrix +end + +# 5 ───────────────────────────────────────────────────────────────────────── +# End-to-end: after `bridge_constrained_variable`, the constraint that lands +# at the inner model has the declared concrete type, and pulling out a +# constraint vector gives a `TriangleVectorization` whose `Factorization`'s +# `.factor` is a `SubArray` with a recoverable parent matrix. +# ───────────────────────────────────────────────────────────────────────── + +function _bridge_into_inner(::Type{B}; kwargs...) where {B} + set, _, _, _ = _weighted_sos_cone(B; kwargs...) + inner = MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()) + model = MOI.Bridges.Variable.SingleBridgeOptimizer{ + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + }(inner) + MOI.add_constrained_variables(model, set) + return inner, set +end + +function test_bridge_emits_declared_set_type_monomial() + inner, set = _bridge_into_inner(MB.Monomial) + BridgeT = MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ) + declared = _added_constraint_set_type(BridgeT) + cts = MOI.get(inner, MOI.ListOfConstraintTypesPresent()) + @test any(((F, S),) -> S === declared, cts) +end + +function test_bridge_emits_declared_set_type_trigonometric() + inner, set = _bridge_into_inner(MB.Trigonometric; gram_degree = 4, n_pts = 9) + BridgeT = MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ) + declared = _added_constraint_set_type(BridgeT) + cts = MOI.get(inner, MOI.ListOfConstraintTypesPresent()) + @test any(((F, S),) -> S === declared, cts) +end + +function test_bridge_stores_subarray_factor_with_recoverable_parent() + # The bridge must use `view(U, j, :)` rather than materialising the row + # so that `parent(.factor) === U` for each `j`. This is what lets the + # downstream batched-FFT path detect "all rank-1 constraints share a + # common parent matrix". + inner, set = _bridge_into_inner(MB.Trigonometric; gram_degree = 4, n_pts = 9) + declared = _added_constraint_set_type( + MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ), + ) + cis = MOI.get( + inner, + MOI.ListOfConstraintIndices{MOI.VectorOfVariables,declared}(), + ) + @test length(cis) == 1 + s = MOI.get(inner, MOI.ConstraintSet(), first(cis)) + @test length(s.vectors) == length(set.basis) + factors = [tv.matrix.factor for tv in s.vectors] + @test all(f -> f isa SubArray, factors) + parents = unique(parent.(factors)) + @test length(parents) == 1 + @test only(parents) isa MB.TrigEvalMatrix +end + +function test_bridge_factor_matches_transformation_row_monomial() + inner, set = _bridge_into_inner(MB.Monomial) + declared = _added_constraint_set_type( + MOI.Bridges.Variable.concrete_bridge_type( + SumOfSquares.Bridges.Variable.LowRankBridge{Float64}, + typeof(set), + ), + ) + ci = first( + MOI.get( + inner, + MOI.ListOfConstraintIndices{MOI.VectorOfVariables,declared}(), + ), + ) + s = MOI.get(inner, MOI.ConstraintSet(), ci) + U = MB.transformation_to(only(set.gram_bases), set.basis) + for j in eachindex(set.basis) + # Each factor is the j-th *row* of `U`, accessed lazily through a + # `SubArray` (no materialisation). + @test s.vectors[j].matrix.factor == view(U, j, :) + end +end + +# 6 ───────────────────────────────────────────────────────────────────────── +# `PolyJuMP.bridges(::Type{<:LRO.SetDotProducts{...}})` — the dispatch hook +# in `src/variables.jl` was retyped to accept *any* `F<:AbstractVector{T}` +# factor type (instead of hardcoded `Vector{T}`). Cover both the previous +# `Vector{T}` and the new `SubArray` shape; both must hit the same bridge +# list so the bridge graph routes them identically. +# ───────────────────────────────────────────────────────────────────────── + +function _polyjump_bridges_for(SetType) + return PolyJuMP.bridges(SetType) +end + +function test_polyjump_bridges_dispatches_for_vector_factor() + T = Float64 + V = LRO.TriangleVectorization{T,LRO.Factorization{T,Vector{T},Array{T,0}}} + SetType = LRO.SetDotProducts{ + LRO.WITHOUT_SET, + MOI.PositiveSemidefiniteConeTriangle, + V, + Vector{V}, + } + bridges = _polyjump_bridges_for(SetType) + @test bridges isa Vector + @test !isempty(bridges) + # The classic chain through `ToPositiveBridge` + `AppendSetBridge`. + @test any(b -> b[1] === LRO.Bridges.Variable.ToPositiveBridge, bridges) +end + +function test_polyjump_bridges_dispatches_for_subarray_factor() + T = Float64 + U = Matrix{Float64}(undef, 3, 2) + F = typeof(view(U, 1, :)) + V = LRO.TriangleVectorization{T,LRO.Factorization{T,F,Array{T,0}}} + SetType = LRO.SetDotProducts{ + LRO.WITHOUT_SET, + MOI.PositiveSemidefiniteConeTriangle, + V, + Vector{V}, + } + bridges = _polyjump_bridges_for(SetType) + @test bridges isa Vector + @test !isempty(bridges) + @test any(b -> b[1] === LRO.Bridges.Variable.ToPositiveBridge, bridges) +end + +function test_polyjump_bridges_dispatches_for_trig_eval_subarray_factor() + T = Float64 + # Points must stay in `[-1, 1]` because `Trigonometric.recurrence_eval` + # computes `sqrt(1 - cos²θ)` for sin terms. The `_weighted_sos_cone` + # helper above keeps points in that range; we mirror its layout here. + _, lag, gram, _ = _weighted_sos_cone(MB.Trigonometric; gram_degree = 4, n_pts = 9) + U = MB.transformation_to(gram, lag) + F = typeof(view(U, 1, :)) + @test F <: SubArray + @test eltype(parent(view(U, 1, :))) === T + V = LRO.TriangleVectorization{T,LRO.Factorization{T,F,Array{T,0}}} + SetType = LRO.SetDotProducts{ + LRO.WITHOUT_SET, + MOI.PositiveSemidefiniteConeTriangle, + V, + Vector{V}, + } + bridges = _polyjump_bridges_for(SetType) + @test bridges isa Vector + @test !isempty(bridges) +end + +# ───────────────────────────────────────────────────────────────────────── + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$(name)", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end + return +end + +end + +TestVariableLowRankParam.runtests() diff --git a/test/Project.toml b/test/Project.toml index 4d7c1243e..e20264312 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,22 +3,32 @@ Clarabel = "61c947e1-3e6d-4ee4-985a-eec8c727bd6e" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" Dualization = "191a621a-6537-11e9-281d-650236a99e60" DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07" +FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" +FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" Hypatia = "b99e6be6-89ff-11e8-14f8-45c827f4f8f2" +JSOSolvers = "10dff2fc-5484-5881-a0e0-c90441020f8a" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LowRankOpt = "607ca3ad-272e-43c8-bcbe-fc71b56c935c" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" MultivariateBases = "be282fd4-ad43-11e9-1d11-8bd9d7e43378" MultivariateMoments = "f4abf1af-0426-5881-a0da-e2f168889b5e" MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3" MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0" +NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" PolyJuMP = "ddf597a6-d67e-5340-b84c-e37d84115374" SemialgebraicSets = "8e049039-38e8-557d-ae3a-bc521ccf6204" +SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StarAlgebras = "0c0c59c1-dc5f-42e9-9a8b-b5dc384a6cd1" SumOfSquares = "4b9e565b-77fc-50a5-a571-1244f986bda1" SymbolicWedderburn = "858aa9a9-4c7c-4c62-b466-2421203962a2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TrigPolys = "bbdedc48-cb31-4a37-9fe3-b015aecc8dd3" TypedPolynomials = "afbbf031-7a57-5f58-a1b9-b774a0fad08d" [sources]