Add an extension for CliqueTrees.jl

Project.toml

@@ -19,12 +19,15 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 
 [extensions]
 MathOptInterfaceBenchmarkToolsExt = "BenchmarkTools"
+MathOptInterfaceCliqueTreesExt = "CliqueTrees"
 
 [compat]
 BenchmarkTools = "1"
+CliqueTrees = "1.17"
 CodecBzip2 = "0.6, 0.7, 0.8"
 CodecZlib = "0.6, 0.7"
 ForwardDiff = "1"
@@ -33,7 +36,7 @@ JSONSchema = "1"
 LinearAlgebra = "1"
 MutableArithmetics = "1"
 NaNMath = "0.3, 1"
-OrderedCollections = "1"
+OrderedCollections = "1.1"
 ParallelTestRunner = "2.4.1"
 PrecompileTools = "1"
 Printf = "1"
@@ -44,8 +47,9 @@ julia = "1.10"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 JSONSchema = "7d188eb4-7ad8-530c-ae41-71a32a6d4692"
 ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 
 [targets]
-test = ["BenchmarkTools", "JSONSchema", "ParallelTestRunner"]
+test = ["BenchmarkTools", "CliqueTrees", "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.Bridges.Constraint.is_defined(::MOI.Bridges.Constraint._CliqueTrees) = true
+
+function MOI.Bridges.Constraint._compute_sparse_sqrt(
+    ::MOI.Bridges.Constraint._CliqueTrees,
+    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

src/Bridges/Constraint/bridges/QuadtoSOCBridge.jl

@@ -60,31 +60,95 @@ end
 const QuadtoSOC{T,OT<:MOI.ModelLike} =
     SingleBridgeOptimizer{QuadtoSOCBridge{T},OT}
 
-function compute_sparse_sqrt(Q, func, set)
+abstract type _AbstractExt end
+
+is_defined(::_AbstractExt) = false
+
+struct _CliqueTrees <: _AbstractExt end
+
+struct _LinearAlgebra <: _AbstractExt end
+
+is_defined(::_LinearAlgebra) = true
+
+function _compute_sparse_sqrt(::_LinearAlgebra, 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 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`. 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))
+    # `check = false`. The try-catch isn't a performance concern because the
+    # alternative is not being able to reformulate the problem.
+    for ext in (_LinearAlgebra(), _CliqueTrees())
+        if is_defined(ext)
+            try
+                if (ret = _compute_sparse_sqrt(ext, Q)) !== nothing
+                    return ret
+                end
+            catch
+            end
+        end
+    end
+    return nothing
+end
+
+function _get_sqrt_error_message(is_clique_trees_defined::Bool)
+    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.
+    """
+    if is_clique_trees_defined
+        return msg
     end
+    clique_trees = """
+
+    ## CliqueTrees.jl
+
+    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.
+    """
+    return msg * clique_trees
 end
 
 function bridge_constraint(
@@ -111,8 +175,12 @@ 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 = _compute_sparse_sqrt(LinearAlgebra.Symmetric(Q))
+    if sqrt_ret === nothing
+        msg = _get_sqrt_error_message(is_defined(_CliqueTrees()))
+        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(