Add an extension for CliqueTrees.jl and LDLFactorizations.jl

Project.toml

@@ -18,17 +18,27 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 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]
+MathOptInterfaceCliqueTreesExt = "CliqueTrees"
+MathOptInterfaceLDLFactorizationsExt = "LDLFactorizations"
+
 [compat]
 BenchmarkTools = "1"
+CliqueTrees = "1.17"
 CodecBzip2 = "0.6, 0.7, 0.8"
 CodecZlib = "0.6, 0.7"
 ForwardDiff = "1"
 JSON = "0.21, 1"
 JSONSchema = "1"
+LDLFactorizations = "0.10"
 LinearAlgebra = "1"
 MutableArithmetics = "1"
 NaNMath = "0.3, 1"
-OrderedCollections = "1"
+OrderedCollections = "1.1"
 ParallelTestRunner = "2.4.1"
 PrecompileTools = "1"
 Printf = "1"
@@ -38,8 +48,10 @@ Test = "1"
 julia = "1.10"
 
 [extras]
+CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 JSONSchema = "7d188eb4-7ad8-530c-ae41-71a32a6d4692"
+LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
 ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 
 [targets]
-test = ["JSONSchema", "ParallelTestRunner"]
+test = ["CliqueTrees", "LDLFactorizations", "JSONSchema", "ParallelTestRunner"]

ext/MathOptInterfaceCliqueTreesExt.jl

@@ -0,0 +1,33 @@
+# 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.
+
+module MathOptInterfaceCliqueTreesExt
+
+import CliqueTrees
+import LinearAlgebra
+import MathOptInterface as MOI
+import SparseArrays
+
+MOI.Utilities.is_defined(::MOI.Utilities.CliqueTreesExt) = true
+
+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)
+        return nothing
+    end
+    return SparseArrays.findnz(U)
+end
+
+end  # module

ext/MathOptInterfaceLDLFactorizationsExt.jl

@@ -0,0 +1,30 @@
+# 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.
+
+module MathOptInterfaceLDLFactorizationsExt
+
+import LDLFactorizations
+import LinearAlgebra
+import MathOptInterface as MOI
+import SparseArrays
+
+MOI.Utilities.is_defined(::MOI.Utilities.LDLFactorizationsExt) = true
+
+function MOI.Utilities.compute_sparse_sqrt(
+    ::MOI.Utilities.LDLFactorizationsExt,
+    Q::AbstractMatrix,
+)
+    n = LinearAlgebra.checksquare(Q)
+    factor = LDLFactorizations.ldl(Q)
+    if !LDLFactorizations.factorized(factor) || minimum(factor.D) < 0
+        return nothing
+    end
+    L = sqrt.(factor.D) * LinearAlgebra.UnitLowerTriangular(factor.L)
+    J, I, V = SparseArrays.findnz(SparseArrays.sparse(L))
+    return I, factor.P[J], V
+end
+
+end  # module

src/Bridges/Constraint/bridges/QuadtoSOCBridge.jl

@@ -60,33 +60,6 @@ end
 const QuadtoSOC{T,OT<:MOI.ModelLike} =
     SingleBridgeOptimizer{QuadtoSOCBridge{T},OT}
 
-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)
-        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
-        msg = """
-        Unable to transform a quadratic constraint into a SecondOrderCone
-        constraint because the quadratic constraint is not strongly convex and
-        our Cholesky decomposition failed.
-        """
-        throw(MOI.UnsupportedConstraint{typeof(func),typeof(set)}(msg))
-    end
-end
-
 function bridge_constraint(
     ::Type{QuadtoSOCBridge{T}},
     model,
@@ -111,8 +84,46 @@ function bridge_constraint(
             MOI.ScalarAffineTerm(scale * term.coefficient, term.variable),
         ) for term in func.affine_terms
     ]
-    I, J, V = compute_sparse_sqrt(LinearAlgebra.Symmetric(Q), func, set)
-    for (i, j, v) in zip(I, J, V)
+    sqrt_ret = MOI.Utilities.compute_sparse_sqrt(LinearAlgebra.Symmetric(Q))
+    if sqrt_ret === nothing
+        msg = """
+        ## SecondOrderCone reformulation
+
+        We tried to reformulate the quadratic constraint into a SecondOrderCone,
+        but this failed because the quadratic constraint is not strongly convex
+        and our matrix factorization failed.
+
+        ## Package extensions
+
+        If the constraint is convex but not strongly convex, you can work-around
+        this issue by manually installing and loading one of the following
+        packages.
+
+        ### LDLFactorizations.jl
+
+        Currently active: $(MOI.Utilities.is_defined(MOI.Utilities.LDLFactorizationsExt()))
+
+        ```julia
+        import Pkg; Pkg.add("LDLFactorizations")
+        using LDLFactorizations
+        ```
+        LDLFactorizations.jl is not included by default because it is licensed
+        under the LGPL.
+
+        ### CliqueTrees.jl
+
+        Currently active: $(MOI.Utilities.is_defined(MOI.Utilities.CliqueTreesExt()))
+
+        ```julia
+        import Pkg; Pkg.add("CliqueTrees")
+        using CliqueTrees
+        ```
+        CliqueTrees.jl is not included by default because it contains a number of
+        heavy dependencies.
+        """
+        return throw(MOI.UnsupportedConstraint{typeof(func),typeof(set)}(msg))
+    end
+    for (i, j, v) in zip(sqrt_ret[1], sqrt_ret[2], sqrt_ret[3])
         push!(
             vector_terms,
             MOI.VectorAffineTerm(

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