Support (un)whiten and (inv)quad with static arrays

src/generics.jl

@@ -33,19 +33,6 @@ LinearAlgebra.checksquare(a::AbstractPDMat) = size(a, 1)
 
 ## whiten and unwhiten
 
-whiten!(a::AbstractMatrix, x::AbstractVecOrMat) = whiten!(x, a, x)
-unwhiten!(a::AbstractMatrix, x::AbstractVecOrMat) = unwhiten!(x, a, x)
-
-function whiten!(r::AbstractVecOrMat, a::AbstractMatrix, x::AbstractVecOrMat)
-    v = _rcopy!(r, x)
-    ldiv!(chol_lower(cholesky(a)), v)
-end
-
-function unwhiten!(r::AbstractVecOrMat, a::AbstractMatrix, x::AbstractVecOrMat)
-    v = _rcopy!(r, x)
-    lmul!(chol_lower(cholesky(a)), v)
-end
-
 """
     whiten(a::AbstractMatrix, x::AbstractVecOrMat)
     unwhiten(a::AbstractMatrix, x::AbstractVecOrMat)
@@ -80,35 +67,41 @@ julia> W * W'
  0.0  1.0
 ```
 """
-whiten(a::AbstractMatrix, x::AbstractVecOrMat) = whiten!(similar(x), a, x)
-unwhiten(a::AbstractMatrix, x::AbstractVecOrMat) = unwhiten!(similar(x), a, x)
+whiten(a::AbstractMatrix{<:Real}, x::AbstractVecOrMat) = whiten(AbstractPDMat(a), x)
+unwhiten(a::AbstractMatrix{<:Real}, x::AbstractVecOrMat) = unwhiten(AbstractPDMat(a), x)
 
+whiten!(a::AbstractMatrix{<:Real}, x::AbstractVecOrMat) = whiten!(x, a, x)
+unwhiten!(a::AbstractMatrix{<:Real}, x::AbstractVecOrMat) = unwhiten!(x, a, x)
+
+function whiten!(r::AbstractVecOrMat, a::AbstractMatrix{<:Real}, x::AbstractVecOrMat)
+    return whiten!(r, AbstractPDMat(a), x)
+end
+function unwhiten!(r::AbstractVecOrMat, a::AbstractMatrix{<:Real}, x::AbstractVecOrMat)
+    return unwhiten!(r, AbstractPDMat(a), x)
+end
 
 ## quad
 
 """
     quad(a::AbstractMatrix, x::AbstractVecOrMat)
 
-Return the value of the quadratic form defined by `a` applied to `x`
+Return the value of the quadratic form defined by `a` applied to `x`.
 
 If `x` is a vector the quadratic form is `x' * a * x`.  If `x` is a matrix
 the quadratic form is applied column-wise.
 """
-function quad(a::AbstractMatrix{T}, x::AbstractMatrix{S}) where {T<:Real, S<:Real}
-    @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
-    quad!(Array{promote_type(T, S)}(undef, size(x,2)), a, x)
+function quad(a::AbstractMatrix{<:Real}, x::AbstractVecOrMat)
+    return quad(AbstractPDMat(a), x)
 end
 
-quad(a::AbstractMatrix, x::AbstractVector) = sum(abs2, chol_upper(cholesky(a)) * x)
-invquad(a::AbstractMatrix, x::AbstractVector) = sum(abs2, chol_lower(cholesky(a)) \ x)
-
 """
     quad!(r::AbstractArray, a::AbstractMatrix, x::AbstractMatrix)
 
-Overwrite `r` with the value of the quadratic form defined by `a` applied columnwise to `x`
+Overwrite `r` with the value of the quadratic form defined by `a` applied columnwise to `x`.
 """
-quad!(r::AbstractArray, a::AbstractMatrix, x::AbstractMatrix) = colwise_dot!(r, x, a * x)
-
+function quad!(r::AbstractArray, a::AbstractMatrix{<:Real}, x::AbstractMatrix)
+    return quad!(r, AbstractPDMat(a), x)
+end
 
 """
     invquad(a::AbstractMatrix, x::AbstractVecOrMat)
@@ -120,15 +113,16 @@ For most `PDMat` types this is done in a way that does not require evaluation of
 If `x` is a vector the quadratic form is `x' * a * x`.  If `x` is a matrix
 the quadratic form is applied column-wise.
 """
-invquad(a::AbstractMatrix, x::AbstractVecOrMat) = x' / a * x
-function invquad(a::AbstractMatrix{T}, x::AbstractMatrix{S}) where {T<:Real, S<:Real}
-    @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
-    invquad!(Array{promote_type(T, S)}(undef, size(x,2)), a, x)
+function invquad(a::AbstractMatrix{<:Real}, x::AbstractVecOrMat)
+    return invquad(AbstractPDMat(a), x)
 end
 
 """
     invquad!(r::AbstractArray, a::AbstractMatrix, x::AbstractMatrix)
 
 Overwrite `r` with the value of the quadratic form defined by `inv(a)` applied columnwise to `x`
 """
-invquad!(r::AbstractArray, a::AbstractMatrix, x::AbstractMatrix) = colwise_dot!(r, x, a \ x)
+function invquad!(r::AbstractArray, a::AbstractMatrix{<:Real}, x::AbstractMatrix)
+    return invquad!(r, AbstractPDMat(a), x)
+end
+

src/pdiagmat.jl

@@ -91,45 +91,38 @@ LinearAlgebra.sqrt(a::PDiagMat) = PDiagMat(map(sqrt, a.diag))
 
 ### whiten and unwhiten
 
-function whiten!(r::StridedVector, a::PDiagMat, x::StridedVector)
-    n = a.dim
-    @check_argdims length(r) == length(x) == n
-    v = a.diag
-    for i = 1:n
-        r[i] = x[i] / sqrt(v[i])
-    end
-    return r
+function whiten!(r::AbstractVecOrMat, a::PDiagMat, x::AbstractVecOrMat)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
+    return r .= x ./ sqrt.(a.diag)
 end
-
-function unwhiten!(r::StridedVector, a::PDiagMat, x::StridedVector)
-    n = a.dim
-    @check_argdims length(r) == length(x) == n
-    v = a.diag
-    for i = 1:n
-        r[i] = x[i] * sqrt(v[i])
-    end
-    return r
+function unwhiten!(r::AbstractVecOrMat, a::PDiagMat, x::AbstractVecOrMat)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
+    return r .= x .* sqrt.(a.diag)
 end
 
-function whiten!(r::StridedMatrix, a::PDiagMat, x::StridedMatrix)
-    r .= x ./ sqrt.(a.diag)
-    return r
+function whiten(a::PDiagMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return x ./ sqrt.(a.diag)
 end
-
-function unwhiten!(r::StridedMatrix, a::PDiagMat, x::StridedMatrix)
-    r .= x .* sqrt.(a.diag)
-    return r
+function unwhiten(a::PDiagMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return x .* sqrt.(a.diag)
 end
 
-
-whiten!(r::AbstractVecOrMat, a::PDiagMat, x::AbstractVecOrMat) = r .= x ./ sqrt.(a.diag)
-unwhiten!(r::AbstractVecOrMat, a::PDiagMat, x::AbstractVecOrMat) = r .= x .* sqrt.(a.diag)
-
-
 ### quadratic forms
 
-quad(a::PDiagMat, x::AbstractVector) = wsumsq(a.diag, x)
-invquad(a::PDiagMat, x::AbstractVector) = invwsumsq(a.diag, x)
+function quad(a::PDiagMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    if x isa AbstractVector
+        return wsumsq(a.diag, x)
+    else
+        # map(Base.Fix1(invquad, a), eachcol(x)) or similar alternatives
+        # do NOT return a `SVector` for inputs `x::SMatrix`.
+        return vec(sum(abs2.(x) .* a.diag; dims = 1))
+    end
+end
 
 function quad!(r::AbstractArray, a::PDiagMat, x::AbstractMatrix)
     ad = a.diag
@@ -145,8 +138,18 @@ function quad!(r::AbstractArray, a::PDiagMat, x::AbstractMatrix)
     r
 end
 
+function invquad(a::PDiagMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    if x isa AbstractVector
+        return invwsumsq(a.diag, x)
+    else
+        # map(Base.Fix1(invquad, a), eachcol(x)) or similar alternatives
+        # do NOT return a `SVector` for inputs `x::SMatrix`.
+        return vec(sum(abs2.(x) ./ a.diag; dims = 1))
+    end
+end
+
 function invquad!(r::AbstractArray, a::PDiagMat, x::AbstractMatrix)
-    m, n = size(x)
     ad = a.diag
     @check_argdims eachindex(ad) == axes(x, 1)
     @check_argdims eachindex(r) == axes(x, 2)
@@ -186,3 +189,18 @@ function Xt_invA_X(a::PDiagMat, x::AbstractMatrix)
     z = x ./ sqrt.(a.diag)
     transpose(z) * z
 end
+
+### Specializations for `Array` arguments with reduced allocations
+
+function quad(a::PDiagMat{<:Real,<:Vector}, x::Matrix)
+    @check_argdims a.dim == size(x, 1)
+    T = typeof(zero(eltype(a)) * abs2(zero(eltype(x))))
+    return quad!(Vector{T}(undef, size(x, 2)), a, x)
+end
+
+function invquad(a::PDiagMat{<:Real,<:Vector}, x::Matrix)
+    @check_argdims a.dim == size(x, 1)
+    T = typeof(abs2(zero(eltype(x))) / zero(eltype(a)))
+    return invquad!(Vector{T}(undef, size(x, 2)), a, x)
+end
+

src/pdmat.jl

@@ -86,6 +86,78 @@ LinearAlgebra.eigmin(a::PDMat) = eigmin(a.mat)
 Base.kron(A::PDMat, B::PDMat) = PDMat(kron(A.mat, B.mat), Cholesky(kron(A.chol.U, B.chol.U), 'U', A.chol.info))
 LinearAlgebra.sqrt(A::PDMat) = PDMat(sqrt(Hermitian(A.mat)))
 
+### (un)whitening
+
+function whiten!(r::AbstractVecOrMat, a::PDMat, x::AbstractVecOrMat)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
+    v = _rcopy!(r, x)
+    return ldiv!(chol_lower(cholesky(a)), v)
+end
+function unwhiten!(r::AbstractVecOrMat, a::PDMat, x::AbstractVecOrMat)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
+    v = _rcopy!(r, x)
+    return lmul!(chol_lower(cholesky(a)), v)
+end
+
+function whiten(a::PDMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return chol_lower(cholesky(a)) \ x
+end
+function unwhiten(a::PDMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return chol_lower(cholesky(a)) * x
+end
+
+## quad/invquad
+
+function quad(a::PDMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    aU_x = chol_upper(cholesky(a)) * x
+    if x isa AbstractVector
+        return sum(abs2, aU_x)
+    else
+        return vec(sum(abs2, aU_x; dims = 1))
+    end
+end
+
+function quad!(r::AbstractArray, a::PDMat, x::AbstractMatrix)
+    @check_argdims axes(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
+    @inbounds for i in axes(x, 2)
+        copyto!(z, view(x, :, i))
+        lmul!(aU, z)
+        r[i] = sum(abs2, z)
+    end
+    return r
+end
+
+function invquad(a::PDMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    inv_aL_x = chol_lower(cholesky(a)) \ x
+    if x isa AbstractVector
+        return sum(abs2, inv_aL_x)
+    else
+        return vec(sum(abs2, inv_aL_x; dims = 1))
+    end
+end
+
+function invquad!(r::AbstractArray, a::PDMat, x::AbstractMatrix)
+    @check_argdims axes(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
+    @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::PDMat, x::AbstractMatrix)
@@ -111,3 +183,18 @@ function Xt_invA_X(a::PDMat, x::AbstractMatrix)
     z = chol_lower(a.chol) \ x
     return transpose(z) * z
 end
+
+### Specializations for `Array` arguments with reduced allocations
+
+function quad(a::PDMat{<:Real,<:Vector}, x::Matrix)
+    @check_argdims a.dim == size(x, 1)
+    T = typeof(zero(eltype(a)) * abs2(zero(eltype(x))))
+    return quad!(Vector{T}(undef, size(x, 2)), a, x)
+end
+
+function invquad(a::PDMat{<:Real,<:Vector}, x::Matrix)
+    @check_argdims a.dim == size(x, 1)
+    T = typeof(abs2(zero(eltype(x))) / zero(eltype(a)))
+    return invquad!(Vector{T}(undef, size(x, 2)), a, x)
+end
+

src/pdsparsemat.jl

@@ -78,37 +78,84 @@ LinearAlgebra.sqrt(A::PDSparseMat) = PDMat(sqrt(Hermitian(Matrix(A))))
 ### whiten and unwhiten
 
 function whiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
     # Can't use `ldiv!` due to missing support in SparseArrays
     return copyto!(r, chol_lower(a.chol) \ x)
 end
 
 function unwhiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
     # `*` is not defined for `PtL` factor components,
     # so we can't use `chol_lower(a.chol) * x`
     C = a.chol
     PtL = sparse(C.L)[C.p, :]
-    # Can't use `lmul!` due to missing support in SparseArrays
     return copyto!(r, PtL * x)
 end
 
+function whiten(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return chol_lower(cholesky(a)) \ x
+end
+
+function unwhiten(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    # `*` is not defined for `PtL` factor components,
+    # so we can't use `chol_lower(a.chol) * x`
+    C = a.chol
+    PtL = sparse(C.L)[C.p, :]
+    return PtL * x
+end
 
 ### quadratic forms
 
-quad(a::PDSparseMat, x::AbstractVector) = dot(x, a * x)
-invquad(a::PDSparseMat, x::AbstractVector) = dot(x, a \ x)
+function quad(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    # https://github.com/JuliaLang/julia/commit/2425ae760fb5151c5c7dd0554e87c5fc9e24de73
+    if VERSION < v"1.4.0-DEV.92"
+        z = a.mat * x
+        return x isa AbstractVector ? dot(x, z) : map(dot, eachcol(x), eachcol(z))
+    else
+        return x isa AbstractVector ? dot(x, a.mat, x) : map(Base.Fix1(quad, a), eachcol(x))
+    end
+end
 
 function quad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
-    @check_argdims eachindex(r) == axes(x, 2)
-    for i in axes(x, 2)
-        r[i] = quad(a, x[:,i])
+    @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)
+            r[i] = dot(xi, a.mat, xi)
+        end
     end
     return r
 end
 
+function invquad(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    z = a.chol \ x
+    return x isa AbstractVector ? dot(x, z) : map(dot, eachcol(x), eachcol(z))
+end
+
 function invquad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
-    @check_argdims eachindex(r) == axes(x, 2)
-    for i in axes(x, 2)
-        r[i] = invquad(a, x[:,i])
+    @check_argdims axes(r) == axes(x, 2)
+    @check_argdims a.dim == size(x, 1)
+    z = similar(r, a.dim) # buffer to save allocations
+    @inbounds for i in axes(x, 2)
+        xi = view(x, :, i)
+        copyto!(z, xi)
+        ldiv!(a.chol, z)
+        r[i] = dot(xi, z)
     end
     return r
 end