Merge pull request #101 from JuliaGaussianProcesses/tgf/sumkernel

Refactor treatment of `KernelSum`

Project.toml

@@ -1,7 +1,7 @@
 name = "TemporalGPs"
 uuid = "e155a3c4-0841-43e1-8b83-a0e4f03cc18f"
 authors = ["willtebbutt <[email protected]> and contributors"]
-version = "0.6.0"
+version = "0.6.1"
 
 [deps]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"

src/gp/lti_sde.jl

@@ -283,66 +283,78 @@ end
 # Sum
 
 function lgssm_components(k::KernelSum, ts::AbstractVector, storage_type::StorageType)
-    As_l, as_l, Qs_l, emission_proj_l, x0_l = lgssm_components(k.kernels[1], ts, storage_type)
-    As_r, as_r, Qs_r, emission_proj_r, x0_r = lgssm_components(k.kernels[2], ts, storage_type)
-
-    As = _map(blk_diag, As_l, As_r)
-    as = _map(vcat, as_l, as_r)
-    Qs = _map(blk_diag, Qs_l, Qs_r)
-    emission_projections = _sum_emission_projections(emission_proj_l, emission_proj_r)
-    x0 = Gaussian(vcat(x0_l.m, x0_r.m), blk_diag(x0_l.P, x0_r.P))
-
+    lgssms = lgssm_components.(k.kernels, Ref(ts), Ref(storage_type))
+    As_kernels = getindex.(lgssms, 1)
+    as_kernels = getindex.(lgssms, 2)
+    Qs_kernels = getindex.(lgssms, 3)
+    emission_proj_kernels = getindex.(lgssms, 4)
+    x0_kernels = getindex.(lgssms, 5)
+
+    As = _map(block_diagonal, As_kernels...)
+    as = _map(vcat, as_kernels...)
+    Qs = _map(block_diagonal, Qs_kernels...)
+    emission_projections = _sum_emission_projections(emission_proj_kernels...)
+    x0 = Gaussian(mapreduce(x -> getproperty(x, :m), vcat, x0_kernels), block_diagonal(getproperty.(x0_kernels, :P)...))
     return As, as, Qs, emission_projections, x0
 end
 
-function _sum_emission_projections(
-    (Hs_l, hs_l)::Tuple{AbstractVector, AbstractVector},
-    (Hs_r, hs_r)::Tuple{AbstractVector, AbstractVector},
-)
-    return map(vcat, Hs_l, Hs_r), hs_l + hs_r
+function _sum_emission_projections(Hs_hs::Tuple{AbstractVector, AbstractVector}...)
+    return map(vcat, first.(Hs_hs)...), sum(last.(Hs_hs))
 end
 
 function _sum_emission_projections(
-    (Cs_l, cs_l, Hs_l, hs_l)::Tuple{AbstractVector, AbstractVector, AbstractVector, AbstractVector},
-    (Cs_r, cs_r, Hs_r, hs_r)::Tuple{AbstractVector, AbstractVector, AbstractVector, AbstractVector},
+    Cs_cs_Hs_hs::Tuple{AbstractVector, AbstractVector, AbstractVector, AbstractVector}...,
 )
-    Cs = _map(vcat, Cs_l, Cs_r)
-    cs = cs_l + cs_r
-    Hs = _map(blk_diag, Hs_l, Hs_r)
-    hs = _map(vcat, hs_l, hs_r)
-    return Cs, cs, Hs, hs
+    Cs = getindex.(Cs_cs_Hs_hs, 1)
+    cs = getindex.(Cs_cs_Hs_hs, 2)
+    Hs = getindex.(Cs_cs_Hs_hs, 3)
+    hs = getindex.(Cs_cs_Hs_hs, 4)
+    C = _map(vcat, Cs...)
+    c = sum(cs)
+    H = _map(block_diagonal, Hs...)
+    h = _map(vcat, hs...)
+    return C, c, H, h
 end
 
 Base.vcat(x::Zeros{T, 1}, y::Zeros{T, 1}) where {T} = Zeros{T}(length(x) + length(y))
 
-function blk_diag(A::AbstractMatrix{T}, B::AbstractMatrix{T}) where {T}
-    return hvcat(
-        (2, 2),
-        A, zeros(T, size(A, 1), size(B, 2)), zeros(T, size(B, 1), size(A, 2)), B,
-    )
-end
-
-function ChainRulesCore.rrule(::typeof(blk_diag), A, B)
-    blk_diag_rrule(Δ::AbstractThunk) = blk_diag_rrule(unthunk(Δ))
-    function blk_diag_rrule(Δ)
-        ΔA = Δ[1:size(A, 1), 1:size(A, 2)]
-        ΔB = Δ[size(A, 1)+1:end, size(A, 2)+1:end]
-        return NoTangent(), ΔA, ΔB
+function block_diagonal(As::AbstractMatrix{T}...) where {T}
+    nblocks = length(As)
+    sizes = size.(As)
+    Xs = [i == j ? As[i] : Zeros{T}(sizes[j][1], sizes[i][2]) for i in 1:nblocks, j in 1:nblocks]
+    return hvcat(ntuple(_ -> nblocks, nblocks), Xs...)
+end
+
+function ChainRulesCore.rrule(::typeof(block_diagonal), As::AbstractMatrix...)
+    szs = size.(As)
+    row_szs = (0, cumsum(first.(szs))...)
+    col_szs = (0, cumsum(last.(szs))...)
+    block_diagonal_rrule(Δ::AbstractThunk) = block_diagonal_rrule(unthunk(Δ))
+    function block_diagonal_rrule(Δ)
+        ΔAs = ntuple(length(As)) do i
+            Δ[(row_szs[i]+1):row_szs[i+1], (col_szs[i]+1):col_szs[i+1]]
+        end
+        return NoTangent(), ΔAs...
     end
-    return blk_diag(A, B), blk_diag_rrule
+    return block_diagonal(As...), block_diagonal_rrule
 end
 
-function blk_diag(A::SMatrix{DA, DA, T}, B::SMatrix{DB, DB, T}) where {DA, DB, T}
-    zero_AB = zeros(SMatrix{DA, DB, T})
-    zero_BA = zeros(SMatrix{DB, DA, T})
-    return [[A zero_AB]; [zero_BA B]]
+function block_diagonal(As::SMatrix...)
+    nblocks = length(As)
+    sizes = size.(As)
+    Xs = [i == j ? As[i] : zeros(SMatrix{sizes[j][1], sizes[i][2]}) for i in 1:nblocks, j in 1:nblocks]
+    return hcat(Base.splat(vcat).(eachrow(Xs))...)
 end
 
-function ChainRulesCore.rrule(::typeof(blk_diag), A::SMatrix{DA, DA, T}, B::SMatrix{DB, DB, T}) where {DA, DB, T}
-    function blk_diag_adjoint(Δ)
-        ΔA = Δ[SVector{DA}(1:DA), SVector{DA}(1:DA)]
-        ΔB = Δ[SVector{DB}((DA+1):(DA+DB)), SVector{DB}((DA+1):(DA+DB))]
-        return NoTangent(), ΔA, ΔB
+function ChainRulesCore.rrule(::typeof(block_diagonal), As::SMatrix...)
+    szs = size.(As)
+    row_szs = (0, cumsum(first.(szs))...)
+    col_szs = (0, cumsum(last.(szs))...)
+    function block_diagonal_rrule(Δ)
+        ΔAs = ntuple(length(As)) do i
+            Δ[SVector{szs[i][1]}((row_szs[i]+1):row_szs[i+1]), SVector{szs[i][2]}((col_szs[i]+1):col_szs[i+1])]
+        end
+        return NoTangent(), ΔAs...
     end
-    return blk_diag(A, B), blk_diag_adjoint
+    return block_diagonal(As...), block_diagonal_rrule
 end

src/space_time/pseudo_point.jl

@@ -383,11 +383,12 @@ function dtc_post_emissions(k::ScaledKernel, x_new::AbstractVector, storage::Sto
 end
 
 function dtc_post_emissions(k::KernelSum, x_new::AbstractVector, storage::StorageType)
-    (Cs_l, cs_l, Hs_l, hs_l), Σs_l = dtc_post_emissions(k.kernels[1], x_new, storage)
-    (Cs_r, cs_r, Hs_r, hs_r), Σs_r = dtc_post_emissions(k.kernels[2], x_new, storage)
-    Cs = _map(vcat, Cs_l, Cs_r)
-    cs = cs_l + cs_r
-    Hs = _map(blk_diag, Hs_l, Hs_r)
-    hs = _map(vcat, hs_l, hs_r)
-    return (Cs, cs, Hs, hs), _map(+, Σs_l, Σs_r)
+    post_emissions = dtc_post_emissions.(k.kernels, Ref(x_new), Ref(storage))
+    Cs_cs_Hs_hs = getindex.(post_emissions, 1)
+    Σs = getindex.(post_emissions, 2)
+    Cs = _map(vcat, getindex.(Cs_cs_Hs_hs, 1)...)
+    cs = sum(getindex.(Cs_cs_Hs_hs, 2))
+    Hs = _map(block_diagonal, getindex.(Cs_cs_Hs_hs, 3)...)
+    hs = _map(vcat, getindex.(Cs_cs_Hs_hs, 4)...)
+    return (Cs, cs, Hs, hs), sum(Σs)
 end

test/gp/lti_sde.jl

@@ -15,11 +15,12 @@ end
 println("lti_sde:")
 @testset "lti_sde" begin
 
-    @testset "blk_diag" begin
+    @testset "block_diagonal" begin
         A = randn(2, 2)
         B = randn(3, 3)
-        test_rrule(TemporalGPs.blk_diag, A, B; check_inferred=false)
-        test_rrule(TemporalGPs.blk_diag, SMatrix{2, 2}(A), SMatrix{3, 3}(B))
+        C = randn(5, 5)
+        test_rrule(TemporalGPs.block_diagonal, A, B, C; check_inferred=false)
+        test_rrule(TemporalGPs.block_diagonal, SMatrix{2, 2}(A), SMatrix{3, 3}(B), SMatrix{5, 5}(C); check_inferred=false)
     end
 
     @testset "SimpleKernel parameter types" begin
@@ -71,12 +72,19 @@ println("lti_sde:")
                 (name="stretched-λ=$λ", val=Matern32Kernel() ∘ ScaleTransform(λ))
             end,
 
+            # TEST_TOFIX
             # Summed kernels.
             # (
-                # name="sum-Matern12Kernel-Matern32Kernel",
-                # val=1.5 * Matern12Kernel() ∘ ScaleTransform(0.1) +
-                    # 0.3 * Matern32Kernel() ∘ ScaleTransform(1.1),
-            # ), # TEST_TOFIX
+            #     name="sum-Matern12Kernel-Matern32Kernel",
+            #     val=1.5 * Matern12Kernel() ∘ ScaleTransform(0.1) +
+            #         0.3 * Matern32Kernel() ∘ ScaleTransform(1.1),
+            # ), 
+            # (
+            #     name="sum-Matern32Kernel-Matern52Kernel-ConstantKernel",
+            #     val = 2.0 * Matern32Kernel() +
+            #         0.5 * Matern52Kernel() +
+            #         1.0 * ConstantKernel(),
+            # ),
         )
 
         # Construct a Gauss-Markov model with either dense storage or static storage.

test/runtests.jl

@@ -6,7 +6,7 @@ ENV["TESTING"] = "TRUE"
 # ["test util", "test models" "test models-lgssm" "test gp" "test space_time"]
 # Select any of this to test a particular aspect.
 # To test everything, simply set GROUP to "all"
-# ENV["GROUP"] = "test space_time"
+ENV["GROUP"] = "test gp"
 const GROUP = get(ENV, "GROUP", "test")
 OUTER_GROUP = first(split(GROUP, ' '))
 

test/util/mul.jl

@@ -11,8 +11,8 @@ using LinearAlgebra: mul!
     A_Matrix = randn(rng, P, Q)
     At_Matrix = collect(A_Matrix')
 
-    A_blk_diag = BlockDiagonal([randn(rng, P, P), randn(rng, P + 1, P + 1)])
-    At_blk_diag = BlockDiagonal(map(collect ∘ transpose, blocks(A_blk_diag)))
+    A_block_diag = BlockDiagonal([randn(rng, P, P), randn(rng, P + 1, P + 1)])
+    At_block_diag = BlockDiagonal(map(collect ∘ transpose, blocks(A_block_diag)))
 
     settings = [
         (
@@ -24,17 +24,17 @@ using LinearAlgebra: mul!
         ),
         (
             name="BlockDiagonal{Float64, Matrix{Float64}}",
-            A=A_blk_diag,
-            At=At_blk_diag,
-            B=randn(rng, size(A_blk_diag, 2), Q),
-            C=randn(rng, size(A_blk_diag, 1), Q),
+            A=A_block_diag,
+            At=At_block_diag,
+            B=randn(rng, size(A_block_diag, 2), Q),
+            C=randn(rng, size(A_block_diag, 1), Q),
         ),
         (
             name="BlockDiagonal{Float64, BlockDiagonal{Float64, Matrix{Float64}}}",
-            A=BlockDiagonal([A_blk_diag, A_blk_diag]),
-            At=BlockDiagonal([At_blk_diag, At_blk_diag]),
-            B=randn(rng, 2 * size(A_blk_diag, 2), Q),
-            C=randn(rng, 2 * size(A_blk_diag, 1), Q),
+            A=BlockDiagonal([A_block_diag, A_block_diag]),
+            At=BlockDiagonal([At_block_diag, At_block_diag]),
+            B=randn(rng, 2 * size(A_block_diag, 2), Q),
+            C=randn(rng, 2 * size(A_block_diag, 1), Q),
         ),
     ]