Merge pull request #886 from chriselrod/checkchol

Added check option to cholesky

Author: andreasnoack

src/cholesky.jl

@@ -1,20 +1,29 @@
 # Generic Cholesky decomposition for fixed-size matrices, mostly unrolled
 non_hermitian_error() = throw(LinearAlgebra.PosDefException(-1))
-@inline function LinearAlgebra.cholesky(A::StaticMatrix)
+@inline function LinearAlgebra.cholesky(A::StaticMatrix; check::Bool = true)
     ishermitian(A) || non_hermitian_error()
-    C = _cholesky(Size(A), A)
-    return Cholesky(C, 'U', 0)
+    _cholesky(Size(A), A, check)
 end
 
-@inline function LinearAlgebra.cholesky(A::LinearAlgebra.RealHermSymComplexHerm{<:Real, <:StaticMatrix})
-    C = _cholesky(Size(A), A.data)
-    return Cholesky(C, 'U', 0)
+@inline function LinearAlgebra.cholesky(A::LinearAlgebra.RealHermSymComplexHerm{<:Real, <:StaticMatrix}; check::Bool = true)
+    _cholesky(Size(A), A.data, check)
 end
 @inline LinearAlgebra._chol!(A::StaticMatrix, ::Type{UpperTriangular}) = (cholesky(A).U, 0)
 
-@generated function _cholesky(::Size{S}, A::StaticMatrix{M,M}) where {S,M}
+@inline function _chol_failure(A, info, check)
+    if check
+        throw(LinearAlgebra.PosDefException(info))
+    else
+        Cholesky(A, 'U', info)
+    end
+end
+# x < zero(x) is check used in `sqrt`, letting LLVM eliminate that check and remove error code.
+@inline _nonpdcheck(x::Real) = x ≥ zero(x)
+@inline _nonpdcheck(x) = x == x
+
+@generated function _cholesky(::Size{S}, A::StaticMatrix{M,M}, check::Bool) where {S,M}
     @assert (M,M) == S
-    M > 24 && return :(_cholesky_large(Size{$S}(), A))
+    M > 24 && return :(_cholesky_large(Size{$S}(), A, check))
     q = Expr(:block)
     for n ∈ 1:M
         for m ∈ n:M
@@ -28,7 +37,9 @@ end
             push!(q.args, :($L_m_n = muladd(-$L_m_k, $L_n_k', $L_m_n)))
         end
         L_n_n = Symbol(:L_,n,:_,n)
-        push!(q.args, :($L_n_n = sqrt($L_n_n)))
+        L_n_n_ltz = Symbol(:L_,n,:_,n,:_,:ltz)
+        push!(q.args, :(_nonpdcheck($L_n_n) || return _chol_failure(A, $n, check)))
+        push!(q.args, :($L_n_n = Base.FastMath.sqrt_fast($L_n_n)))
         Linv_n_n = Symbol(:Linv_,n,:_,n)
         push!(q.args, :($Linv_n_n = inv($L_n_n)))
         for m ∈ n+1:M
@@ -46,13 +57,15 @@ end
             push!(ret.args, :(zero(T)))
         end
     end
-    push!(q.args, :(similar_type(A, T)($ret)))
+    push!(q.args, :(Cholesky(similar_type(A, T)($ret), 'U', 0)))
     return Expr(:block, Expr(:meta, :inline), q)
 end
 
 # Otherwise default algorithm returning wrapped SizedArray
-@inline _cholesky_large(::Size{S}, A::StaticArray) where {S} =
-    similar_type(A)(cholesky(Hermitian(Matrix(A))).U)
+@inline function _cholesky_large(::Size{S}, A::StaticArray, check::Bool) where {S}
+    C = cholesky(Hermitian(Matrix(A)); check=check)
+    Cholesky(similar_type(A)(C.U), 'U', C.info)
+end
 
 LinearAlgebra.hermitian_type(::Type{SA}) where {T, S, SA<:SArray{S,T}} = Hermitian{T,SA}
 

test/chol.jl

@@ -34,6 +34,7 @@ using LinearAlgebra: PosDefException
             m_a = randn(elty, 4,4)
             #non hermitian
             @test_throws PosDefException cholesky(SMatrix{4,4}(m_a))
+            
             m_a = m_a*m_a'
             m = SMatrix{4,4}(m_a)
             @test cholesky(m).L ≈ cholesky(m_a).L
@@ -86,6 +87,21 @@ using LinearAlgebra: PosDefException
             @test (@inferred c \ v) isa SVector{3,elty}
             @test c \ v ≈ c_a \ v_a
         end
+
+        @testset "Check" begin
+            for i ∈ [1,3,7,25]
+                x = SVector(ntuple(elty, i))
+                nonpd = x * x'
+                if i > 1
+                    @test_throws PosDefException cholesky(nonpd)
+                    @test !issuccess(cholesky(nonpd,check=false))
+                    @test_throws PosDefException cholesky(Hermitian(nonpd))
+                    @test !issuccess(cholesky(Hermitian(nonpd),check=false))
+                else
+                    @test issuccess(cholesky(Hermitian(nonpd),check=false))
+                end
+            end
+        end
     end
 
     @testset "static blockmatrix" for i = 1:10