Simplify kron(::PDiagMat, ::PDiagMat) (#164)

* Simplify `kron(::PDiagMat, ::PDiagMat)` * Update Project.toml * Add test with StaticArrays
JuliaStats · Jun 18, 2022 · 8d88c15 · 8d88c15 · devmotion · Jun 18, 2022
1 parent da9fe95
commit 8d88c15
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "PDMats"
 uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
-version = "0.11.12"
+version = "0.11.13"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

diff --git a/src/pdiagmat.jl b/src/pdiagmat.jl
@@ -57,7 +57,7 @@ function /(x::AbstractVecOrMat, a::PDiagMat)
     # return matrix for 1-element vectors `x`, consistent with LinearAlgebra
     return reshape(x, Val(2)) ./ permutedims(a.diag) # = (a' \ x')'
 end
-Base.kron(A::PDiagMat, B::PDiagMat) = PDiagMat( vcat([A.diag[i] * B.diag for i in 1:dim(A)]...) )
+Base.kron(A::PDiagMat, B::PDiagMat) = PDiagMat(vec(permutedims(A.diag) .* B.diag))
 
 ### Algebra
 

diff --git a/test/kron.jl b/test/kron.jl
@@ -3,7 +3,7 @@ using Test
 using LinearAlgebra: LinearAlgebra
 
 function _pd_kron_compare(A::AbstractPDMat, B::AbstractPDMat)
-    PDAkB_kron = kron(A, B)
+    PDAkB_kron = @inferred kron(A, B)
     PDAkB_dense = PDMat( kron( Matrix(A), Matrix(B) ) )
     _pd_compare(PDAkB_kron, PDAkB_dense)
 end

diff --git a/test/specialarrays.jl b/test/specialarrays.jl
@@ -18,6 +18,7 @@ using StaticArrays
         # Diagonal matrix
         D = PDiagMat(@SVector(rand(4)))
         @test D isa PDiagMat{Float64, <:SVector{4, Float64}}
+        @test @inferred(kron(D, D)) isa PDiagMat{Float64, <:SVector{16, Float64}}
 
         x = @SVector rand(4)
         X = @SMatrix rand(10, 4)