Extend Woodbury multiplication to Diagonals (#14)

* Minor changes

* Bump version

* Add check that Diagonal matricies are posdef

* Add ArgumentError test

* Update Project.toml

Co-authored-by: Miha Zgubic <[email protected]>

Co-authored-by: Miha Zgubic <[email protected]>

Author: AlexRobson

Project.toml

@@ -1,7 +1,7 @@
 name = "PDMatsExtras"
 uuid = "2c7acb1b-7338-470f-b38f-951d2bcb9193"
 authors = ["Invenia Technical Computing"]
-version = "2.2.0"
+version = "2.3.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/woodbury_pd_mat.jl

@@ -88,3 +88,14 @@ end
 # NOTE: the parameterisation to scale up the Woodbury matrix is not unique. Here we
 # implement one way to scale it.
 *(a::WoodburyPDMat, c::Real) = WoodburyPDMat(a.A, a.D * c, a.S * c)
+*(c::Real, a::WoodburyPDMat) = a * c
+function *(a::WoodburyPDMat, c::Diagonal{T}) where {T<:Real}
+    isposdef(c) || throw(ArgumentError("c must be positive definite"))
+    WoodburyPDMat(sqrt(c) * a.A, a.D, a.S * c)
+end
+*(c::Diagonal{T}, a::WoodburyPDMat) where {T<:Real} = a * c
+function *(c1::Diagonal{T}, a::WoodburyPDMat, c2::Diagonal{T}) where {T<:Real}
+    isposdef(c1) || throw(ArgumentError("c1 must be positive definite"))
+    isposdef(c2) || throw(ArgumentError("c2 must be positive definite"))
+    WoodburyPDMat(sqrt(c1) * sqrt(c2) * a.A, a.D, c1 * a.S * c2)
+end

test/woodbury_pd_mat.jl

@@ -43,6 +43,25 @@
         c = 2.0
         @test c * W == W * c
         @test c * W_dense ≈ c * W atol=1e-6
+        @test (c * W) isa WoodburyPDMat
+
+        c = Diagonal(2.0 * ones(4,))
+        @test c * W == W * c
+        @test c * W_dense ≈ c * W atol=1e-6
+        @test (c * W) isa WoodburyPDMat
+
+        c1 = Diagonal(2.0 * ones(4,))
+        c2 = Diagonal(3.0 * ones(4,))
+        c_neg = Diagonal([1,2,-2,3])
+
+        @test c2 * W * c1 == c1 * W * c2
+        @test c1 * W * c2 ≈ c1 * W_dense * c2
+        @test (c1 * W * c2) isa WoodburyPDMat
+
+        @test_throws(ArgumentError, c_neg * W)
+        @test_throws(ArgumentError, c_neg * W * c2)
+        @test_throws(ArgumentError, c1 * W * c_neg)
+
     end
 
     @testset "MvNormal logpdf" begin