add doubleMM testset

Author: bgctw

.gitignore

@@ -6,3 +6,4 @@
 /docs/build/
 test/Manifest*.toml
 dev/Manifest*.toml
+tmp/

dev/doubleMM.jl

@@ -1,14 +1,3 @@
-() -> begin
-    using SimpleChains, BenchmarkTools, Static, OptimizationOptimisers
-    import Zygote
-    using StatsFuns: logistic
-    using UnicodePlots
-    using Distributions
-    using StableRNGs
-    using LinearAlgebra, StatsBase, Combinatorics
-    using Random
-end
-
 using Test
 using HybridVariationalInference
 using StableRNGs

ext/HybridVariationalInferenceLuxExt.jl

@@ -16,7 +16,7 @@ function HVI.construct_LuxApplicator(m::Chain; device = gpu_device())
     st = st |> device
     stateful_layer = StatefulLuxLayer{true}(m, nothing, st)
     #stateful_layer(x_o_gpu[:, 1:n_site_batch], ps_ca)
-    int_ϕ = ComponentArrayInterpreter(ps_ca)
+    int_ϕ = get_concrete(ComponentArrayInterpreter(ps_ca))
     LuxApplicator(stateful_layer, int_ϕ)
 end
 

test/Project.toml

@@ -1,12 +1,18 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"

test/runtests.jl

@@ -7,6 +7,9 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml
         @time @safetestset "test_gencovar" include("test_gencovar.jl")
         #@safetestset "test" include("test/test_SimpleChains.jl")
         @time @safetestset "test_SimpleChains" include("test_SimpleChains.jl")
+        #@safetestset "test" include("test/test_doubleMM.jl")
+        @time @safetestset "test_doubleMM" include("test_doubleMM.jl")
+        #
         #@safetestset "test" include("test/test_Flux.jl")
         @time @safetestset "test_Flux" include("test_Flux.jl")
         #@safetestset "test" include("test/test_Lux.jl")

test/test_doubleMM.jl

@@ -0,0 +1,94 @@
+using Test
+using HybridVariationalInference
+using StableRNGs
+using Random
+using Statistics
+using ComponentArrays: ComponentArrays as CA
+
+using SimpleChains
+using MLUtils
+import Zygote
+
+using OptimizationOptimisers
+
+const case = DoubleMM.DoubleMMCase()
+const MLengine = Val(nameof(SimpleChains))
+scenario = (:default,)
+
+par_templates = get_hybridcase_par_templates(case; scenario)
+
+(; n_covar, n_site, n_batch, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
+
+rng = StableRNG(111)
+(; xM, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o
+) = gen_hybridcase_synthetic(case, rng; scenario);
+
+@testset "gen_hybridcase_synthetic" begin
+    @test isapprox(
+        vec(mean(CA.getdata(θMs_true); dims = 2)), CA.getdata(par_templates.θM), rtol = 0.02)
+    @test isapprox(vec(std(CA.getdata(θMs_true); dims = 2)),
+        CA.getdata(par_templates.θM) .* 0.1, rtol = 0.02)
+
+    # test same results for same rng
+    rng2 = StableRNG(111)
+    gen2 = gen_hybridcase_synthetic(case, rng2; scenario);
+    @test gen2.y_o == y_o
+end
+
+@testset "loss_g" begin
+    g, ϕg0 = gen_hybridcase_MLapplicator(case, MLengine; scenario);
+
+    function loss_g(ϕg, x, g)
+        ζMs = g(x, ϕg) # predict the log of the parameters
+        θMs = exp.(ζMs)
+        loss = sum(abs2, θMs .- θMs_true)
+        return loss, θMs
+    end
+    loss_g(ϕg0, xM, g)
+    Zygote.gradient(x -> loss_g(x, xM, g)[1], ϕg0);
+
+    optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg, xM, g)[1],
+        Optimization.AutoZygote())
+    optprob = Optimization.OptimizationProblem(optf, ϕg0);
+    res = Optimization.solve(optprob, Adam(0.02), callback = callback_loss(100), maxiters = 600);
+
+    ϕg_opt1 = res.u;
+    pred = loss_g(ϕg_opt1, xM, g)
+    θMs_pred = pred[2]
+    #scatterplot(vec(θMs_true), vec(θMs_pred))
+    @test cor(vec(θMs_true), vec(θMs_pred)) > 0.9
+end
+
+@testset "loss_gf" begin
+    #----------- fit g and θP to y_o
+    g, ϕg0 = gen_hybridcase_MLapplicator(case, MLengine; scenario);
+    f = gen_hybridcase_PBmodel(case; scenario)
+
+    int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
+        ϕg = 1:length(ϕg0), θP = par_templates.θP))
+    p = p0 = vcat(ϕg0, par_templates.θP .* 0.8);  # slightly disturb θP_true
+
+    # Pass the site-data for the batches as separate vectors wrapped in a tuple
+    train_loader = MLUtils.DataLoader((xM, xP, y_o), batchsize = n_batch)
+
+    loss_gf = get_loss_gf(g, f, y_global_o, int_ϕθP)
+    l1 = loss_gf(p0, train_loader.data...)[1]
+
+    optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
+        Optimization.AutoZygote())
+    optprob = OptimizationProblem(optf, p0, train_loader)
+
+    res = Optimization.solve(
+        optprob, Adam(0.02), callback = callback_loss(100), maxiters = 1000);
+
+    l1, y_pred_global, y_pred, θMs_pred = loss_gf(res.u, train_loader.data...)
+    @test isapprox(par_templates.θP, int_ϕθP(res.u).θP, rtol = 0.11)
+    @test cor(vec(θMs_true), vec(θMs_pred)) > 0.9
+
+    () -> begin
+        scatterplot(vec(θMs_true), vec(θMs_pred))
+        scatterplot(log.(vec(θMs_true)), log.(vec(θMs_pred)))
+        scatterplot(vec(y_pred), vec(y_o))
+        hcat(par_templates.θP, int_ϕθP(p0).θP, int_ϕθP(res.u).θP)
+    end
+end