Create and refactor test library with tags

GeneralisedFilters/Project.toml

@@ -31,18 +31,15 @@ HypothesisTests = "0.11"
 NNlib = "0.9"
 LogExpFunctions = "0.3"
 OffsetArrays = "1.14.1"
-SSMProblems = "0.4"
 Statistics = "1.11.1"
 StatsBase = "0.34.3"
 
 [extras]
-HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 
 [targets]
-test = ["HypothesisTests", "LinearAlgebra", "StableRNGs", "Test", "TestItemRunner", "TestItems", "PDMats"]
+test = ["PDMats", "StableRNGs", "Test", "TestItemRunner", "TestItems"]

GeneralisedFilters/src/GFTest/GFTest.jl

@@ -0,0 +1,15 @@
+module GFTest
+
+using CUDA
+using LinearAlgebra
+using NNlib
+using Random
+
+using GeneralisedFilters
+using SSMProblems
+
+include("utils.jl")
+include("models/linear_gaussian.jl")
+include("models/dummy_linear_gaussian.jl")
+
+end

GeneralisedFilters/test/utils/models/dummy_linear_gaussian.jl → GeneralisedFilters/src/GFTest/models/dummy_linear_gaussian.jl

@@ -14,6 +14,8 @@
     dynamics, this model can be used in Rao-Blackwellised settings.
 """
 
+export InnerDynamics, create_dummy_linear_gaussian_model
+
 """
     Inner dynamics of the dummy linear Gaussian model.
 

GeneralisedFilters/src/GFTest/utils.jl

@@ -0,0 +1,4 @@
+function rand_cov(rng::AbstractRNG, T::Type{<:Real}, d::Int)
+    Σ = rand(rng, T, d, d)
+    return Σ * Σ'
+end

GeneralisedFilters/src/GeneralisedFilters.jl

@@ -147,4 +147,7 @@ include("algorithms/kalman.jl")
 include("algorithms/forward.jl")
 include("algorithms/rbpf.jl")
 
+# Unit-testing helper module
+include("GFTest/GFTest.jl")
+
 end

GeneralisedFilters/src/algorithms/kalman.jl

@@ -97,32 +97,32 @@ function gemm_static_shmem_element_kernel!(
     C, A, B, batch_size::Int32, ::Val{M}, ::Val{N}, ::Val{P}, ::Val{K}
 ) where {M,N,P,K}
     tid = threadIdx().x
-    
+
     # Loading phase - coalesced reads down columns
-    shmem_a = @cuStaticSharedMem(Float32, (M, N, K))
-    shmem_b = @cuStaticSharedMem(Float32, (P, N, K))
+    shmem_a = CuStaticSharedArray(Float32, (M, N, K))
+    shmem_b = CuStaticSharedArray(Float32, (P, N, K))
 
     # Load A
-    for i in tid:blockDim().x:M*N*K
-        matrix = div(i - 1i32, M*N) + 1i32
-        mn_idx = mod1(i, M*N)
+    for i in tid:(blockDim().x):(M * N * K)
+        matrix = div(i - 1i32, M * N) + 1i32
+        mn_idx = mod1(i, M * N)
         m = mod1(mn_idx, M)
         n = div(mn_idx - 1i32, M) + 1i32
         matrix_idx = (blockIdx().x - 1i32) * K + matrix
-        
+
         if matrix_idx <= batch_size
             shmem_a[m, n, matrix] = A[m, n, matrix_idx]
         end
     end
 
     # Load B - separate loop with no dependency on M
-    for i in tid:blockDim().x:P*N*K
-        matrix = div(i - 1i32, P*N) + 1i32
-        pn_idx = mod1(i, P*N)
+    for i in tid:(blockDim().x):(P * N * K)
+        matrix = div(i - 1i32, P * N) + 1i32
+        pn_idx = mod1(i, P * N)
         p = mod1(pn_idx, P)
         n = div(pn_idx - 1i32, P) + 1i32
         matrix_idx = (blockIdx().x - 1i32) * K + matrix
-        
+
         if matrix_idx <= batch_size
             shmem_b[p, n, matrix] = B[p, n, matrix_idx]
         end
@@ -131,8 +131,8 @@ function gemm_static_shmem_element_kernel!(
     sync_threads()
 
     # Computing phase - each thread computes one element of C
-    local_matrix = div(tid - 1i32, M*P) + 1i32
-    matrix_thread = mod1(tid, M*P)
+    local_matrix = div(tid - 1i32, M * P) + 1i32
+    matrix_thread = mod1(tid, M * P)
     row = mod1(matrix_thread, M)
     col = div(matrix_thread - 1i32, M) + 1i32
     matrix_idx = (blockIdx().x - 1i32) * K + local_matrix

GeneralisedFilters/test/batch_kalman_test.jl

@@ -1,4 +1,4 @@
-@testitem "Batch Kalman test" begin
+@testitem "Batch Kalman test" tags = [:gpu] begin
     using GeneralisedFilters
     using Distributions
     using LinearAlgebra

GeneralisedFilters/test/resamplers.jl

@@ -13,7 +13,7 @@
     rng = StableRNG(SEED)
 
     xs = rand(rng, Normal(0, 1), N)
-    ws = pdf(Normal(1, 1), xs) ./ pdf(Normal(0, 1), xs)
+    ws = map(x -> pdf(Normal(1, 1), x) / pdf(Normal(0, 1), x), xs)
     ws ./= sum(ws)
 
     μ0 = sum(ws .* xs)
@@ -59,29 +59,29 @@ end
     @test μ0 ≈ μ1 rtol = 1e-1
 end
 
-@testitem "Test GPU systematic resampling" setup = [GPUResamplingTestSetup] begin
+@testitem "Test GPU systematic resampling" setup = [GPUResamplingTestSetup] tags = [:gpu] begin
     idxs = GeneralisedFilters.sample_ancestors(rng, Systematic(), ws)
     @test length(idxs) == N
     μ1 = sum(xs[idxs]) / N
     @test μ0 ≈ μ1 rtol = 1e-1
 end
 
-@testitem "Test GPU stratified resampling" setup = [GPUResamplingTestSetup] begin
+@testitem "Test GPU stratified resampling" setup = [GPUResamplingTestSetup] tags = [:gpu] begin
     idxs = GeneralisedFilters.sample_ancestors(rng, Stratified(), ws)
     @test length(idxs) == N
     μ1 = sum(xs[idxs]) / N
     @test μ0 ≈ μ1 rtol = 1e-1
 end
 
-@testitem "Test GPU offspring-to-ancestors" begin
+@testitem "Test GPU offspring-to-ancestors" tags = [:gpu] begin
     using CUDA
     offspring = CuVector{Int}([0, 2, 2, 3, 5])
     true_ancestors = CuVector{Int}([2, 2, 4, 5, 5])
     ancestors = GeneralisedFilters.offspring_to_ancestors(offspring)
     @test ancestors == true_ancestors
 end
 
-@testitem "Test GPU ancestors-to-offspring" begin
+@testitem "Test GPU ancestors-to-offspring" tags = [:gpu] begin
     using CUDA
     ancestors = CuVector{Int}([4, 2, 2, 3, 1])
     true_offspring = CuVector{Int}([1, 2, 1, 1, 0])

GeneralisedFilters/test/runtests.jl

@@ -6,7 +6,7 @@ using TestItemRunner
 include("batch_kalman_test.jl")
 include("resamplers.jl")
 
-@testitem "Kalman filter test" setup = [TestModels] begin
+@testitem "Kalman filter test" begin
     using GeneralisedFilters
     using Distributions
     using LinearAlgebra
@@ -18,13 +18,13 @@ include("resamplers.jl")
 
     for Dy in Dys
         rng = StableRNG(1234)
-        model = TestModels.create_linear_gaussian_model(rng, Dx, Dy)
+        model = GeneralisedFilters.GFTest.create_linear_gaussian_model(rng, Dx, Dy)
         _, _, ys = sample(rng, model, 1)
 
         filtered, ll = GeneralisedFilters.filter(rng, model, KalmanFilter(), ys)
 
         # Let Z = [X0, X1, Y1] be the joint state vector
-        μ_Z, Σ_Z = TestModels._compute_joint(model, 1)
+        μ_Z, Σ_Z = GeneralisedFilters.GFTest._compute_joint(model, 1)
 
         # Condition on observations using formula for MVN conditional distribution. See: 
         # https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Conditional_distributions
@@ -46,7 +46,7 @@ include("resamplers.jl")
     end
 end
 
-@testitem "Kalman smoother test" setup = [TestModels] begin
+@testitem "Kalman smoother test" begin
     using GeneralisedFilters
     using Distributions
     using LinearAlgebra
@@ -58,13 +58,13 @@ end
 
     for Dy in Dys
         rng = StableRNG(1234)
-        model = TestModels.create_linear_gaussian_model(rng, Dx, Dy)
+        model = GeneralisedFilters.GFTest.create_linear_gaussian_model(rng, Dx, Dy)
         _, _, ys = sample(rng, model, 2)
 
         states, ll = GeneralisedFilters.smooth(rng, model, KalmanSmoother(), ys)
 
         # Let Z = [X0, X1, X2, Y1, Y2] be the joint state vector
-        μ_Z, Σ_Z = TestModels._compute_joint(model, 2)
+        μ_Z, Σ_Z = GeneralisedFilters.GFTest._compute_joint(model, 2)
 
         # Condition on observations using formula for MVN conditional distribution. See: 
         # https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Conditional_distributions
@@ -182,9 +182,8 @@ end
     @test ll ≈ log(marginal)
 end
 
-@testitem "Kalman-RBPF test" setup = [TestModels] begin
+@testitem "Kalman-RBPF test" begin
     using LogExpFunctions: softmax
-    using SSMProblems
     using StableRNGs
     using StatsBase
 
@@ -197,7 +196,7 @@ end
 
     rng = StableRNG(SEED)
 
-    full_model, hier_model = TestModels.create_dummy_linear_gaussian_model(
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
         rng, D_outer, D_inner, D_obs
     )
     _, _, ys = sample(rng, full_model, T)
@@ -223,7 +222,7 @@ end
     @test last(kf_states.μ) ≈ sum(last.(getproperty.(zs, :μ)) .* weights) rtol = 1e-2
 end
 
-@testitem "GPU Kalman-RBPF test" setup = [TestModels] begin
+@testitem "GPU Kalman-RBPF test" tags = [:gpu] begin
     using CUDA
     using LinearAlgebra
     using NNlib
@@ -240,7 +239,7 @@ end
 
     rng = StableRNG(SEED)
 
-    full_model, hier_model = TestModels.create_dummy_linear_gaussian_model(
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
         rng, D_outer, D_inner, D_obs, ET
     )
     _, _, ys = sample(rng, full_model, T)
@@ -266,7 +265,7 @@ end
     @test eltype(xs) == ET
 end
 
-@testitem "RBPF ancestory test" setup = [TestModels] begin
+@testitem "RBPF ancestory test" begin
     using SSMProblems
     using StableRNGs
 
@@ -275,7 +274,9 @@ end
     N_particles = 100
 
     rng = StableRNG(SEED)
-    full_model, hier_model = TestModels.create_dummy_linear_gaussian_model(rng, 1, 1, 1)
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
+        rng, 1, 1, 1
+    )
     _, _, ys = sample(rng, full_model, T)
 
     # Manually create tree to force expansion on second step
@@ -294,9 +295,8 @@ end
     paths = GeneralisedFilters.get_ancestry(tree)
 end
 
-@testitem "Dense ancestry test" setup = [TestModels] begin
+@testitem "Dense ancestry test" begin
     using GeneralisedFilters
-    using SSMProblems
     using StableRNGs
     using PDMats
     using LinearAlgebra
@@ -319,7 +319,7 @@ end
     N_particles = max(10, K + 2)
 
     rng = StableRNG(SEED)
-    model = TestModels.create_linear_gaussian_model(rng, 1, 1)
+    model = GeneralisedFilters.GFTest.create_linear_gaussian_model(rng, 1, 1)
     _, _, ys = sample(rng, model, K)
 
     ref_traj = OffsetVector([rand(rng, 1) for _ in 0:K], -1)
@@ -337,9 +337,8 @@ end
     @test GeneralisedFilters.get_ancestry(cb.container, 1) == ref_traj
 end
 
-@testitem "CSMC test" setup = [TestModels] begin
+@testitem "CSMC test" begin
     using GeneralisedFilters
-    using SSMProblems
     using StableRNGs
     using PDMats
     using LinearAlgebra
@@ -360,7 +359,7 @@ end
     N_sample = 10000
 
     rng = StableRNG(SEED)
-    model = TestModels.create_linear_gaussian_model(rng, Dx, Dy)
+    model = GeneralisedFilters.GFTest.create_linear_gaussian_model(rng, Dx, Dy)
     _, _, ys = sample(rng, model, K)
 
     # Kalman smoother
@@ -390,11 +389,9 @@ end
     @test csmc_mean ≈ state.μ rtol = 1e-1
 end
 
-@testitem "RBCSMC test" setup = [TestModels] begin
+@testitem "RBCSMC test" begin
     using GeneralisedFilters
-    using SSMProblems
     using GaussianDistributions
-    using SSMProblems
     using StableRNGs
     using PDMats
     using LinearAlgebra
@@ -416,7 +413,7 @@ end
     N_sample = 500
 
     rng = StableRNG(SEED)
-    full_model, hier_model = TestModels.create_dummy_linear_gaussian_model(
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
         rng, D_outer, D_inner, D_obs, T
     )
     _, _, ys = sample(rng, full_model, K)
@@ -459,7 +456,7 @@ end
     @test state.μ[2] ≈ only(mean(getproperty.(z_trajectories, :μ))) rtol = 1e-1
 end
 
-@testitem "GPU Conditional Kalman-RBPF execution test" setup = [TestModels] begin
+@testitem "GPU Conditional Kalman-RBPF execution test" tags = [:gpu] begin
     using CUDA
     using OffsetArrays
     using SSMProblems
@@ -475,7 +472,7 @@ end
 
     rng = StableRNG(1234)
 
-    full_model, hier_model = TestModels.create_dummy_linear_gaussian_model(
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
         rng, D_outer, D_inner, D_obs, T
     )
     _, _, ys = sample(rng, full_model, K)
@@ -486,7 +483,9 @@ end
             CuArray(rand(rng, T, D_outer, 1)),
             GeneralisedFilters.BatchGaussianDistribution(
                 CuArray(rand(rng, T, D_inner, 1)),
-                CuArray(reshape(TestModels.rand_cov(rng, T, D_inner), Val(3))),
+                CuArray(
+                    reshape(GeneralisedFilters.GFTest.rand_cov(rng, T, D_inner), Val(3))
+                ),
             ),
         ) for _ in 0:K
     ]
@@ -499,7 +498,7 @@ end
     @test typeof(ll) == T
 end
 
-@testitem "GPU-RBPF ancestory test" setup = [TestModels] begin
+@testitem "GPU-RBPF ancestory test" tags = [:gpu] begin
     using GeneralisedFilters
     using CUDA
     using LinearAlgebra
@@ -516,7 +515,7 @@ end
 
     rng = StableRNG(1234)
 
-    full_model, hier_model = TestModels.create_dummy_linear_gaussian_model(
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
         rng, D_outer, D_inner, D_obs, T
     )
     _, _, ys = sample(rng, full_model, K)
@@ -542,7 +541,7 @@ end
     ancestry = GeneralisedFilters.get_ancestry(tree, K)
 end
 
-@testitem "GPU Conditional Kalman-RBPF validity test" setup = [TestModels] begin
+@testitem "GPU Conditional Kalman-RBPF validity test" tags = [:gpu, :long] begin
     using GeneralisedFilters
     using CUDA
     using NNlib
@@ -563,7 +562,7 @@ end
 
     rng = StableRNG(1234)
 
-    full_model, hier_model = TestModels.create_dummy_linear_gaussian_model(
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
         rng, D_outer, D_inner, D_obs, T
     )
     _, _, ys = sample(rng, full_model, K)