From af621146af86d5481a87c399b2cebc98247b362e Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Tue, 26 Sep 2023 09:38:50 +0200
Subject: [PATCH 1/7] Support StaticArrays in X(t)_(inv)A_X(t) with ScalMat

---
 Project.toml          | 2 +-
 src/scalmat.jl        | 8 ++++----
 test/specialarrays.jl | 5 ++++-
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index 3942ff7..ae30fd8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "PDMats"
 uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
-version = "0.11.19"
+version = "0.11.20"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
diff --git a/src/scalmat.jl b/src/scalmat.jl
index 6e8092e..e1c5038 100644
--- a/src/scalmat.jl
+++ b/src/scalmat.jl
@@ -99,20 +99,20 @@ invquad!(r::AbstractArray, a::ScalMat, x::AbstractMatrix) = colwise_sumsqinv!(r,
 
 function X_A_Xt(a::ScalMat, x::AbstractMatrix)
     @check_argdims LinearAlgebra.checksquare(a) == size(x, 2)
-    lmul!(a.value, x * transpose(x))
+    a.value * (x * transpose(x))
 end
 
 function Xt_A_X(a::ScalMat, x::AbstractMatrix)
     @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
-    lmul!(a.value, transpose(x) * x)
+    a.value * (transpose(x) * x)
 end
 
 function X_invA_Xt(a::ScalMat, x::AbstractMatrix)
     @check_argdims LinearAlgebra.checksquare(a) == size(x, 2)
-    _rdiv!(x * transpose(x), a.value)
+    (x * transpose(x)) / a.value
 end
 
 function Xt_invA_X(a::ScalMat, x::AbstractMatrix)
     @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
-    _rdiv!(transpose(x) * x, a.value)
+    (transpose(x) * x) / a.value
 end
diff --git a/test/specialarrays.jl b/test/specialarrays.jl
index 06621e8..b812a80 100644
--- a/test/specialarrays.jl
+++ b/test/specialarrays.jl
@@ -20,11 +20,14 @@ using StaticArrays
         @test D isa PDiagMat{Float64, <:SVector{4, Float64}}
         @test @inferred(kron(D, D)) isa PDiagMat{Float64, <:SVector{16, Float64}}
 
+        # Scaled identity matrix
+        E = ScalMat(4, 1.2)
+
         x = @SVector rand(4)
         X = @SMatrix rand(10, 4)
         Y = @SMatrix rand(4, 10)
 
-        for A in (PDS, D)
+        for A in (PDS, D, E)
             @test A * x isa SVector{4, Float64}
             @test A * x ≈ Matrix(A) * Vector(x)
 

From f9b4a375981771ac0f8a26fb1c5281314ce76089 Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Wed, 27 Sep 2023 16:24:08 +0200
Subject: [PATCH 2/7] Add specializations for `Matrix` with reduced allocations

---
 src/scalmat.jl | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/scalmat.jl b/src/scalmat.jl
index e1c5038..ea672ec 100644
--- a/src/scalmat.jl
+++ b/src/scalmat.jl
@@ -116,3 +116,24 @@ function Xt_invA_X(a::ScalMat, x::AbstractMatrix)
     @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
     (transpose(x) * x) / a.value
 end
+
+# Specializations for `x::Matrix` with reduced allocations
+function X_A_Xt(a::ScalMat, x::Matrix)
+    @check_argdims LinearAlgebra.checksquare(a) == size(x, 2)
+    lmul!(a.value, x * transpose(x))
+end
+
+function Xt_A_X(a::ScalMat, x::Matrix)
+    @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
+    lmul!(a.value, transpose(x) * x)
+end
+
+function X_invA_Xt(a::ScalMat, x::Matrix)
+    @check_argdims LinearAlgebra.checksquare(a) == size(x, 2)
+    _rdiv!(x * transpose(x), a.value)
+end
+
+function Xt_invA_X(a::ScalMat, x::Matrix)
+    @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
+    _rdiv!(transpose(x) * x, a.value)
+end

From 5b376f27997ff1b3faace1bfbbffd117d92a71b3 Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Tue, 26 Sep 2023 16:31:46 +0200
Subject: [PATCH 3/7] Support `(un)whiten` and `(inv)quad` with static arrays

---
 src/generics.jl       | 54 +++++++++++++++------------------
 src/pdiagmat.jl       | 65 +++++++++++++++++++++-------------------
 src/pdmat.jl          | 70 +++++++++++++++++++++++++++++++++++++++++++
 src/pdsparsemat.jl    | 33 ++++++++++++++++----
 src/scalmat.jl        | 54 +++++++++++++++++++++++++++++----
 test/specialarrays.jl | 24 +++++++++++++++
 6 files changed, 227 insertions(+), 73 deletions(-)

diff --git a/src/generics.jl b/src/generics.jl
index 686d190..78e95d2 100644
--- a/src/generics.jl
+++ b/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,10 +113,8 @@ 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
 
 """
@@ -131,4 +122,7 @@ end
 
 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
+
diff --git a/src/pdiagmat.jl b/src/pdiagmat.jl
index df6f74f..2147883 100644
--- a/src/pdiagmat.jl
+++ b/src/pdiagmat.jl
@@ -86,45 +86,48 @@ 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 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 quad!(r::AbstractArray, a::PDiagMat, x::AbstractMatrix)
     ad = a.diag
diff --git a/src/pdmat.jl b/src/pdmat.jl
index b8fd077..98a75ba 100644
--- a/src/pdmat.jl
+++ b/src/pdmat.jl
@@ -78,6 +78,76 @@ 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 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 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!(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)
diff --git a/src/pdsparsemat.jl b/src/pdsparsemat.jl
index d7f10ea..68e4121 100644
--- a/src/pdsparsemat.jl
+++ b/src/pdsparsemat.jl
@@ -70,20 +70,41 @@ LinearAlgebra.sqrt(A::PDSparseMat) = PDMat(sqrt(Hermitian(Matrix(A))))
 ### whiten and unwhiten
 
 function whiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
-    r[:] = sparse(chol_lower(a.chol)) \ x
-    return r
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
+    # ldiv! is not defined for SparseMatrixCSC
+    return copyto!(r, sparse(chol_lower(a.chol)) \ x)
 end
 
 function unwhiten!(r::AbstractVecOrMat, a::PDSparseMat, x::AbstractVecOrMat)
-    r[:] = sparse(chol_lower(a.chol)) * x
-    return r
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
+    # lmul! is not defined for SparseMatrixCSC
+    return copyto!(r, sparse(chol_lower(a.chol)) * x)
 end
 
+function whiten(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return sparse(chol_lower(cholesky(a))) \ x
+end
+
+function unwhiten(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return sparse(chol_lower(cholesky(a))) * 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)
+    z = sparse(chol_lower(cholesky(a)))' * x
+    return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 1))
+end
+function invquad(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    z = sparse(chol_lower(cholesky(a))) \ x
+    return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 1))
+end
 
 function quad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
     @check_argdims eachindex(r) == axes(x, 2)
diff --git a/src/scalmat.jl b/src/scalmat.jl
index ea672ec..cb15ca1 100644
--- a/src/scalmat.jl
+++ b/src/scalmat.jl
@@ -76,23 +76,65 @@ LinearAlgebra.sqrt(a::ScalMat) = ScalMat(a.dim, sqrt(a.value))
 ### whiten and unwhiten
 
 function whiten!(r::AbstractVecOrMat, a::ScalMat, x::AbstractVecOrMat)
-    @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
     _ldiv!(r, sqrt(a.value), x)
 end
 
 function unwhiten!(r::AbstractVecOrMat, a::ScalMat, x::AbstractVecOrMat)
-    @check_argdims LinearAlgebra.checksquare(a) == size(x, 1)
+    @check_argdims axes(r) == axes(x)
+    @check_argdims a.dim == size(x, 1)
     mul!(r, x, sqrt(a.value))
 end
 
+function whiten(a::ScalMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return x / sqrt(a.value)
+end
+function unwhiten(a::ScalMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    return sqrt(a.value) * x
+end
 
 ### quadratic forms
 
-quad(a::ScalMat, x::AbstractVector) = sum(abs2, x) * a.value
-invquad(a::ScalMat, x::AbstractVector) = sum(abs2, x) / a.value
+function quad(a::ScalMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    if x isa AbstractVector
+        return sum(abs2, x) * a.value
+    else
+        # map(Base.Fix1(quad, a), eachcol(x)) or similar alternatives
+        # do NOT return a `SVector` for inputs `x::SMatrix`.
+        wsq = let w = a.value
+            x -> w * abs2(x)
+        end 
+        return vec(sum(wsq, x; dims=1))
+    end
+end
+function invquad(a::ScalMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    if x isa AbstractVector
+        return sum(abs2, x) / a.value
+    else
+        # map(Base.Fix1(invquad, a), eachcol(x)) or similar alternatives
+        # do NOT return a `SVector` for inputs `x::SMatrix`.
+        wsq = let w = a.value
+            x -> abs2(x) / w
+        end 
+        return vec(sum(wsq, x; dims=1))
+    end
+end
 
-quad!(r::AbstractArray, a::ScalMat, x::AbstractMatrix) = colwise_sumsq!(r, x, a.value)
-invquad!(r::AbstractArray, a::ScalMat, x::AbstractMatrix) = colwise_sumsqinv!(r, x, a.value)
+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))
+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))
+end
 
 
 ### tri products
diff --git a/test/specialarrays.jl b/test/specialarrays.jl
index b812a80..3b6417e 100644
--- a/test/specialarrays.jl
+++ b/test/specialarrays.jl
@@ -43,6 +43,30 @@ using StaticArrays
             @test A \ Y isa SMatrix{4, 10, Float64}
             @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, Y) isa SMatrix{4, 10, Float64}
+            @test whiten(A, Y) ≈ cholesky(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, Y) isa SMatrix{4, 10, Float64}
+            @test unwhiten(A, Y) ≈ cholesky(Matrix(A)).L * Matrix(Y)
+
+            @test quad(A, x) isa Float64
+            @test quad(A, x) ≈ Vector(x)' * Matrix(A) * Vector(x)
+
+            @test quad(A, Y) isa SVector{10, Float64}
+            @test quad(A, Y) ≈ diag(Matrix(Y)' * Matrix(A) * Matrix(Y))
+
+            @test invquad(A, x) isa Float64
+            @test invquad(A, x) ≈ Vector(x)' * (Matrix(A) \ Vector(x))
+
+            @test invquad(A, Y) isa SVector{10, Float64}
+            @test invquad(A, Y) ≈ diag(Matrix(Y)' * (Matrix(A) \ Matrix(Y)))
+
             @test X_A_Xt(A, X) isa SMatrix{10, 10, Float64}
             @test X_A_Xt(A, X) ≈ Matrix(X) * Matrix(A) *  Matrix(X)'
 

From 0aca2f39b2cb80108932095ce94f93be035d471e Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Wed, 27 Sep 2023 17:45:18 +0200
Subject: [PATCH 4/7] Add specialization for `Array` arguments

---
 src/pdiagmat.jl    | 37 ++++++++++++++++++++++++++-----------
 src/pdmat.jl       | 35 ++++++++++++++++++++++++++---------
 src/pdsparsemat.jl | 11 ++++++-----
 src/scalmat.jl     | 12 +++++++-----
 4 files changed, 65 insertions(+), 30 deletions(-)

diff --git a/src/pdiagmat.jl b/src/pdiagmat.jl
index 2147883..b384b44 100644
--- a/src/pdiagmat.jl
+++ b/src/pdiagmat.jl
@@ -118,16 +118,6 @@ function quad(a::PDiagMat, x::AbstractVecOrMat)
         return vec(sum(abs2.(x) .* a.diag; dims = 1))
     end
 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 quad!(r::AbstractArray, a::PDiagMat, x::AbstractMatrix)
     ad = a.diag
@@ -143,8 +133,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)
@@ -184,3 +184,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
+
diff --git a/src/pdmat.jl b/src/pdmat.jl
index 98a75ba..855058e 100644
--- a/src/pdmat.jl
+++ b/src/pdmat.jl
@@ -113,15 +113,6 @@ function quad(a::PDMat, x::AbstractVecOrMat)
         return vec(sum(abs2, aU_x; dims = 1))
     end
 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 quad!(r::AbstractArray, a::PDMat, x::AbstractMatrix)
     @check_argdims axes(r) == axes(x, 2)
@@ -135,6 +126,17 @@ function quad!(r::AbstractArray, a::PDMat, x::AbstractMatrix)
     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)
@@ -173,3 +175,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
+
diff --git a/src/pdsparsemat.jl b/src/pdsparsemat.jl
index 68e4121..47fd427 100644
--- a/src/pdsparsemat.jl
+++ b/src/pdsparsemat.jl
@@ -100,11 +100,6 @@ function quad(a::PDSparseMat, x::AbstractVecOrMat)
     z = sparse(chol_lower(cholesky(a)))' * x
     return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 1))
 end
-function invquad(a::PDSparseMat, x::AbstractVecOrMat)
-    @check_argdims a.dim == size(x, 1)
-    z = sparse(chol_lower(cholesky(a))) \ x
-    return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 1))
-end
 
 function quad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
     @check_argdims eachindex(r) == axes(x, 2)
@@ -114,6 +109,12 @@ function quad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
     return r
 end
 
+function invquad(a::PDSparseMat, x::AbstractVecOrMat)
+    @check_argdims a.dim == size(x, 1)
+    z = sparse(chol_lower(cholesky(a))) \ x
+    return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 1))
+end
+
 function invquad!(r::AbstractArray, a::PDSparseMat, x::AbstractMatrix)
     @check_argdims eachindex(r) == axes(x, 2)
     for i in axes(x, 2)
diff --git a/src/scalmat.jl b/src/scalmat.jl
index cb15ca1..a0a5ae3 100644
--- a/src/scalmat.jl
+++ b/src/scalmat.jl
@@ -111,6 +111,13 @@ function quad(a::ScalMat, x::AbstractVecOrMat)
         return vec(sum(wsq, x; dims=1))
     end
 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))
+end
+
 function invquad(a::ScalMat, x::AbstractVecOrMat)
     @check_argdims a.dim == size(x, 1)
     if x isa AbstractVector
@@ -125,11 +132,6 @@ function invquad(a::ScalMat, x::AbstractVecOrMat)
     end
 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))
-end
 function invquad!(r::AbstractArray, a::ScalMat, x::AbstractMatrix)
     @check_argdims eachindex(r) == axes(x, 2)
     @check_argdims a.dim == size(x, 1)

From 61f801f35dabd492c72c7ee7b39db98db516cd73 Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Tue, 3 Oct 2023 09:36:49 +0200
Subject: [PATCH 5/7] Fixes for PDSparseMat

---
 src/pdsparsemat.jl | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/pdsparsemat.jl b/src/pdsparsemat.jl
index 2b927b8..0f43ae6 100644
--- a/src/pdsparsemat.jl
+++ b/src/pdsparsemat.jl
@@ -104,7 +104,12 @@ end
 
 function quad(a::PDSparseMat, x::AbstractVecOrMat)
     @check_argdims a.dim == size(x, 1)
-    z = sparse(chol_lower(cholesky(a)))' * x
+    # `*` is not defined for `UP` factor components,
+    # so we can't use `chol_upper(a.chol) * x`
+    # Moreover, `sparse` is only defined for `L` factor components
+    C = a.chol
+    UP = transpose(sparse(C.L))[:, C.p]
+    z = UP * x
     return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 1))
 end
 
@@ -118,7 +123,7 @@ end
 
 function invquad(a::PDSparseMat, x::AbstractVecOrMat)
     @check_argdims a.dim == size(x, 1)
-    z = sparse(chol_lower(cholesky(a))) \ x
+    z = chol_lower(cholesky(a)) \ x
     return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 1))
 end
 

From 3f17bdee92702d81ffc1376912b57910bd65ed70 Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Tue, 3 Oct 2023 09:52:17 +0200
Subject: [PATCH 6/7] More fixes for PDSparseMat

---
 src/pdsparsemat.jl | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/pdsparsemat.jl b/src/pdsparsemat.jl
index 0f43ae6..4adabc6 100644
--- a/src/pdsparsemat.jl
+++ b/src/pdsparsemat.jl
@@ -114,9 +114,17 @@ function quad(a::PDSparseMat, x::AbstractVecOrMat)
 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)
+    # `*` is not defined for `UP` factor components,
+    # so we can't use `chol_upper(a.chol) * x`
+    # Moreover, `sparse` is only defined for `L` factor components
+    C = a.chol
+    UP = transpose(sparse(C.L))[:, C.p]
+    z = similar(r, a.dim) # buffer to save allocations
+    @inbounds for i in axes(x, 2)
+        copyto!(z, view(x, :, i))
+        lmul!(UP, z)
+        r[i] = sum(abs2, z)
     end
     return r
 end
@@ -128,9 +136,15 @@ function invquad(a::PDSparseMat, x::AbstractVecOrMat)
 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)
+    aL = chol_lower(cholesky(a))
+    # `\` with `PtL` factor components is not implemented for views
+    # and the generic fallback requires indexing which is not supported
+    xi = similar(x, a.dim)
+    @inbounds for i in axes(x, 2)
+        copyto!(xi, view(x, :, i))
+        r[i] = sum(abs2, aL \ xi)
     end
     return r
 end

From 3b050b9b459917f6be0983eca64a5898d695e79d Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Wed, 4 Oct 2023 02:03:12 +0200
Subject: [PATCH 7/7] Fix tests

---
 src/pdsparsemat.jl | 53 ++++++++++++++++++++++++----------------------
 1 file changed, 28 insertions(+), 25 deletions(-)

diff --git a/src/pdsparsemat.jl b/src/pdsparsemat.jl
index 4adabc6..72cab5b 100644
--- a/src/pdsparsemat.jl
+++ b/src/pdsparsemat.jl
@@ -104,47 +104,50 @@ end
 
 function quad(a::PDSparseMat, x::AbstractVecOrMat)
     @check_argdims a.dim == size(x, 1)
-    # `*` is not defined for `UP` factor components,
-    # so we can't use `chol_upper(a.chol) * x`
-    # Moreover, `sparse` is only defined for `L` factor components
-    C = a.chol
-    UP = transpose(sparse(C.L))[:, C.p]
-    z = UP * x
-    return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 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 axes(r) == axes(x, 2)
-    # `*` is not defined for `UP` factor components,
-    # so we can't use `chol_upper(a.chol) * x`
-    # Moreover, `sparse` is only defined for `L` factor components
-    C = a.chol
-    UP = transpose(sparse(C.L))[:, C.p]
-    z = similar(r, a.dim) # buffer to save allocations
-    @inbounds for i in axes(x, 2)
-        copyto!(z, view(x, :, i))
-        lmul!(UP, z)
-        r[i] = sum(abs2, z)
+    # 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 = chol_lower(cholesky(a)) \ x
-    return x isa AbstractVector ? sum(abs2, z) : vec(sum(abs2, z; dims = 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 axes(r) == axes(x, 2)
     @check_argdims a.dim == size(x, 1)
-    aL = chol_lower(cholesky(a))
-    # `\` with `PtL` factor components is not implemented for views
-    # and the generic fallback requires indexing which is not supported
-    xi = similar(x, a.dim)
+    z = similar(r, a.dim) # buffer to save allocations
     @inbounds for i in axes(x, 2)
-        copyto!(xi, view(x, :, i))
-        r[i] = sum(abs2, aL \ xi)
+        xi = view(x, :, i)
+        copyto!(z, xi)
+        ldiv!(a.chol, z)
+        r[i] = dot(xi, z)
     end
     return r
 end