Refactor how extensions for sparse square root are handled

Author: odow

Project.toml

@@ -4,7 +4,6 @@ version = "1.49.0"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 CodecBzip2 = "523fee87-0ab8-5b00-afb7-3ecf72e48cfd"
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -20,6 +19,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
+CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
 
 [extensions]
@@ -38,7 +38,7 @@ LDLFactorizations = "0.10"
 LinearAlgebra = "1"
 MutableArithmetics = "1"
 NaNMath = "0.3, 1"
-OrderedCollections = "1"
+OrderedCollections = "1.1"
 ParallelTestRunner = "2.4.1"
 PrecompileTools = "1"
 Printf = "1"

ext/MathOptInterfaceCliqueTreesExt.jl

@@ -6,30 +6,28 @@
 
 module MathOptInterfaceCliqueTreesExt
 
-import CliqueTrees.Multifrontal: ChordalCholesky
-import LinearAlgebra: RowMaximum, cholesky!
+import CliqueTrees
+import LinearAlgebra
 import MathOptInterface as MOI
-import SparseArrays: findnz, sparse
+import SparseArrays
 
-function MOI.Bridges.Constraint.compute_sparse_sqrt_fallback(
-    Q::AbstractMatrix,
-    ::F,
-    ::S,
-) where {F<:MOI.ScalarQuadraticFunction,S<:MOI.AbstractSet}
-    G = cholesky!(ChordalCholesky(Q), RowMaximum())
-    U = sparse(G.U) * G.P
+MOI.Utilities.is_defined(::MOI.Utilities.CliqueTreesExt) = true
 
-    # Verify the factorization reconstructs Q. This catches indefinite
-    # matrices where the diagonal is all zeros (e.g., [0 -1; -1 0]).
+function MOI.Utilities.compute_sparse_sqrt(
+    ::MOI.Utilities.CliqueTreesExt,
+    Q::AbstractMatrix,
+)
+    G = LinearAlgebra.cholesky!(
+        CliqueTrees.Multifrontal.ChordalCholesky(Q),
+        LinearAlgebra.RowMaximum(),
+    )
+    U = SparseArrays.sparse(G.U) * G.P
+    # Verify the factorization reconstructs Q. We don't have something like
+    # LinearAlgebra.issuccess(G)
     if !isapprox(Q, U' * U; atol = 1e-10)
-        msg = """
-        Unable to transform a quadratic constraint into a SecondOrderCone
-        constraint because the quadratic constraint is not convex.
-        """
-        throw(MOI.UnsupportedConstraint{F,S}(msg))
+        return nothing
     end
-
-    return findnz(U)
+    return SparseArrays.findnz(U)
 end
 
 end  # module

ext/MathOptInterfaceLDLFactorizationsExt.jl

@@ -11,39 +11,19 @@ import LinearAlgebra
 import MathOptInterface as MOI
 import SparseArrays
 
-# The type signature of this function is not important, so long as it is more
-# specific than the (untyped) generic fallback with the error pointing to
-# LDLFactorizations.jl
-function MOI.Bridges.Constraint.compute_sparse_sqrt_fallback(
+MOI.Utilities.is_defined(::MOI.Utilities.LDLFactorizationsExt) = true
+
+function MOI.Utilities.compute_sparse_sqrt(
+    ::MOI.Utilities.LDLFactorizationsExt,
     Q::AbstractMatrix,
-    ::F,
-    ::S,
-) where {F<:MOI.ScalarQuadraticFunction,S<:MOI.AbstractSet}
+)
     n = LinearAlgebra.checksquare(Q)
     factor = LDLFactorizations.ldl(Q)
-    # Ideally we should use `LDLFactorizations.factorized(factor)` here, but it
-    # has some false negatives. Instead we check that the factorization appeared
-    # to work. This is a heuristic. There might be other cases where check is
-    # insufficient.
-    if minimum(factor.D) < 0 || any(issubnormal, factor.D)
-        msg = """
-        Unable to transform a quadratic constraint into a SecondOrderCone
-        constraint because the quadratic constraint is not convex.
-        """
-        throw(MOI.UnsupportedConstraint{F,S}(msg))
-    end
-    # We have Q = P' * L * D * L' * P. We want to find Q = U' * U, so
-    # U = sqrt(D) * L' * P. First, compute L'. Note I and J are reversed:
-    J, I, V = SparseArrays.findnz(factor.L)
-    # Except L doesn't include the identity along the diagonal. Add it back.
-    append!(J, 1:n)
-    append!(I, 1:n)
-    append!(V, ones(n))
-    # Now scale by sqrt(D)
-    for (k, i) in enumerate(I)
-        V[k] *= sqrt(factor.D[i, i])
+    if !LDLFactorizations.factorized(factor) || minimum(factor.D) < 0
+        return nothing
     end
-    # Finally, permute the columns of L'. The rows stay in the same order.
+    L = sqrt.(factor.D) * LinearAlgebra.UnitLowerTriangular(factor.L)
+    J, I, V = SparseArrays.findnz(SparseArrays.sparse(L))
     return I, factor.P[J], V
 end
 

src/Bridges/Constraint/bridges/QuadtoSOCBridge.jl

@@ -60,8 +60,8 @@ end
 const QuadtoSOC{T,OT<:MOI.ModelLike} =
     SingleBridgeOptimizer{QuadtoSOCBridge{T},OT}
 
-function compute_sparse_sqrt_fallback(Q, ::F, ::S) where {F,S}
-    msg = """
+function _error_msg(::MOI.Utilities.LDLFactorizationsExt)
+    return """
     Unable to transform a quadratic constraint into a SecondOrderCone
     constraint because the quadratic constraint is not strongly convex and
     our Cholesky decomposition failed.
@@ -76,36 +76,41 @@ function compute_sparse_sqrt_fallback(Q, ::F, ::S) where {F,S}
     LDLFactorizations.jl is not included by default because it is licensed
     under the LGPL.
     """
-    return throw(MOI.AddConstraintNotAllowed{F,S}(msg))
 end
 
-function compute_sparse_sqrt(Q, func, set)
-    # There's a big try-catch here because Cholesky can fail even if
-    # `check = false`. As one example, it currently (v1.12) fails with
-    # `BigFloat`. Similarly, we want to guard against errors in
-    # `compute_sparse_sqrt_fallback`.
-    #
-    # The try-catch isn't a performance concern because the alternative is not
-    # being able to reformulate the problem.
-    try
-        factor = LinearAlgebra.cholesky(Q; check = false)
-        if !LinearAlgebra.issuccess(factor)
-            return compute_sparse_sqrt_fallback(Q, func, set)
-        end
-        L, p = SparseArrays.sparse(factor.L), factor.p
-        # We have Q = P' * L * L' * P. We want to find Q = U' * U, so U = L' * P
-        # First, compute L'. Note I and J are reversed
-        J, I, V = SparseArrays.findnz(L)
-        # Then, we want to permute the columns of L'. The rows stay in the same
-        # order.
-        return I, p[J], V
-    catch err
-        if err isa MOI.AddConstraintNotAllowed
-            rethrow(err)
-        end
-        msg = "There was an error computing a matrix square root"
-        throw(MOI.UnsupportedConstraint{typeof(func),typeof(set)}(msg))
+function _error_msg(::MOI.Utilities.CliqueTreesExt)
+    return """
+    Unable to transform a quadratic constraint into a SecondOrderCone
+    constraint because the quadratic constraint is not strongly convex and
+    our Cholesky decomposition failed.
+
+    If the constraint is convex but not strongly convex, you can work-around
+    this issue by manually installing and loading `CliqueTrees.jl`:
+    ```julia
+    import Pkg; Pkg.add("CliqueTrees")
+    using CliqueTrees
+    ```
+
+    CliqueTrees.jl is not included by default because it contains a number of
+    heavy dependencies.
+    """
+end
+
+function compute_sparse_sqrt(Q, ::F, ::S) where {F,S}
+    if (ret = MOI.Utilities.compute_sparse_sqrt(Q)) !== nothing
+        return ret
+    elseif !MOI.Utilities.is_defined(MOI.Utilities.LDLFactorizationsExt())
+        msg = _error_msg(MOI.Utilities.LDLFactorizationsExt())
+        return throw(MOI.AddConstraintNotAllowed{F,S}(msg))
+    elseif !MOI.Utilities.is_defined(MOI.Utilities.CliqueTreesExt())
+        msg = _error_msg(MOI.Utilities.CliqueTreesExt())
+        return throw(MOI.AddConstraintNotAllowed{F,S}(msg))
     end
+    msg = """
+    Unable to transform a quadratic constraint into a SecondOrderCone
+    constraint because the quadratic constraint is not convex.
+    """
+    return throw(MOI.UnsupportedConstraint{F,S}(msg))
 end
 
 function bridge_constraint(

src/Utilities/Utilities.jl

@@ -70,5 +70,6 @@ include("lazy_iterators.jl")
 include("set_dot.jl")
 
 include("distance_to_set.jl")
+include("sparse_square_root.jl")
 
 end # module

src/Utilities/sparse_square_root.jl

@@ -0,0 +1,72 @@
+# Copyright (c) 2017: Miles Lubin and contributors
+# Copyright (c) 2017: Google Inc.
+#
+# 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.
+
+abstract type AbstractExt end
+
+is_defined(::AbstractExt) = false
+
+struct LDLFactorizationsExt <: AbstractExt end
+
+struct CliqueTreesExt <: AbstractExt end
+
+struct LinearAlgebraExt <: AbstractExt end
+
+is_defined(::LinearAlgebraExt) = true
+
+function compute_sparse_sqrt(::LinearAlgebraExt, Q::AbstractMatrix)
+    factor = LinearAlgebra.cholesky(Q; check = false)
+    if !LinearAlgebra.issuccess(factor)
+        return nothing
+    end
+    L, p = SparseArrays.sparse(factor.L), factor.p
+    # We have Q = P' * L * L' * P. We want to find Q = U' * U, so U = L' * P
+    # First, compute L'. Note I and J are reversed
+    J, I, V = SparseArrays.findnz(L)
+    # Then, we want to permute the columns of L'. The rows stay in the same
+    # order.
+    return I, p[J], V
+end
+
+"""
+    compute_sparse_sqrt(Q::AbstractMatrix)
+
+Attempts to compute a sparse square root such that `Q = A' * A`.
+
+## Return value
+
+If successful, this function returns an `(I, J, V)` triplet of the sparse `A`
+matrix.
+
+If unsuccessful, this function returns `nothing`.
+
+## Extensions
+
+By default, this function attempts to use a Cholesky decomposition. If that
+fails, it may optionally use various extension packages.
+
+These extension packages must be loaded before calling `compute_sparse_sqrt`.
+
+The extensions currently supported are:
+
+ * The LDL routine in `LDLFactorizations.jl`
+ * The pivoted Cholesky in `CliqueTrees.jl`
+"""
+function compute_sparse_sqrt(Q::AbstractMatrix)
+    # There's a big try-catch here because Cholesky can fail even if
+    # `check = false`. The try-catch isn't a performance concern because the
+    # alternative is not being able to reformulate the problem.
+    for ext in (LinearAlgebraExt(), LDLFactorizationsExt(), CliqueTreesExt())
+        if is_defined(ext)
+            try
+                if (ret = compute_sparse_sqrt(ext, Q)) !== nothing
+                    return ret
+                end
+            catch
+            end
+        end
+    end
+    return nothing
+end