Improve broadcasting of PDMat and PDiagMat (#197)

Author: devmotion

Project.toml

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

src/pdiagmat.jl

@@ -32,6 +32,9 @@ Base.Matrix(a::PDiagMat) = Matrix(Diagonal(a.diag))
 LinearAlgebra.diag(a::PDiagMat) = copy(a.diag)
 LinearAlgebra.cholesky(a::PDiagMat) = Cholesky(Diagonal(map(sqrt, a.diag)), 'U', 0)
 
+### Treat as a `Diagonal` matrix in broadcasting since that is better supported
+Base.broadcastable(a::PDiagMat) = Base.broadcastable(Diagonal(a.diag))
+
 ### Inheriting from AbstractMatrix
 
 function Base.getindex(a::PDiagMat, i::Integer)

src/pdmat.jl

@@ -46,6 +46,9 @@ Base.Matrix(a::PDMat) = Matrix(a.mat)
 LinearAlgebra.diag(a::PDMat) = diag(a.mat)
 LinearAlgebra.cholesky(a::PDMat) = a.chol
 
+### Work with the underlying matrix in broadcasting
+Base.broadcastable(a::PDMat) = Base.broadcastable(a.mat)
+
 ### Inheriting from AbstractMatrix
 
 Base.getindex(a::PDMat, i::Int) = getindex(a.mat, i)

test/pdmtypes.jl

@@ -256,4 +256,32 @@ using Test
             @test_throws DimensionMismatch PDSparseMat(A[1:(end - 1), 1:(end - 1)], C)
         end
     end
+
+    @testset "Subtraction" begin
+        # This falls back to the generic method in Julia based on broadcasting
+        dim = 4
+        x = rand(dim, dim)
+        A = PDMat(x' * x + I)
+        @test Base.broadcastable(A) == A.mat
+
+        B = PDiagMat(rand(dim))
+        @test Base.broadcastable(B) == Diagonal(B.diag)
+
+        for X in (A, B), Y in (A, B)
+            @test X - Y isa (X === Y === B ? Diagonal{Float64, Vector{Float64}} : Matrix{Float64})
+            @test X - Y ≈ Matrix(X) - Matrix(Y)
+        end
+
+        C = ScalMat(dim, rand())
+        @test A - C isa Matrix{Float64}
+        @test A - C ≈ Matrix(A) - Matrix(C)
+        @test C - A isa Matrix{Float64}
+        @test C - A ≈ Matrix(C) - Matrix(A)
+
+        # ScalMat does not behave nicely with PDiagMat
+        @test_broken B - C isa Diagonal{Float64, Vector{Float64}}
+        @test B - C ≈ Matrix(B) - Matrix(C)
+        @test_broken C - B isa Diagonal{Float64, Vector{Float64}}
+        @test C - B ≈ Matrix(C) - Matrix(B)
+    end
 end

test/specialarrays.jl

@@ -84,6 +84,22 @@ using StaticArrays
             @test Xt_invA_X(A, Y) isa Symmetric{Float64,<:SMatrix{10, 10, Float64}}
             @test Xt_invA_X(A, Y) ≈ Matrix(Y)' * (Matrix(A) \ Matrix(Y))
         end
+
+        # Subtraction falls back to the generic method in Base which is based on broadcasting
+        @test Base.broadcastable(PDS) == PDS.mat
+        @test Base.broadcastable(D) == Diagonal(D.diag)
+        for A in (PDS, D), B in (PDS, D)
+            @test A - B isa SMatrix{4, 4, Float64}
+            @test A - B ≈ Matrix(A) - Matrix(B)
+        end
+
+        # ScalMat does not behave nicely with broadcasting currently
+        for A in (PDS, D)
+            @test_broken A - E isa SMatrix{4, 4, Float64}
+            @test_broken E - A isa SMatrix{4, 4, Float64}
+            @test A - E ≈ Matrix(A) - Matrix(E)
+            @test E - A ≈ Matrix(E) - Matrix(A)
+        end
     end
 
     @testset "BandedMatrices" begin