Merge pull request #17 from invenia/gm/submat

Define submat for PDMat subtypes

Project.toml

@@ -1,12 +1,14 @@
 name = "PDMatsExtras"
 uuid = "2c7acb1b-7338-470f-b38f-951d2bcb9193"
 authors = ["Invenia Technical Computing"]
-version = "2.4.0"
+version = "2.5.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 
 [compat]
 ChainRulesCore = "0.9.17, 0.10"

src/PDMatsExtras.jl

@@ -2,11 +2,15 @@ module PDMatsExtras
 using ChainRulesCore
 using LinearAlgebra
 using PDMats
+using SparseArrays: sparse
+using SuiteSparse
 import Base: *, \
 
 export PSDMat, WoodburyPDMat
+export submat
 
 include("psd_mat.jl")
 include("woodbury_pd_mat.jl")
+include("utils.jl")
 
 end

src/utils.jl

@@ -0,0 +1,16 @@
+"""
+    submat(A::AbstractPDMat, inds)
+
+Select a square submatrix of an `AbstractPDMat` specified by `inds`.
+Returns a `AbstractPDMat` of the same type.
+"""
+submat(A::PDMat, inds) = PDMat(A[inds, inds])
+submat(A::PSDMat, inds) = PSDMat(A[inds, inds])
+submat(A::PDiagMat, inds) = PDiagMat(A.diag[inds])
+submat(A::PDSparseMat, inds) = PDSparseMat(sparse(A[inds, inds]))
+submat(A::WoodburyPDMat, inds) = WoodburyPDMat(A.A[inds, :], A.D, Diagonal(A.S.diag[inds]))
+function submat(A::ScalMat, inds)
+    # inds should not be able to "access" non-existent indices
+    checkbounds(Bool, A, inds) || throw(BoundsError(A, inds))
+    return ScalMat(length(inds), A.value)
+end

test/psd_mat.jl

@@ -1,21 +1,3 @@
-
-# NOTE: We could probably do a more thorough testing job if we able to override the
-# `t_tripod` and `t_whiten` test in PDMats.
-test_matrices = Dict(
-    "Positive definite" => [
-        0.796911  0.602112  0.766136  0.247788
-        0.602112  0.480312  0.605538  0.218218
-        0.766136  0.605538  1.28666   0.290052
-        0.247788  0.218218  0.290052  0.130588
-    ],
-    "Positive semi-definite" => [
-        10.8145   -9.27226   1.67126   4.02515
-        -9.27226   8.08443  -1.48168  -4.27258
-        1.67126  -1.48168   1.31866   1.43293
-        4.02515  -4.27258   1.43293   6.76801
-    ]
-)
-
 # We write the tests like this to match the style of `testutils.jl`.
 # We define this here, rather than in `testutils.jl` because we're not allowed to change the
 # contents of that file, which should always match the PDMats v0.10 version of that file
@@ -35,7 +17,7 @@ end
 @testset "PSDMat" begin
     verbose = 1
     @testset "Positive definite" begin
-        M = test_matrices["Positive definite"]
+        M = TEST_MATRICES["Positive definite"]
         pivoted = cholesky(M, Val(true))
         C = PSDMat(M, pivoted)
         test_pdmat(
@@ -49,7 +31,7 @@ end
         pdtest_kron(C, C.mat, verbose)
     end
     @testset "Positive semi-definite" begin
-        M = test_matrices["Positive semi-definite"]
+        M = TEST_MATRICES["Positive semi-definite"]
         @test !isposdef(M)
         pivoted = cholesky(M, Val(true); check=false)
         C = PSDMat(M, pivoted)
@@ -74,7 +56,7 @@ end
     n_tsamples = 10^6
 
     @testset "Positive definite" begin
-        g = MvNormal(means, PSDMat(test_matrices["Positive definite"]))
+        g = MvNormal(means, PSDMat(TEST_MATRICES["Positive definite"]))
         d = length(g)
         μ = mean(g)
         Σ = cov(g)
@@ -119,7 +101,7 @@ end
     end
 
     @testset "Positive semi-definite" begin
-        g = MvNormal(means, PSDMat(test_matrices["Positive semi-definite"]))
+        g = MvNormal(means, PSDMat(TEST_MATRICES["Positive semi-definite"]))
         d = length(g)
         μ = mean(g)
         Σ = cov(g)

test/runtests.jl

@@ -4,15 +4,35 @@ using Distributions
 using FiniteDifferences
 using LinearAlgebra
 using PDMats
+using SparseArrays: sparse
 using Random
 using Test
 using Zygote
 
 Random.seed!(1)
+
+# NOTE: We could probably do a more thorough testing job if we able to override the
+# `t_tripod` and `t_whiten` test in PDMats.
+const TEST_MATRICES = Dict(
+    "Positive definite" => [
+        0.796911  0.602112  0.766136  0.247788
+        0.602112  0.480312  0.605538  0.218218
+        0.766136  0.605538  1.28666   0.290052
+        0.247788  0.218218  0.290052  0.130588
+    ],
+    "Positive semi-definite" => [
+        10.8145   -9.27226   1.67126   4.02515
+        -9.27226   8.08443  -1.48168  -4.27258
+        1.67126  -1.48168   1.31866   1.43293
+        4.02515  -4.27258   1.43293   6.76801
+    ]
+)
+
 @testset "PDMatsExtras.jl" begin
     include("testutils.jl")
     include("test_ad.jl")
 
     include("psd_mat.jl")
     include("woodbury_pd_mat.jl")
+    include("utils.jl")
 end

test/utils.jl

@@ -0,0 +1,69 @@
+@testset "utils.jl" begin
+    @testset "submat" begin
+
+        @testset "PDMat" begin
+            M = PDMat(TEST_MATRICES["Positive definite"])
+            M_dense = Matrix(M)
+
+            subM = submat(M, [1, 3])
+            @test subM isa PDMat
+            @test Matrix(subM) == M_dense[[1, 3], [1, 3]]
+        end
+
+        @testset "PDDiagMat" begin
+            M = PDiagMat(diag(TEST_MATRICES["Positive definite"]))
+            M_dense = Matrix(M)
+
+            subM = submat(M, [1, 3])
+            @test subM isa PDiagMat
+            @test Matrix(subM) == M_dense[[1, 3], [1, 3]]
+        end
+
+        @testset "PDSparseMat" begin
+            M = PDSparseMat(sparse(TEST_MATRICES["Positive definite"]))
+            M_dense = Matrix(M)
+
+            subM = submat(M, [1, 3])
+            @test subM isa PDSparseMat
+            @test Matrix(subM) == M_dense[[1, 3], [1, 3]]
+        end
+
+        @testset "ScalMat" begin
+            M = ScalMat(3, 0.1)
+            M_dense = Matrix(M)
+
+            subM = submat(M, [1, 3])
+            @test subM isa ScalMat
+            @test Matrix(subM) == M_dense[[1, 3], [1, 3]]
+            @test_throws BoundsError submat(M, [1, 10])
+        end
+
+        @testset "PSDMat" begin
+            SM = TEST_MATRICES["Positive semi-definite"]
+            pivoted = cholesky(SM, Val(true); check=false)
+            M = PSDMat(SM, pivoted)
+            M_dense = Matrix(M)
+
+            subM = submat(M, [1, 3])
+            @test subM isa PSDMat
+            @test Matrix(subM) == M_dense[[1, 3], [1, 3]]
+        end
+
+        @testset "WoodburyPDMat" begin
+            A = randn(4, 2)
+            D = Diagonal(randn(2).^2 .+ 1)
+            S = Diagonal(randn(4).^2 .+ 1)
+            x = randn(size(A, 1))
+
+            W = WoodburyPDMat(A, D, S)
+            W_dense = PDMat(Symmetric(Matrix(W)))
+
+            subW = submat(W, [1, 3])
+            @test subW isa WoodburyPDMat
+            @test subW.A == W.A[[1, 3], :]
+            @test subW.D == W.D
+            @test subW.S == Diagonal(W.S.diag[[1, 3]])
+            @test Matrix(subW) == W_dense[[1, 3], [1, 3]]
+        end
+    end
+end