Fix and test broken quad! and invquad! (#190)

* Fix and test broken `quad!` and `invquad!`

* Fix spurious test failures

* Better fixes for older Julia versions

* Address PR review

Author: devmotion

src/PDMats.jl

@@ -39,8 +39,8 @@ module PDMats
 
     # source files
 
-    include("chol.jl")
     include("utils.jl")
+    include("chol.jl")
 
     include("pdmat.jl")
 

src/chol.jl

@@ -46,45 +46,73 @@ for T in (:AbstractVector, :AbstractMatrix)
 end
 
 # quad
-quad(A::Cholesky, x::AbstractVector) = sum(abs2, chol_upper(A) * x)
+function quad(A::Cholesky, x::AbstractVector)
+    @check_argdims size(A, 1) == length(x)
+    return sum(abs2, chol_upper(A) * x)
+end
 function quad(A::Cholesky, X::AbstractMatrix)
+    @check_argdims size(A, 1) == size(X, 1)
     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))
+    @check_argdims eachindex(r) == axes(X, 2)
+    @check_argdims size(A, 1) == size(X, 1)
+    aU = chol_upper(A)
+    z = similar(r, size(A, 1)) # buffer to save allocations
+    @inbounds for i in axes(X, 2)
+        copyto!(z, view(X, :, i))
+        lmul!(aU, z)
+        r[i] = sum(abs2, z)
+    end
+    return r
 end
 
 # invquad
-invquad(A::Cholesky, x::AbstractVector) = sum(abs2, chol_lower(A) \ x)
-function invquad(A::Cholesky, X::AbstractMatrix)
+function invquad(A::Cholesky, x::AbstractVector)
+    @check_argdims size(A, 1) == size(x, 1)
+    return sum(abs2, chol_lower(A) \ x)
+end
+function invquad(A::Cholesky, X::AbstractMatrix) 
+    @check_argdims size(A, 1) == size(X, 1)
     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))
+    @check_argdims eachindex(r) == axes(X, 2)
+    @check_argdims size(A, 1) == size(X, 1)
+    aL = chol_lower(A)
+    z = similar(r, size(A, 1)) # buffer to save allocations
+    @inbounds for i in axes(X, 2)
+        copyto!(z, view(X, :, i))
+        ldiv!(aL, z)
+        r[i] = sum(abs2, z)
+    end
+    return r
 end
 
 # tri products
 
 function X_A_Xt(A::Cholesky, X::AbstractMatrix)
+    @check_argdims size(A, 1) == size(X, 2)
     Z = X * chol_lower(A)
     return Z * transpose(Z)
 end
 
 function Xt_A_X(A::Cholesky, X::AbstractMatrix)
+    @check_argdims size(A, 1) == size(X, 1)
     Z = chol_upper(A) * X
     return transpose(Z) * Z
 end
 
 function X_invA_Xt(A::Cholesky, X::AbstractMatrix)
+    @check_argdims size(A, 1) == size(X, 2)
     Z = X / chol_upper(A)
     return Z * transpose(Z)
 end
 
 function Xt_invA_X(A::Cholesky, X::AbstractMatrix)
+    @check_argdims size(A, 1) == size(X, 1)
     Z = chol_lower(A) \ X
     return transpose(Z) * Z
 end

src/pdmat.jl

@@ -123,7 +123,7 @@ function quad(a::PDMat, x::AbstractVecOrMat)
 end
 
 function quad!(r::AbstractArray, a::PDMat, x::AbstractMatrix)
-    @check_argdims axes(r) == axes(x, 2)
+    @check_argdims eachindex(r) == axes(x, 2)
     @check_argdims a.dim == size(x, 1)
     aU = chol_upper(cholesky(a))
     z = similar(r, a.dim) # buffer to save allocations
@@ -146,7 +146,7 @@ function invquad(a::PDMat, x::AbstractVecOrMat)
 end
 
 function invquad!(r::AbstractArray, a::PDMat, x::AbstractMatrix)
-    @check_argdims axes(r) == axes(x, 2)
+    @check_argdims eachindex(r) == axes(x, 2)
     @check_argdims a.dim == size(x, 1)
     aL = chol_lower(cholesky(a))
     z = similar(r, a.dim) # buffer to save allocations

src/pdsparsemat.jl

@@ -122,19 +122,14 @@ function quad(a::PDSparseMat, x::AbstractVecOrMat)
 end
 
 function quad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
-    @check_argdims axes(r) == axes(x, 2)
-    # https://github.com/JuliaLang/julia/commit/2425ae760fb5151c5c7dd0554e87c5fc9e24de73
-    if VERSION < v"1.4.0-DEV.92"
-       z = similar(r, a.dim) # buffer to save allocations
-        @inbounds for i in axes(x, 2)
-            xi = view(x, :, i)
-            copyto!(z, xi)
-            lmul!(a.mat, z)
-            r[i] = dot(xi, z)
-        end
-    else
-        @inbounds for i in axes(x, 2)
-            xi = view(x, :, i)
+    @check_argdims eachindex(r) == axes(x, 2)
+    @inbounds for i in axes(x, 2)
+        xi = view(x, :, i)
+        # https://github.com/JuliaLang/julia/commit/2425ae760fb5151c5c7dd0554e87c5fc9e24de73
+        if VERSION < v"1.4.0-DEV.92"
+            # Can't use `lmul!` with buffer due to missing support in SparseArrays
+            r[i] = dot(xi, a.mat * xi)
+        else
             r[i] = dot(xi, a.mat, xi)
         end
     end
@@ -148,14 +143,12 @@ function invquad(a::PDSparseMat, x::AbstractVecOrMat)
 end
 
 function invquad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
-    @check_argdims axes(r) == axes(x, 2)
+    @check_argdims eachindex(r) == axes(x, 2)
     @check_argdims a.dim == size(x, 1)
-    z = similar(r, a.dim) # buffer to save allocations
+    # Can't use `ldiv!` with buffer due to missing support in SparseArrays
     @inbounds for i in axes(x, 2)
         xi = view(x, :, i)
-        copyto!(z, xi)
-        ldiv!(a.chol, z)
-        r[i] = dot(xi, z)
+        r[i] = dot(xi, a.chol \ xi)
     end
     return r
 end

src/scalmat.jl

@@ -115,7 +115,10 @@ end
 function quad!(r::AbstractArray, a::ScalMat, x::AbstractMatrix)
     @check_argdims eachindex(r) == axes(x, 2)
     @check_argdims a.dim == size(x, 1)
-    return map!(Base.Fix1(quad, a), r, eachcol(x))
+    @inbounds for i in axes(x, 2)
+        r[i] = quad(a, view(x, :, i))
+    end
+    return r
 end
 
 function invquad(a::ScalMat, x::AbstractVecOrMat)
@@ -135,7 +138,10 @@ end
 function invquad!(r::AbstractArray, a::ScalMat, x::AbstractMatrix)
     @check_argdims eachindex(r) == axes(x, 2)
     @check_argdims a.dim == size(x, 1)
-    return map!(Base.Fix1(invquad, a), r, eachcol(x))
+    @inbounds for i in axes(x, 2)
+        r[i] = invquad(a, view(x, :, i))
+    end
+    return r
 end
 
 

test/specialarrays.jl

@@ -47,16 +47,16 @@ using StaticArrays
             @test A \ Y ≈ Matrix(A) \ Matrix(Y)
 
             @test whiten(A, x) isa SVector{4, Float64}
-            @test whiten(A, x) ≈ cholesky(Matrix(A)).L \ Vector(x)
+            @test whiten(A, x) ≈ cholesky(Symmetric(Matrix(A))).L \ Vector(x)
 
             @test whiten(A, Y) isa SMatrix{4, 10, Float64}
-            @test whiten(A, Y) ≈ cholesky(Matrix(A)).L \ Matrix(Y)
+            @test whiten(A, Y) ≈ cholesky(Symmetric(Matrix(A))).L \ Matrix(Y)
 
             @test unwhiten(A, x) isa SVector{4, Float64}
-            @test unwhiten(A, x) ≈ cholesky(Matrix(A)).L * Vector(x)
+            @test unwhiten(A, x) ≈ cholesky(Symmetric(Matrix(A))).L * Vector(x)
 
             @test unwhiten(A, Y) isa SMatrix{4, 10, Float64}
-            @test unwhiten(A, Y) ≈ cholesky(Matrix(A)).L * Matrix(Y)
+            @test unwhiten(A, Y) ≈ cholesky(Symmetric(Matrix(A))).L * Matrix(Y)
 
             @test quad(A, x) isa Float64
             @test quad(A, x) ≈ Vector(x)' * Matrix(A) * Vector(x)

test/testutils.jl

@@ -272,6 +272,9 @@ function pdtest_quad(C, Cmat::Matrix, Imat::Matrix, X::Matrix, verbose::Int)
         @test quad(C, view(X,:,i)) ≈ r_quad[i]
     end
     @test quad(C, X) ≈ r_quad
+    r = similar(r_quad)
+    @test quad!(r, C, X) === r
+    @test r ≈ r_quad
 
     _pdt(verbose, "invquad")
     r_invquad = zeros(eltype(C),n)
@@ -282,6 +285,9 @@ function pdtest_quad(C, Cmat::Matrix, Imat::Matrix, X::Matrix, verbose::Int)
         @test invquad(C, view(X,:,i)) ≈ r_invquad[i]
     end
     @test invquad(C, X) ≈ r_invquad
+    r = similar(r_invquad)
+    @test invquad!(r, C, X) === r
+    @test r ≈ r_invquad
 end