unify prior parameterization and computation

Author: bgctw

src/HybridSolver.jl

@@ -13,9 +13,9 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve
     ad_backend_loss = AutoZygote(),
     epochs,
     is_omitting_NaNbatches = false,
-    is_omit_priors::Val{is_omit_prior} = Val(false),
+    is_omit_priors::Val{omit_priors} = Val(false),
     kwargs...
-) where {is_infer, is_omit_prior}
+) where {is_infer, omit_priors}
     gdevs = isnothing(gdevs) ? get_gdev_MP(scenario) : gdevs
     pt = get_hybridproblem_par_templates(prob; scenario)
     g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario)
@@ -63,13 +63,17 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve
     priors = get_hybridproblem_priors(prob; scenario)
     priorsP = Tuple(priors[k] for k in keys(pt.θP))
     priorsM = Tuple(priors[k] for k in keys(pt.θM))
+    zero_prior_logdensity = omit_priors ? 0f0 : get_zero_prior_logdensity(
+        priorsP, priorsM, pt.θP, pt.θM)     
     #intP = ComponentArrayInterpreter(pt.θP)
     loss_gf = get_loss_gf(g_dev, transM, transP, f_dev,  py, intϕ;
         n_site_batch=n_batch, 
-        cdev=infer_cdev(gdevs), pbm_covars, priorsP, priorsM, is_omit_priors,)
+        cdev=infer_cdev(gdevs), pbm_covars, 
+        priorsP, priorsM, is_omit_priors, zero_prior_logdensity,)
     loss_gf_test = get_loss_gf(g_dev, transM, transP, ftest_dev,  py, intϕ;
         n_site_batch=n_site_test,
-        cdev=infer_cdev(gdevs), pbm_covars, priorsP, priorsM, is_omit_priors,)
+        cdev=infer_cdev(gdevs), pbm_covars, 
+        priorsP, priorsM, is_omit_priors, zero_prior_logdensity,)
     # call loss function once
     l1 = is_infer ? 
         Test.@inferred(loss_gf(ϕ0_dev, first(train_loader_dev)...; is_testmode=true))[1] : 

src/elbo.jl

@@ -134,9 +134,9 @@ function neg_elbo_ζtf(ζsP, ζsMs, σ, f, py, xP, y_ob, y_unc;
     transMs=StackedArray(transM, size(ζsMs, 2)),
     priorsP, priorsM,
     floss_penalty, ϕg, ϕq,
-    is_omit_priors::Val{omit_priors},
+    is_omit_priors::Val,
     zero_prior_logdensity,
-) where omit_priors
+) 
     n_MC = size(ζsP,2)
     f_sample = (ζP, ζMs) -> begin    
             θP, θMs, logjac_i = transform_and_logjac_ζ(ζP, ζMs; transP, transMs)
@@ -150,16 +150,8 @@ function neg_elbo_ζtf(ζsP, ζsMs, σ, f, py, xP, y_ob, y_unc;
             # @descend_code_warntype f(θP, θMs, xP)
             nLy_i = py(y_ob, y_pred_i, y_unc)
             loss_penalty_i = convert(eltype(nLy_i),floss_penalty(y_pred_i, θMs, θP, ϕg, ϕq))
-            neg_log_prior_i = if omit_priors
-                zero_prior_logdensity
-            elseif (θP isa AbstractGPUArray) || (θMs isa AbstractGPUArray)
-                @warn("neg_elbo_ζtf: Cannot apply priors to gpu array. Piors are omitted. " *
-                "either compute PBM on CPU or omit priors.")
-                zero_prior_logdensity
-            else
-                compute_priors_logdensity(priorsP, priorsM, θP, θMs, zero_prior_logdensity)
-            end
-
+            neg_log_prior_i = compute_priors_logdensity(priorsP, priorsM, θP, θMs,
+                is_omit_priors, zero_prior_logdensity)
             # make sure names to not match outer, otherwise Box type instability
             (nLy_i, neg_log_prior_i, -logjac_i, loss_penalty_i)
             #(nLy_i, 0.0, 0.0, 0.0)
@@ -207,6 +199,19 @@ function neg_elbo_ζtf(ζsP, ζsMs, σ, f, py, xP, y_ob, y_unc;
     (;nLjoint, entropy_ζ, loss_penalty, nLy, neg_log_prior, neg_log_jac)
 end
 
+function compute_priors_logdensity(priorsP, priorsM, θP, θMs,
+        ::Val{omit_priors}, zero_prior_logdensity) where {omit_priors}
+    if omit_priors
+        zero_prior_logdensity
+    elseif (θP isa AbstractGPUArray) || (θMs isa AbstractGPUArray)
+        @warn("neg_elbo_ζtf: Cannot apply priors to gpu array. Piors are omitted. "*
+        "either compute PBM on CPU or omit priors.")
+        zero_prior_logdensity
+    else
+        compute_priors_logdensity(priorsP, priorsM, θP, θMs, zero_prior_logdensity)
+    end
+end
+
 function compute_priors_logdensity(priorsP, priorsM, θP, θMs, zero_prior_logdensity)
     logpdf_t = (prior, θ) -> logpdf(prior, θ)::eltype(θP)
     function logpdf_tv_sum(prior, θ::AbstractVector{T}) where T 

src/gf.jl

@@ -231,16 +231,20 @@ function get_loss_gf(g, transM, transP, f, py,
         intϕ(1:length(intϕ)).ϕP);
     cdev=cpu_device(),
     pbm_covars, n_site_batch, 
-    priorsP, priorsM, floss_penalty = zero_penalty_loss,
-    is_omit_priors::Val{is_omit_prior} = Val(false),
-    kwargs...) where is_omit_prior
+    floss_penalty = zero_penalty_loss,
+    priorsP, priorsM, 
+    is_omit_priors::Val = Val(false),
+    zero_prior_logdensity,
+    kwargs...) 
 
     let g = g, transM = transM, transP = transP, f = f, 
         intϕ = get_concrete(intϕ),
         transMs = StackedArray(transM, n_site_batch),
         cdev = cdev,
         pbm_covar_indices = CA.getdata(intP(1:length(intP))[pbm_covars]),
-        priorsP = priorsP, priorsM = priorsM, floss_penalty = floss_penalty,
+        zero_prior_logdensity = zero_prior_logdensity, is_omit_priors = is_omit_priors,
+        priorsP = priorsP, priorsM = priorsM, 
+        floss_penalty = floss_penalty,
         cpu_dev = cpu_device() # real cpu, different form infer_cdev(gdevs) that maybe idenetity
         #, intP = get_concrete(intP)
         #inv_transP = inverse(transP), kwargs = kwargs
@@ -273,29 +277,8 @@ function get_loss_gf(g, transM, transP, f, py,
             logpdf_tv = (prior, θ::AbstractVector) -> begin
                 map(Base.Fix1(logpdf, prior), θ)::Vector{eltype(θP_pred)}
             end
-            #Main.@infiltrate_main
-            #Maybe: move priors to GPU, for now need to move θ to cpu
-            # currently does not work on gpu, moving to dpu has problems with gradient
-            #    need to specify is_omit_priors if PBM is on GPU
-            neg_log_prior = if is_omit_prior
-                    zero(nLy)
-            else
-                nLP = if isempty(θP_pred) 
-                    zero(nLy)
-                else
-                    θP_pred_cpu = CA.getdata(θP_pred)
-                    -sum(logpdf_t.(priorsP, θP_pred_cpu))
-                end
-                θMs_pred_cpu = CA.getdata(θMs_pred)
-                nLM = -sum(map((priorMi, θMi) -> sum(
-                    logpdf_tv(priorMi, θMi)), priorsM, eachcol(θMs_pred_cpu)))
-                nLP + nLM
-            end
-            # neg_log_prior = is_omit_priors ? zero(nLy) :
-            #     (isempty() ? zero(nLy) : ) +
-            #     -sum(map((priorMi, θMi) -> sum(
-            #         logpdf_tv(priorMi, θMi)), priorsM, eachcol(θMs_pred_cpu))) 
-            #neg_log_prior = min(sqrt(floatmax(neg_log_prior0)), neg_log_prior0)                
+            neg_log_prior = compute_priors_logdensity(priorsP, priorsM, θP_pred, θMs_pred,
+                is_omit_priors, zero_prior_logdensity)
             if !isfinite(neg_log_prior)
                 @info "loss_gf: encountered non-finite prior density"
                 @show θP_pred, θMs_pred, ϕc.ϕP

test/Project.toml

@@ -27,5 +27,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"

test/test_HybridProblem.jl

@@ -1,6 +1,7 @@
 using Test
 using HybridVariationalInference
 using HybridVariationalInference: HybridVariationalInference as CP
+using UnPack
 using StableRNGs
 using Random
 using Statistics
@@ -19,6 +20,7 @@ using Suppressor
 
 using Functors
 
+
 cdev = cpu_device()
 
 #scenario = Val((:default,))
@@ -106,9 +108,7 @@ end
     function test_f_doubleMM(θ::AbstractVector{ET}, x; intθ1) where ET
         # extract parameters not depending on order, i.e whether they are in θP or θM
         θc = intθ1(θ)
-        (r0, r1, K1, K2) = map((:r0, :r1, :K1, :K2)) do par
-            CA.getdata(θc[par])::ET
-        end
+        @unpack r0, r1, K1, K2 = θc
         r0 .+ r1 .* x.S1 ./ (K1 .+ x.S1) .* x.S2 ./ (K2 .+ x.S2)
     end    
     y = @inferred test_f_doubleMM(CA.getdata(θ2), xP1; intθ1 = int_θdoubleMM)
@@ -151,14 +151,16 @@ test_without_flux = (scenario) -> begin
         intϕ = ComponentArrayInterpreter(CA.ComponentVector(
             ϕg=1:length(ϕg0), ϕP=par_templates.θP))
         priors = get_hybridproblem_priors(prob; scenario)
-        priorsP = [priors[k] for k in keys(par_templates.θP)]
-        priorsM = [priors[k] for k in keys(par_templates.θM)]
+        priorsP = Tuple(priors[k] for k in keys(par_templates.θP))
+        priorsM = Tuple(priors[k] for k in keys(par_templates.θM))
         # slightly disturb θP_true
         p = p0 = vcat(ϕg0, par_templates.θP .* convert(eltype(ϕg0), 0.8))  
 
         # Pass the site-data for the batches as separate vectors wrapped in a tuple
+        zero_prior_logdensity = CP.get_zero_prior_logdensity(
+            priorsP, priorsM, par_templates.θP, par_templates.θM)     
         loss_gf = get_loss_gf(g, transM, transP, f, py, intϕ; 
-            pbm_covars, n_site_batch = n_batch, priorsP, priorsM,
+            pbm_covars, n_site_batch = n_batch, priorsP, priorsM, zero_prior_logdensity,
             )
         (_xM, _xP, _y_o, _y_unc, _i_sites) = first(train_loader)
         #l1 = loss_gf(p0, _xM, _xP, _y_o, _y_unc, _i_sites; is_testmode = false)
@@ -312,13 +314,13 @@ test_with_flux_gpu = (scenario) -> begin
                 (; y, θsP, θsMs) = predict_hvi(rng, probo; scenario = scenf, n_sample_pred);            
                 # to inspect correlations among θP and θMs construct ComponentVector
                 # TODO redo get_int_PMst_site
-    # get_ca_int_PMs = let
-    #     function get_ca_int_PMs_inner(n_site)
-    #         ComponentArrayInterpreter(CA.ComponentVector(; P = θP,
-    #             Ms = CA.ComponentMatrix(
-    #                 zeros(n_θM, n_site), first(CA.getaxes(θM)), CA.Shaped1DAxis((n_site,)))))
-    #     end
-    # end
+                # get_ca_int_PMs = let
+                #     function get_ca_int_PMs_inner(n_site)
+                #         ComponentArrayInterpreter(CA.ComponentVector(; P = θP,
+                #             Ms = CA.ComponentMatrix(
+                #                 zeros(n_θM, n_site), first(CA.getaxes(θM)), CA.Shaped1DAxis((n_site,)))))
+                #     end
+                # end
                 int_mPMs = stack_ca_int(Val((n_sample_pred,)), get_int_PMst_site(hpints))
                 θs =  int_mPMs(CP.flatten_hybrid_pars(θsP, θsMs))
                 mean_θ = CA.ComponentVector(vec(mean(CA.getdata(θs), dims=1)), last(CA.getaxes(θs)))

test/test_doubleMM.jl

@@ -204,8 +204,8 @@ end
     f2 = create_nsite_applicator(f, n_site)
     py = get_hybridproblem_neg_logden_obs(prob; scenario)
     priors = get_hybridproblem_priors(prob; scenario)
-    priorsP = [priors[k] for k in keys(par_templates.θP)]
-    priorsM = [priors[k] for k in keys(par_templates.θM)]
+    priorsP = Tuple(priors[k] for k in keys(par_templates.θP))
+    priorsM = Tuple(priors[k] for k in keys(par_templates.θM))
 
     intϕ = ComponentArrayInterpreter(CA.ComponentVector(
         ϕg = 1:length(ϕg0), ϕP = par_templates.θP))
@@ -222,10 +222,12 @@ end
     pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario)
 
     #loss_gf = get_loss_gf(g, transM, f,  intϕ; gdev = identity)
+    zero_prior_logdensity = CP.get_zero_prior_logdensity(
+        priorsP, priorsM, par_templates.θP, par_templates.θM)     
     loss_gf = get_loss_gf(g, transM, transP, f,  py, intϕ;
-        pbm_covars, n_site_batch = n_batch, priorsP, priorsM)
+        pbm_covars, n_site_batch = n_batch, priorsP, priorsM, zero_prior_logdensity)
     loss_gf_site = get_loss_gf(g, transM, transP, f2, py, intϕ;
-        pbm_covars, n_site_batch = n_site, priorsP, priorsM)
+        pbm_covars, n_site_batch = n_site, priorsP, priorsM, zero_prior_logdensity)
     nLjoint = @inferred first(loss_gf(p0, first(train_loader)...; is_testmode=true))
     (xM_batch, xP_batch, y_o_batch, y_unc_batch, i_sites_batch) = first(train_loader)
     # @usingany Cthulhu