implement HybridPointSolver on cpu

Author: bgctw

Project.toml

@@ -9,10 +9,12 @@ BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
@@ -34,12 +36,14 @@ BlockDiagonals = "0.1.42"
 CUDA = "5.5.2"
 ChainRulesCore = "1.25"
 Combinatorics = "1.0.2"
+CommonSolve = "0.2.4"
 ComponentArrays = "0.15.19"
 Flux = "v0.15.2, 0.16"
 GPUArraysCore = "0.1, 0.2"
 LinearAlgebra = "1.10.0"
 Lux = "1.4.2"
 MLUtils = "0.4.5"
+Optimization = "3.19.3, 4"
 Random = "1.10.0"
 SimpleChains = "0.4"
 StatsBase = "0.34.4"

dev/doubleMM.jl

@@ -5,31 +5,123 @@ using StableRNGs
 using Random
 using Statistics
 using ComponentArrays: ComponentArrays as CA
-
+using Optimization
+using OptimizationOptimisers # Adam
+using UnicodePlots
 using SimpleChains
-import Flux 
+using Flux
+using MLUtils
+
+rng = StableRNG(114)
+scenario = NTuple{0, Symbol}()
+#scenario = (:use_Flux,)
+
+#------ setup synthetic data and training data loader
+(; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc
+) = gen_hybridcase_synthetic(rng, DoubleMM.DoubleMMCase(); scenario);
+get_train_loader = (rng; n_batch, kwargs...) -> MLUtils.DataLoader((xM, xP, y_o, y_unc), batchsize = n_batch)
+σ_o = exp(first(y_unc)/2)
+
+# assign the train_loader, otherwise it eatch time creates another version of synthetic data
+prob0 = update(HybridProblem(DoubleMM.DoubleMMCase(); scenario); get_train_loader)
+
+#------- pointwise hybrid model fit
+#solver = HybridPointSolver(; alg = Adam(0.02), n_batch = 30)
+solver = HybridPointSolver(; alg = Adam(0.01), n_batch = 10)
+#solver = HybridPointSolver(; alg = Adam(), n_batch = 200)
+(; ϕ, resopt) = solve(prob0, solver; scenario,
+    rng, callback = callback_loss(100), maxiters = 1200);
+prob0o = update(prob0; ϕg=ϕ.ϕg, θP=ϕ.θP)
+y_pred_global, y_pred, θMs = gf(prob0o, xM, xP);
+scatterplot(θMs_true[1,:], θMs[1,:])
+scatterplot(θMs_true[2,:], θMs[2,:])
+
+# do a few steps without minibatching, 
+#   by providing the data rather than the DataLoader
+# train_loader0 = get_hybridproblem_train_dataloader(rng, prob0; scenario, n_batch=1000)
+# get_train_loader_data = (args...; kwargs...) -> train_loader0.data
+# prob1 = update(prob0o; get_train_loader = get_train_loader_data)
+prob1 = prob0o
+
+#solver1 = HybridPointSolver(; alg = Adam(0.05), n_batch = n_site)
+solver1 = HybridPointSolver(; alg = Adam(0.01), n_batch = n_site)
+(; ϕ, resopt) = solve(prob1, solver1; scenario, rng, 
+    callback = callback_loss(20), maxiters = 600);
+prob1o = update(prob1; ϕg=ϕ.ϕg, θP=ϕ.θP)
+y_pred_global, y_pred, θMs = gf(prob1o, xM, xP);
+scatterplot(θMs_true[1,:], θMs[1,:])
+scatterplot(θMs_true[2,:], θMs[2,:])
+prob1o.θP
+scatterplot(vec(y_true), vec(y_pred))
+
+() -> begin # with more iterations?
+    prob2 = prob1o
+    (; ϕ, resopt) = solve(prob2, solver1; scenario, rng, 
+        callback = callback_loss(20), maxiters = 600);
+    prob2o = update(prob2; ϕg=ϕ.ϕg, θP=ϕ.θP)
+    y_pred_global, y_pred, θMs = gf(prob2o, xM, xP);
+    prob2o.θP
+end
+
+#----------- fit g to true θMs 
+# and fit gf starting from true parameters
+prob = prob0
+g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario);
+(; transP, transM) = get_hybridproblem_transforms(prob; scenario)
+
+function loss_g(ϕg, x, g, transM)
+    ζMs = g(x, ϕg) # predict the log of the parameters
+    θMs = reduce(hcat, map(transM, eachcol(ζMs))) # transform each column
+    loss = sum(abs2, θMs .- θMs_true)
+    return loss, θMs
+end
+loss_g(ϕg0, xM, g, transM)
+
+optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg, xM, g, transM)[1],
+    Optimization.AutoZygote())
+optprob = Optimization.OptimizationProblem(optf, ϕg0);
+res = Optimization.solve(optprob, Adam(0.015), callback = callback_loss(100), maxiters = 2000);
+
+ϕg_opt1 = res.u;
+l1, θMs = loss_g(ϕg_opt1, xM, g, transM)
+#scatterplot(θMs_true[1,:], θMs[1,:])
+scatterplot(θMs_true[2,:], θMs[2,:]) # able to fit θMs[2,:]
+
+prob3 = update(prob0, ϕg = ϕg_opt1, θP = θP_true)
+solver1 = HybridPointSolver(; alg = Adam(0.01), n_batch = n_site)
+(; ϕ, resopt) = solve(prob3, solver1; scenario, rng, 
+    callback = callback_loss(50), maxiters = 600);
+prob3o = update(prob3; ϕg=ϕ.ϕg, θP=ϕ.θP)
+y_pred_global, y_pred, θMs = gf(prob3o, xM, xP);
+scatterplot(θMs_true[2,:], θMs[2,:])
+prob3o.θP
+scatterplot(vec(y_true), vec(y_pred))
+scatterplot(vec(y_true), vec(y_o))
+scatterplot(vec(y_pred), vec(y_o))
+
+() -> begin # optimized loss is indeed lower than with true parameters
+    int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
+        ϕg = 1:length(prob0.ϕg), θP = prob0.θP))
+    loss_gf = get_loss_gf(prob0.g, prob0.transM, prob0.f, Float32[], int_ϕθP)
+    loss_gf(vcat(prob3.ϕg, prob3.θP), xM, xP, y_o, y_unc)[1]
+    loss_gf(vcat(prob3o.ϕg, prob3o.θP), xM, xP, y_o, y_unc)[1]
+    #
+    loss_gf(vcat(prob2o.ϕg, prob2o.θP), xM, xP, y_o, y_unc)[1]
+end
+    
+#----------- Hybrid Variational inference 
+
 using MLUtils
 import Zygote
 
 using CUDA
-using OptimizationOptimisers
 using Bijectors
-using UnicodePlots
 
-const prob = DoubleMM.DoubleMMCase()
-scenario = (:default,)
-rng = StableRNG(111)
-
-par_templates = get_hybridproblem_par_templates(prob; scenario)
 
 #n_covar = get_hybridproblem_n_covar(prob; scenario)
 #, n_batch, n_θM, n_θP) = get_hybridproblem_sizes(prob; scenario)
 
-(; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc
-) = gen_hybridcase_synthetic(rng, prob; scenario);
-
-n_covar = size(xM,1)
-
+n_covar = size(xM, 1)
 
 #----- fit g to θMs_true
 g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario);
@@ -92,8 +184,6 @@ FT = get_hybridproblem_float_type(prob; scenario)
     θP_true, θMs_true[:, 1], ϕg_opt1, n_batch; transP, transM);
 ϕ_true = ϕ
 
-
-
 () -> begin
     coef_logσ2_logMs = [-5.769 -3.501; -0.01791 0.007951]
     logσ2_logP = CA.ComponentVector(r0 = -8.997, K2 = -5.893)
@@ -245,7 +335,7 @@ y_pred = predict_gf(rng, g_flux, f, res.u, xM_gpu, xP, interpreters;
 size(y_pred) # n_obs x n_site, n_sample_pred
 
 σ_o_post = dropdims(std(y_pred; dims = 3), dims = 3);
-σ_o = exp.(y_unc[:,1] / 2)
+σ_o = exp.(y_unc[:, 1] / 2)
 
 #describe(σ_o_post)
 hcat(σ_o, fill(mean_σ_o_MC, length(σ_o)),

src/AbstractHybridProblem.jl

@@ -123,7 +123,7 @@ function get_hybridproblem_float_type(prob::AbstractHybridProblem; scenario=())
 end
 
 """
-    get_hybridproblem_train_dataloader([rng,] ::AbstractHybridProblem; scenario)
+    get_hybridproblem_train_dataloader([rng,] ::AbstractHybridProblem; scenario, n_batch)
 
 Return a DataLoader that provides a tuple of
 - `xM`: matrix of covariates, with one column per site
@@ -132,9 +132,8 @@ Return a DataLoader that provides a tuple of
 - `y_unc`: matrix `sizeof(y_o)` of uncertainty information 
 """
 function get_hybridproblem_train_dataloader(rng::AbstractRNG, prob::AbstractHybridProblem; 
-    scenario = ())
+    scenario = (), n_batch = 10)
     (; xM, xP, y_o, y_unc) = gen_hybridcase_synthetic(rng, prob; scenario)
-    n_batch = 10
     xM_gpu = :use_Flux ∈ scenario ? CuArray(xM) : xM
     train_loader = MLUtils.DataLoader((xM_gpu, xP, y_o, y_unc), batchsize = n_batch)
     return(train_loader)

src/DoubleMM/f_doubleMM.jl

@@ -1,11 +1,11 @@
 struct DoubleMMCase <: AbstractHybridProblem end
 
 
-θP = CA.ComponentVector{Float32}(r0 = 0.3, K2 = 2.0)
-θM = CA.ComponentVector{Float32}(r1 = 0.5, K1 = 0.2)
+const θP = CA.ComponentVector{Float32}(r0 = 0.3, K2 = 2.0)
+const θM = CA.ComponentVector{Float32}(r1 = 0.5, K1 = 0.2)
 
-transP = elementwise(exp)
-transM = Stacked(elementwise(identity), elementwise(exp))
+const transP = elementwise(exp)
+const transM = Stacked(elementwise(identity), elementwise(exp))
 
 
 const int_θdoubleMM = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM)))
@@ -54,13 +54,13 @@ end
 #     return Float32
 # end
 
-const xP_S1 = Float32[1.0, 1.0, 1.0, 1.0, 0.4, 0.3, 0.1]
+const xP_S1 = Float32[0.5, 0.5, 0.5, 0.5, 0.4, 0.3, 0.1]
 const xP_S2 = Float32[1.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0]
 
 function HVI.gen_hybridcase_synthetic(rng::AbstractRNG, prob::DoubleMMCase;
         scenario = ())
     n_covar_pc = 2
-    n_site = 200
+    n_site = 800
     n_covar = 5
     n_θM = length(θM)
     FloatType = get_hybridproblem_float_type(prob; scenario)

src/HybridProblem.jl

@@ -1,37 +1,77 @@
-struct HybridProblem <: AbstractHybridProblem 
+struct HybridProblem <: AbstractHybridProblem
     θP
     θM
     f
     g
     ϕg
-    py    
+    py
     transP
     transM
     cor_starts # = (P=(1,),M=(1,))
-    train_loader
+    get_train_loader
     # inner constructor to constrain the types
     function HybridProblem(
-        θP::CA.ComponentVector, θM::CA.ComponentVector, 
-        g::AbstractModelApplicator, ϕg::AbstractVector, 
-        f::Function, 
-        py::Function,
-        transM::Union{Function, Bijectors.Transform}, 
-        transP::Union{Function, Bijectors.Transform}, 
-        train_loader::DataLoader,
-        cor_starts::NamedTuple = (P=(1,), M=(1,)))
-        new(θP, θM, f, g, ϕg, py, transM, transP, cor_starts, train_loader)
+            θP::CA.ComponentVector, θM::CA.ComponentVector,
+            g::AbstractModelApplicator, ϕg::AbstractVector,
+            f::Function,
+            py::Function,
+            transM::Union{Function, Bijectors.Transform},
+            transP::Union{Function, Bijectors.Transform},
+            #train_loader::DataLoader,
+            # return a function that constructs the trainloader based on n_batch
+            get_train_loader::Function,
+            cor_starts::NamedTuple = (P = (1,), M = (1,)))
+        new(θP, θM, f, g, ϕg, py, transM, transP, cor_starts, get_train_loader)
     end
 end
 
-function HybridProblem(θP::CA.ComponentVector, θM::CA.ComponentVector, 
-    # note no ϕg argument and g_chain unconstrained
-    g_chain, f::Function, 
-    args...; rng = Random.default_rng(), kwargs...)
+function HybridProblem(θP::CA.ComponentVector, θM::CA.ComponentVector,
+        # note no ϕg argument and g_chain unconstrained
+        g_chain, f::Function,
+        args...; rng = Random.default_rng(), kwargs...)
     # dispatches on type of g_chain
     g, ϕg = construct_ChainsApplicator(rng, g_chain, eltype(θM))
     HybridProblem(θP, θM, g, ϕg, f, args...; kwargs...)
 end
 
+function HybridProblem(prob::AbstractHybridProblem; scenario = ())
+    (; θP, θM) = get_hybridproblem_par_templates(prob; scenario)
+    g, ϕg = get_hybridproblem_MLapplicator(prob; scenario)
+    f = get_hybridproblem_PBmodel(prob; scenario)
+    py = get_hybridproblem_neg_logden_obs(prob; scenario)
+    (; transP, transM) = get_hybridproblem_transforms(prob; scenario)
+    get_train_loader = let prob = prob, scenario = scenario
+        function inner_get_train_loader(rng::AbstractRNG; kwargs...)
+            get_hybridproblem_train_dataloader(rng::AbstractRNG, prob; scenario, kwargs...)
+        end
+    end
+    cor_starts = get_hybridproblem_cor_starts(prob; scenario)
+    HybridProblem(θP, θM, g, ϕg, f, py, transP, transM, get_train_loader, cor_starts)
+end
+
+function update(prob::HybridProblem;
+        θP::CA.ComponentVector = prob.θP,
+        θM::CA.ComponentVector = prob.θM,
+        g::AbstractModelApplicator = prob.g, ϕg::AbstractVector = prob.ϕg,
+        f::Function = prob.f,
+        py::Function = prob.py,
+        transM::Union{Function, Bijectors.Transform} = prob.transM,
+        transP::Union{Function, Bijectors.Transform} = prob.transP,
+        get_train_loader::Function = prob.get_train_loader,
+        cor_starts::NamedTuple = prob.cor_starts)
+    # prob.θP = θP
+    # prob.θM = θM
+    # prob.f = f
+    # prob.g = g
+    # prob.ϕg = ϕg
+    # prob.py = py
+    # prob.transM = transM
+    # prob.transP = transP
+    # prob.cor_starts = cor_starts
+    # prob.get_train_loader = get_train_loader
+    HybridProblem(θP, θM, g, ϕg, f, py, transP, transM, get_train_loader, cor_starts)
+end
+
 function get_hybridproblem_par_templates(prob::HybridProblem; scenario::NTuple = ())
     (; θP = prob.θP, θM = prob.θM)
 end
@@ -54,12 +94,12 @@ function get_hybridproblem_PBmodel(prob::HybridProblem; scenario::NTuple = ())
     prob.f
 end
 
-function get_hybridproblem_MLapplicator(prob::HybridProblem; scenario::NTuple = ());
+function get_hybridproblem_MLapplicator(prob::HybridProblem; scenario::NTuple = ())
     prob.g, prob.ϕg
 end
 
-function get_hybridproblem_train_dataloader(rng::AbstractRNG, prob::HybridProblem; scenario = ())
-    return(prob.train_loader)
+function get_hybridproblem_train_dataloader(rng::AbstractRNG, prob::HybridProblem; kwargs...)
+    return prob.get_train_loader(rng; kwargs...)
 end
 
 function get_hybridproblem_cor_starts(prob::HybridProblem; scenario = ())
@@ -69,6 +109,3 @@ end
 # function get_hybridproblem_float_type(prob::HybridProblem; scenario::NTuple = ()) 
 #     eltype(prob.θM)
 # end
-
-
-