implement broadcasted doubleMM f_PBM

move n_batch from solver to problem

remove problem responsibilty of putting dataloader to gpu but do it in fitting and predicting functions

fix error in solver of properly updating problem.phi_unc

Author: bgctw

dev/doubleMM.jl

@@ -28,35 +28,34 @@ gdev = :use_gpu ∈ scenario ? gpu_device() : identity
 cdev = gdev isa MLDataDevices.AbstractGPUDevice ? cpu_device() : identity
 
 #------ setup synthetic data and training data loader
+prob0_ = HybridProblem(DoubleMM.DoubleMMCase(); scenario);
 (; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc
-) = gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario);
-#n_site = get_hybridproblem_n_site(DoubleMM.DoubleMMCase(); scenario)
+) = gen_hybridproblem_synthetic(rng, prob0_; scenario);
+n_site, n_batch = get_hybridproblem_n_site_and_batch(prob0_; scenario)
 ζP_true, ζMs_true = log.(θP_true), log.(θMs_true)
 i_sites = 1:n_site
-xM_cpu = xM;
-xM = xM_cpu |> gdev;
-get_train_loader = (; n_batch, kwargs...) -> MLUtils.DataLoader(
+n_site, n_batch = get_hybridproblem_n_site_and_batch(prob0_; scenario)
+train_dataloader = MLUtils.DataLoader(
     (xM, xP, y_o, y_unc, 1:n_site);
     batchsize = n_batch, partial = false)
 σ_o = exp.(y_unc[:, 1] / 2)
-
 # assign the train_loader, otherwise it eatch time creates another version of synthetic data
-prob0 = HVI.update(HybridProblem(DoubleMM.DoubleMMCase(); scenario); get_train_loader)
+prob0 = HVI.update(prob0_; train_dataloader);
 #tmp = HVI.get_hybridproblem_ϕunc(prob0; scenario)
 
 #------- pointwise hybrid model fit
-solver_point = HybridPointSolver(; alg = OptimizationOptimisers.Adam(0.01), n_batch = 30)
+solver_point = HybridPointSolver(; alg = OptimizationOptimisers.Adam(0.01))
 #solver_point = HybridPointSolver(; alg = Adam(0.01), n_batch = 30)
 #solver_point = HybridPointSolver(; alg = Adam(0.01), n_batch = 10)
 #solver_point = HybridPointSolver(; alg = Adam(), n_batch = 200)
-n_batches_in_epoch = n_site ÷ solver_point.n_batch
+n_batches_in_epoch = n_site ÷ n_batch
 n_epoch = 80
 (; ϕ, resopt, probo) = solve(prob0, solver_point; scenario,
     rng, callback = callback_loss(n_batches_in_epoch * 10),
     maxiters = n_batches_in_epoch * n_epoch);
 # update the problem with optimized parameters
 prob0o = probo;
-y_pred_global, y_pred, θMs = gf(prob0o, xM, xP; scenario);
+y_pred_global, y_pred, θMs = gf(prob0o, scenario);
 plt = scatterplot(θMs_true[1, :], θMs[1, :]);
 lineplot!(plt, 0, 1)
 scatterplot(θMs_true[2, :], θMs[2, :])
@@ -149,10 +148,10 @@ probh = prob0o  # start from point optimized to infer uncertainty
 #probh = prob1o  # start from point optimized to infer uncertainty
 #probh = prob0  # start from no information
 solver_post = HybridPosteriorSolver(;
-    alg = OptimizationOptimisers.Adam(0.01), n_batch = min(50, n_site), n_MC = 3)
+    alg = OptimizationOptimisers.Adam(0.01), n_MC = 3)
 #solver_point = HybridPointSolver(; alg = Adam(), n_batch = 200)
-n_batches_in_epoch = n_site ÷ solver_post.n_batch
-n_epoch = 80
+n_batches_in_epoch = n_site ÷ n_batch
+n_epoch = 40
 (; ϕ, θP, resopt, interpreters, probo) = solve(probh, solver_post; scenario,
     rng, callback = callback_loss(n_batches_in_epoch * 5),
     maxiters = n_batches_in_epoch * n_epoch,
@@ -213,6 +212,7 @@ end
     n_sample_pred = 400
     (; θ, y, entropy_ζ) = predict_gf(rng, prob2o_indep, xM, xP; scenario, n_sample_pred);
     (θ2_indep, y2_indep) = (θ, y)
+    #(θ2_indep, y2_indep) = (θ2, y2)  # workaround to use covarK2 when loading failed
 end
 
 () -> begin # otpimize using LUX
@@ -246,7 +246,7 @@ exp.(ϕunc_VI.coef_logσ2_logMs[1, :])
 
 # test predicting correct obs-uncertainty of predictive posterior
 n_sample_pred = 400
-(; θ, y, entropy_ζ) = predict_gf(rng, prob2o, xM, xP; scenario, n_sample_pred);
+(; θ, y, entropy_ζ) = predict_gf(rng, prob2o; scenario, n_sample_pred);
 (θ2, y2) = (θ, y)
 size(y) # n_obs x n_site, n_sample_pred
 size(θ)  # n_θP + n_site * n_θM x n_sample
@@ -506,12 +506,13 @@ chain = sample(model, NUTS(), MCMCThreads(), ceil(Integer,n_sample_NUTS/n_thread
     using JLD2
     fname = "intermediate/doubleMM_chain_zeta_$(last(scenario)).jld2"
     jldsave(fname, false, IOStream; chain)
-    chain = load(fname, "chain"; iotype = IOStream)
+    chain = load(fname, "chain"; iotype = IOStream);
 end
 
 #ζi = first(eachrow(Array(chain)))
+f_allsites = get_hybridproblem_PBmodel(prob0; scenario, use_all_sites = true)
 ζs = mapreduce(ζi -> transposeMs(ζi, intm_PMs_gen, true), hcat, eachrow(Array(chain)));
-(; θ, y) = HVI.predict_ζf(ζs, f, xP, trans_PMs_gen, intm_PMs_gen);
+(; θ, y) = HVI.predict_ζf(ζs, f_allsites, xP, trans_PMs_gen, intm_PMs_gen);
 (ζs_hmc, θ_hmc, y_hmc) = (ζs, θ, y);
 
 

dev/negLogDensity.pdf

@@ -12,7 +12,7 @@ For a specific prob, provide functions that specify details
 - `get_hybridproblem_train_dataloader` (may use `construct_dataloader_from_synthetic`)
 - `get_hybridproblem_priors` 
 - `get_hybridproblem_n_covar` 
-- `get_hybridproblem_n_site` 
+- `get_hybridproblem_n_site_and_batch` 
 optionally
 - `gen_hybridproblem_synthetic`
 - `get_hybridproblem_float_type` (defaults to `eltype(θM)`)
@@ -125,11 +125,11 @@ function get_hybridproblem_pbmpar_covars(::AbstractHybridProblem; scenario)
 end
 
 """
-    get_hybridproblem_n_site(::AbstractHybridProblem; scenario)
+    get_hybridproblem_n_site_and_batch(::AbstractHybridProblem; scenario)
 
 Provide the number of sites. 
 """
-function get_hybridproblem_n_site end
+function get_hybridproblem_n_site_and_batch end
 
 
 """
@@ -172,30 +172,51 @@ function get_hybridproblem_train_dataloader end
         scenario = (), n_batch)
 
 Construct a dataloader based on `gen_hybridproblem_synthetic`. 
-gdev is applied to xM.
-If :f_on_gpu is in scenario tuple, gdev is also applied to `xP`, `y_o`, and `y_unc`,
-to put the entire data to gpu.
-Alternatively, gdev could be applied to the dataloader, then for each
-iteration the subset of data is separately transferred to gpu.
 """
 function construct_dataloader_from_synthetic(rng::AbstractRNG, prob::AbstractHybridProblem;
         scenario = (), n_batch, 
-        gdev = :use_gpu ∈ scenario ? gpu_device() : identity,
+        #gdev = :use_gpu ∈ scenario ? gpu_device() : identity,
         )
     (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, prob; scenario)
     n_site = size(xM,2)
     @assert length(xP) == n_site
     @assert size(y_o,2) == n_site
     @assert size(y_unc,2) == n_site
     i_sites = 1:n_site
-    xM_dev = gdev(xM)
-    xP_dev, y_o_dev, y_unc_dev = :f_on_gpu ∈ scenario ? 
-        (gdev(xP), gdev(y_o), gdev(y_unc)) : (xP, y_o, y_unc)
-    train_loader = MLUtils.DataLoader((xM_dev, xP_dev, y_o_dev, y_unc_dev, i_sites);
+    train_loader = MLUtils.DataLoader((xM, xP, y_o, y_unc, i_sites);
         batchsize = n_batch, partial = false)
     return (train_loader)
 end
 
+
+"""
+    gdev_hybridproblem_dataloader(dataloader::MLUtils.DataLoader,
+        scenario = (), 
+        gdev = gpu_device(),
+        gdev_M = :use_gpu ∈ scenario ? gdev : identity,
+        gdev_P = :f_on_gpu ∈ scenario ? gdev : identity,
+        batchsize = dataloader.batchsize,
+        partial = dataloader.partial
+        )
+
+Put relevant parts of the DataLoader to gpu, depending on scenario.
+"""
+function gdev_hybridproblem_dataloader(dataloader::MLUtils.DataLoader;
+    scenario = (), 
+    gdev = gpu_device(),
+    gdev_M = :use_gpu ∈ scenario ? gdev : identity,
+    gdev_P = :f_on_gpu ∈ scenario ? gdev : identity,
+    batchsize = dataloader.batchsize,
+    partial = dataloader.partial
+    )
+    xM, xP, y_o, y_unc, i_sites = dataloader.data
+    xM_dev = gdev_M(xM)
+    xP_dev, y_o_dev, y_unc_dev = (gdev_P(xP), gdev_P(y_o), gdev_P(y_unc)) 
+    train_loader_dev = MLUtils.DataLoader((xM_dev, xP_dev, y_o_dev, y_unc_dev, i_sites);
+        batchsize, partial)
+    return(train_loader_dev)
+end
+
 # function get_hybridproblem_train_dataloader(prob::AbstractHybridProblem; scenario = ())
 #     rng::AbstractRNG = Random.default_rng()
 #     get_hybridproblem_train_dataloader(rng, prob; scenario)

dev/r1_density.pdf

@@ -136,15 +136,49 @@ end
 #     (; n_covar, n_batch, n_θM, n_θP)
 # end
 
+# function HVI.get_hybridproblem_PBmodel(prob::DoubleMMCase; scenario::NTuple = (),
+#     gdev = :f_on_gpu ∈ scenario ? gpu_device() : identity, 
+#     )
+#     #fsite = (θ, x_site) -> f_doubleMM(θ)  # omit x_site drivers
+#     par_templates = get_hybridproblem_par_templates(prob; scenario)
+#     intθ, θFix = setup_PBMpar_interpreter(par_templates.θP, par_templates.θM, θall)
+#     let θFix = gdev(θFix), intθ = get_concrete(intθ)
+#         function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, xP)
+#             pred_sites = applyf(f_doubleMM, θMs, θP, θFix, xP, intθ)
+#             pred_global = eltype(pred_sites)[]
+#             return pred_global, pred_sites
+#         end
+#     end
+# end
+
 function HVI.get_hybridproblem_PBmodel(prob::DoubleMMCase; scenario::NTuple = (),
+    use_all_sites = false,
     gdev = :f_on_gpu ∈ scenario ? gpu_device() : identity, 
     )
+    n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario)
+    n_site_batch = use_all_sites ? n_site : n_batch
     #fsite = (θ, x_site) -> f_doubleMM(θ)  # omit x_site drivers
     par_templates = get_hybridproblem_par_templates(prob; scenario)
-    intθ, θFix = setup_PBMpar_interpreter(par_templates.θP, par_templates.θM, θall)
-    let θFix = gdev(θFix), intθ = get_concrete(intθ)
-        function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, x)
-            pred_sites = applyf(f_doubleMM, θMs, θP, θFix, x, intθ)
+    intθ1, θFix1 = setup_PBMpar_interpreter(par_templates.θP, par_templates.θM, θall)
+    θFix = repeat(θFix1', n_site_batch)
+    intθ = get_concrete(ComponentArrayInterpreter((n_site_batch,), intθ1))
+    isP = repeat(axes(par_templates.θP,1)', n_site_batch)  
+    let θFix = θFix, θFix_dev = gdev(θFix), intθ = get_concrete(intθ), isP=isP, n_site_batch=n_site_batch
+        function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, xP)
+            @assert length(xP) == n_site_batch
+            @assert size(θMs,2) == n_site_batch
+            # convert vector of tuples to tuple of matricesByRows
+            # need to supply xP as vectorOfTuples to work with DataLoader
+            # k = first(keys(xP[1]))
+            xPM = (; zip(keys(xP[1]), map(keys(xP[1])) do k
+                #stack(map(r -> r[k], xP))' 
+                stack(map(r -> r[k], xP); dims = 1)
+            end)...)
+            #xPM = map(transpose, xPM1)
+            # make sure the same order of columns as in intθ
+            θFixd = (θP isa GPUArraysCore.AbstractGPUVector) ? θFix_dev : θFix
+            θ = hcat(CA.getdata(θP[isP]), CA.getdata(θMs)', θFixd)
+            pred_sites = f_doubleMM(θ, xPM, intθ)'
             pred_global = eltype(pred_sites)[]
             return pred_global, pred_sites
         end
@@ -157,6 +191,7 @@ function HVI.get_hybridproblem_PBmodel(prob::DoubleMMCase; scenario::NTuple = ()
     end
 end
 
+
 function HVI.get_hybridproblem_neg_logden_obs(::DoubleMMCase; scenario::NTuple = ())
     neg_logden_indep_normal
 end
@@ -173,25 +208,28 @@ const xP_S2 = Float32[1.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0]
 # const xP_S2 = Float32[1.0, 2.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0]
 
 HVI.get_hybridproblem_n_covar(prob::DoubleMMCase; scenario) = 5
-function HVI.get_hybridproblem_n_site(prob::DoubleMMCase; scenario) 
+function HVI.get_hybridproblem_n_site_and_batch(prob::DoubleMMCase; scenario) 
+    n_batch = 20
+    n_site = 800
     if (:few_sites ∈ scenario)
-         return(100) 
+         n_site = 100
     elseif (:sites20 ∈ scenario)
-        return(20) 
+            n_site = 20
     end
-    800
+    (n_site, n_batch)
 end
 
 function HVI.get_hybridproblem_train_dataloader(prob::DoubleMMCase; scenario = (), 
-    n_batch, rng::AbstractRNG = StableRNG(111), kwargs...
+    rng::AbstractRNG = StableRNG(111), kwargs...
     )
+    n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario)
     construct_dataloader_from_synthetic(rng, prob; scenario, n_batch, kwargs...)
 end
 
 function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase;
         scenario = ())
     n_covar_pc = 2
-    n_site = get_hybridproblem_n_site(prob; scenario)
+    n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario)
     n_covar = get_hybridproblem_n_covar(prob; scenario)
     n_θM = length(θM)
     FloatType = get_hybridproblem_float_type(prob; scenario)
@@ -201,7 +239,7 @@ function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase;
     int_θMs_sites = ComponentArrayInterpreter(θM, (n_site,))
     # normalize to be distributed around the prescribed true values
     θMs_true = int_θMs_sites(scale_centered_at(θMs_true0, θM, FloatType(0.1)))
-    f = get_hybridproblem_PBmodel(prob; scenario, gdev=identity)
+    f = get_hybridproblem_PBmodel(prob; scenario, gdev=identity, use_all_sites = true)
     xP = fill((; S1 = xP_S1, S2 = xP_S2), n_site)
     θP = par_templates.θP
     y_global_true, y_true = f(θP, θMs_true, xP)