Define functions for Cholesky

Project.toml

@@ -1,6 +1,6 @@
 name = "PDMats"
 uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
-version = "0.11.16"
+version = "0.11.17"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

README.md

@@ -217,6 +217,17 @@ While in theory all of them can be defined, at present only the following subset
 
 PRs to implement more generic fallbacks are welcome.
 
+### Fallbacks for `LinearAlgebra.Cholesky`
+
+For Cholesky decompositions of type `Cholesky` the following functions are defined as well:
+
+ - `dim`
+ - `whiten`, `whiten!`
+ - `unwhiten`, `unwhiten!`
+ - `quad`, `quad!`
+ - `invquad`, `invquad!`
+ - `X_A_Xt`, `Xt_A_X`, `X_invA_Xt`, `Xt_invA_X`
+
 ## Define Customized Subtypes
 
 In some situation, it is useful to define a customized subtype of `AbstractPDMat` to capture positive definite matrices with special structures. For this purpose, one has to define a subset of methods (as listed below), and other methods will be automatically provided.

src/chol.jl

@@ -21,3 +21,67 @@ if HAVE_CHOLMOD
 
     chol_lower(cf::CholTypeSparse) = cf.L
 end
+
+# Interface for `Cholesky`
+
+dim(A::Cholesky) = LinearAlgebra.checksquare(A)
+
+# whiten
+whiten(A::Cholesky, x::AbstractVecOrMat) = chol_lower(A) \ x
+whiten!(A::Cholesky, x::AbstractVecOrMat) = ldiv!(chol_lower(A), x)
+
+# unwhiten
+unwhiten(A::Cholesky, x::AbstractVecOrMat) = chol_lower(A) * x
+unwhiten!(A::Cholesky, x::AbstractVecOrMat) = lmul!(chol_lower(A), x)
+
+# 3-argument whiten/unwhiten
+for T in (:AbstractVector, :AbstractMatrix)
+    @eval begin
+        whiten!(r::$T, A::Cholesky, x::$T) = whiten!(A, copyto!(r, x))
+        unwhiten!(r::$T, A::Cholesky, x::$T) = unwhiten!(A, copyto!(r, x))
+    end
+end
+
+# quad
+quad(A::Cholesky, x::AbstractVector) = sum(abs2, chol_upper(A) * x)
+function quad(A::Cholesky, X::AbstractMatrix)
+    Z = chol_upper(A) * X
+    return vec(sum(abs2, Z; dims=1))
+end
+function quad!(r::AbstractArray, A::Cholesky, X::AbstractMatrix)
+    Z = chol_upper(A) * X
+    return map!(Base.Fix1(sum, abs2), r, _eachcol(Z))
+end
+
+# invquad
+invquad(A::Cholesky, x::AbstractVector) = sum(abs2, chol_lower(A) \ x)
+function invquad(A::Cholesky, X::AbstractMatrix)
+    Z = chol_lower(A) \ X
+    return vec(sum(abs2, Z; dims=1))
+end
+function invquad!(r::AbstractArray, A::Cholesky, X::AbstractMatrix)
+    Z = chol_lower(A) * X
+    return map!(Base.Fix1(sum, abs2), r, _eachcol(Z))
+end
+
+# tri products
+
+function X_A_Xt(A::Cholesky, X::AbstractMatrix)
+    Z = X * chol_lower(A)
+    return Z * transpose(Z)
+end
+
+function Xt_A_X(A::Cholesky, X::AbstractMatrix)
+    Z = chol_upper(A) * X
+    return transpose(Z) * Z
+end
+
+function X_invA_Xt(A::Cholesky, X::AbstractMatrix)
+    Z = X / chol_upper(A)
+    return Z * transpose(Z)
+end
+
+function Xt_invA_X(A::Cholesky, X::AbstractMatrix)
+    Z = chol_lower(A) \ X
+    return transpose(Z) * Z
+end

src/utils.jl

@@ -124,3 +124,10 @@ end
 else
     _ldiv!(Y::AbstractArray, s::Number, X::AbstractArray) = ldiv!(Y, s, X)
 end
+
+# eachcol was introduced in Julia 1.1
+@static if VERSION < v"1.1.0-DEV.792"
+    _eachcol(x::AbstractMatrix) = (view(x, :, i) for i in axes(x, 2))
+else
+    const _eachcol = eachcol
+end

test/chol.jl

@@ -2,19 +2,28 @@ using LinearAlgebra, PDMats
 using PDMats: chol_lower, chol_upper
 
 @testset "chol_lower" begin
-    A = rand(100, 100)
+    d = 100
+    A = rand(d, d)
     C = A'A
+    invC = inv(C)
     size_of_one_copy = sizeof(C)
-    @assert size_of_one_copy > 100  # ensure the matrix is large enough that few-byte allocations don't matter
+    @assert size_of_one_copy > d  # ensure the matrix is large enough that few-byte allocations don't matter
 
     chol_lower(C)
     @test (@allocated chol_lower(C)) < 1.05 * size_of_one_copy  # allow 5% overhead
 
+    X = randn(d, 10)
     for uplo in (:L, :U)
         ch = cholesky(Symmetric(C, uplo))
         chol_lower(ch)
         @test (@allocated chol_lower(ch)) < 33  # allow small overhead for wrapper types
         chol_upper(ch)
         @test (@allocated chol_upper(ch)) < 33  # allow small overhead for wrapper types
+
+        # Only test dim, `quad`/`invquad`, `whiten`/`unwhiten`, and tri products
+        @test dim(ch) == size(C, 1)
+        pdtest_quad(ch, C, invC, X, 0)
+        pdtest_triprod(ch, C, invC, X, 0)
+        pdtest_whiten(ch, C, 0)
     end
 end

test/specialarrays.jl

@@ -24,12 +24,15 @@ using StaticArrays
         X = @SMatrix rand(10, 4)
         Y = @SMatrix rand(4, 10)
 
-        for A in (PDS, D)
-            @test A * x isa SVector{4, Float64}
-            @test A * x ≈ Matrix(A) * Vector(x)
+        for A in (PDS, D, C)
+            if !(A isa Cholesky)
+                # `*(::Cholesky, ::SArray)` is not defined
+                @test A * x isa SVector{4, Float64}
+                @test A * x ≈ Matrix(A) * Vector(x)
 
-            @test A * Y isa SMatrix{4, 10, Float64}
-            @test A * Y ≈ Matrix(A) * Matrix(Y)
+                @test A * Y isa SMatrix{4, 10, Float64}
+                @test A * Y ≈ Matrix(A) * Matrix(Y)
+            end
 
             @test X / A isa SMatrix{10, 4, Float64}
             @test X / A ≈ Matrix(X) / Matrix(A)
@@ -52,6 +55,33 @@ using StaticArrays
             @test Xt_invA_X(A, Y) isa SMatrix{10, 10, Float64}
             @test Xt_invA_X(A, Y) ≈ Matrix(Y)' * (Matrix(A) \ Matrix(Y))
         end
+
+        # TODO: Fix for `PDMat` and `PDiagMat`
+        for A in (C,)
+            @test quad(A, x) isa Float64
+            @test quad(A, x) ≈ quad(Matrix(A), Vector(x))
+
+            @test quad(A, Y) isa SVector{10, Float64}
+            @test quad(A, Y) ≈ quad(Matrix(A), Matrix(Y))
+
+            @test invquad(A, x) isa Float64
+            @test invquad(A, x) ≈ invquad(Matrix(A), Vector(x))
+
+            @test invquad(A, Y) isa SVector{10, Float64}
+            @test invquad(A, Y) ≈ invquad(Matrix(A), Matrix(Y))
+
+            @test whiten(A, x) isa SVector{4, Float64}
+            @test whiten(A, x) ≈ whiten(Matrix(A), Vector(x))
+
+            @test whiten(A, Y) isa SMatrix{4, 10, Float64}
+            @test whiten(A, Y) ≈ whiten(Matrix(A), Matrix(Y))
+
+            @test unwhiten(A, x) isa SVector{4, Float64}
+            @test unwhiten(A, x) ≈ unwhiten(Matrix(A), Vector(x))
+
+            @test unwhiten(A, Y) isa SMatrix{4, 10, Float64}
+            @test unwhiten(A, Y) ≈ unwhiten(Matrix(A), Matrix(Y))
+        end
     end
 
     @testset "BandedMatrices" begin

test/testutils.jl

@@ -17,7 +17,7 @@ function test_pdmat(C, Cmat::Matrix;
                     t_cholesky::Bool=true,      # whether to test cholesky method
                     t_scale::Bool=true,         # whether to test scaling
                     t_add::Bool=true,           # whether to test pdadd
-		    t_det::Bool=true,           # whether to test det method
+                    t_det::Bool=true,           # whether to test det method
                     t_logdet::Bool=true,        # whether to test logdet method
                     t_eig::Bool=true,           # whether to test eigmax and eigmin
                     t_mul::Bool=true,           # whether to test multiplication