From 227e82d5516c58a70f09f8442d7ee87bb1815bb1 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Mon, 14 Nov 2022 19:45:41 +0100
Subject: [PATCH] make augmented compute_sos fast

---
 src/certify.jl     | 90 ++++++++++++++++------------------------------
 test/1712.07167.jl | 37 ++++++-------------
 2 files changed, 40 insertions(+), 87 deletions(-)

diff --git a/src/certify.jl b/src/certify.jl
index 089621d..cdebcc9 100644
--- a/src/certify.jl
+++ b/src/certify.jl
@@ -5,80 +5,50 @@ function augment_columns!(Q::AbstractMatrix)
     return Q
 end
 
-function _fma_SOS_thr!(
-    result::AbstractVector{T},
-    mstructure::AbstractMatrix{<:Integer},
-    Q::AbstractMatrix{T},
-    acc_matrix=zeros(T, size(mstructure)...),
-) where {T}
+function __sos_via_sqr!(
+    res::StarAlgebras.AlgebraElement,
+    P::AbstractMatrix;
+    augmented::Bool
+)
+    StarAlgebras.zero!(res)
+    A = parent(res)
+    b = basis(A)
+    @assert size(A.mstructure) == size(P)
+    e = b[one(b[1])]
 
-    s1, s2 = size(mstructure)
-
-    @inbounds for k = 1:s2
-        let k = k, s1 = s1, s2 = s2, Q = Q, acc_matrix = acc_matrix
-            Threads.@threads for j = 1:s2
-                for i = 1:s1
-                    @inbounds acc_matrix[i, j] =
-                        muladd(Q[i, k], Q[j, k], acc_matrix[i, j])
-                end
+    for i in axes(A.mstructure, 1)
+        x = StarAlgebras._istwisted(A.mstructure) ? StarAlgebras.star(b[i]) : b[i]
+        for j in axes(A.mstructure, 2)
+            p = P[i, j]
+            xy = b[A.mstructure[i, j]]
+            # either result += P[x,y]*(x*y)
+            res[xy] += p
+            if augmented
+                # or result += P[x,y]*(1-x)*(1-y) == P[x,y]*(2-x-y+xy)
+                y = b[j]
+                res[e] += p
+                res[x] -= p
+                res[y] -= p
             end
         end
     end
 
-    @inbounds for j = 1:s2
-        for i = 1:s1
-            result[mstructure[i, j]] += acc_matrix[i, j]
-        end
-    end
-
-    return result
+    return res
 end
 
-function _cnstr_sos!(res::StarAlgebras.AlgebraElement, Q::AbstractMatrix, cnstrs)
+function __sos_via_cnstr!(res::StarAlgebras.AlgebraElement, Q²::AbstractMatrix, cnstrs)
     StarAlgebras.zero!(res)
-    Q² = Q' * Q
     for (g, A_g) in cnstrs
         res[g] = dot(A_g, Q²)
     end
     return res
 end
 
-function _augmented_sos!(res::StarAlgebras.AlgebraElement, Q::AbstractMatrix)
-    A = parent(res)
-    StarAlgebras.zero!(res)
-    Q² = Q' * Q
-
-    N = LinearAlgebra.checksquare(A.mstructure)
-    augmented_basis = [A(1) - A(b) for b in @view basis(A)[1:N]]
-    tmp = zero(res)
-
-    for (j, y) in enumerate(augmented_basis)
-        for (i, x) in enumerate(augmented_basis)
-            # res += Q²[i, j] * x * y
-
-            StarAlgebras.mul!(tmp, x, y)
-            StarAlgebras.mul!(tmp, tmp, Q²[i, j])
-            StarAlgebras.add!(res, res, tmp)
-        end
-    end
-    return res
-end
-
 function compute_sos(A::StarAlgebras.StarAlgebra, Q::AbstractMatrix; augmented::Bool)
-    if augmented
-        z = zeros(eltype(Q), length(basis(A)))
-        res = StarAlgebras.AlgebraElement(z, A)
-        return _augmented_sos!(res, Q)
-        cnstrs = constraints(basis(A), A.mstructure; augmented=true)
-        return _cnstr_sos!(res, Q, cnstrs)
-    else
-        @assert size(A.mstructure) == size(Q)
-        z = zeros(eltype(Q), length(basis(A)))
-
-        _fma_SOS_thr!(z, A.mstructure, Q)
-
-        return StarAlgebras.AlgebraElement(z, A)
-    end
+    Q² = Q' * Q
+    res = StarAlgebras.AlgebraElement(zeros(eltype(Q²), length(basis(A))), A)
+    res = __sos_via_sqr!(res, Q², augmented=augmented)
+    return res
 end
 
 function sufficient_λ(residual::StarAlgebras.AlgebraElement, λ; halfradius)
@@ -159,7 +129,7 @@ function certify_solution(
         true, compute_sos(parent(elt), Q_int, augmented=augmented)
     end
 
-    @info "Checking in $(eltype(sos_int)) arithmetic with" λ
+    @info "Checking in $(eltype(sos_int)) arithmetic with" λ_int
 
     λ_certified =
         sufficient_λ(elt, orderunit, λ_int, sos_int, halfradius=halfradius)
diff --git a/test/1712.07167.jl b/test/1712.07167.jl
index 982806e..9a1adcc 100644
--- a/test/1712.07167.jl
+++ b/test/1712.07167.jl
@@ -6,37 +6,20 @@ function check_positivity(elt, unit, wd; upper_bound=Inf, halfradius=2, optimize
     @time status, _ = PropertyT.solve(sos_problem, optimizer)
 
     Q = let Ps = Ps
-        flPs = [real.(sqrt(JuMP.value.(P))) for P in Ps]
-        PropertyT.reconstruct(flPs, wd)
+        Qs = [real.(sqrt(JuMP.value.(P))) for P in Ps]
+        PropertyT.reconstruct(Qs, wd)
     end
 
     λ = JuMP.value(sos_problem[:λ])
 
-    sos = let RG = parent(elt), Q = Q
-        z = zeros(eltype(Q), length(basis(RG)))
-        res = AlgebraElement(z, RG)
-        cnstrs = PropertyT.constraints(basis(RG), RG.mstructure, augmented=true)
-        PropertyT._cnstr_sos!(res, Q, cnstrs)
-    end
-
-    residual = elt - λ * unit - sos
-    λ_fl = PropertyT.sufficient_λ(residual, λ, halfradius=2)
-
-    λ_fl < 0 && return status, false, λ_fl
-
-    sos = let RG = parent(elt), Q = [PropertyT.IntervalArithmetic.@interval(q) for q in Q]
-        z = zeros(eltype(Q), length(basis(RG)))
-        res = AlgebraElement(z, RG)
-        cnstrs = PropertyT.constraints(basis(RG), RG.mstructure, augmented=true)
-        PropertyT._cnstr_sos!(res, Q, cnstrs)
-    end
-
-    λ_int = PropertyT.IntervalArithmetic.@interval(λ)
-
-    residual_int = elt - λ_int * unit - sos
-    λ_int = PropertyT.sufficient_λ(residual_int, λ_int, halfradius=2)
-
-    return status, λ_int > 0, PropertyT.IntervalArithmetic.inf(λ_int)
+    certified, λ_cert = PropertyT.certify_solution(
+        elt,
+        unit,
+        λ,
+        Q,
+        halfradius=2
+    )
+    return status, certified, λ_cert
 end
 
 @testset "1712.07167 Examples" begin