provide HVIApproximation type to sampling

to allow different implementations

Author: bgctw

src/HVIApproximation.jl

@@ -0,0 +1,14 @@
+"""
+    AbstractHVIApproximation
+
+Provides a type hierarchy to distinguish different forms and
+parameterizations of posterior approximations.    
+"""
+abstract type AbstractHVIApproximation end
+
+abstract type AbstractMeanHVIApproximation <: AbstractHVIApproximation end
+
+struct MeanHVIApproximation <: AbstractMeanHVIApproximation end
+struct MeanHVIApproximationMat <: AbstractMeanHVIApproximation end
+
+

src/HybridProblem.jl

@@ -39,6 +39,7 @@ struct HybridProblem <: AbstractHybridProblem
     n_site::Int
     n_batch::Int
     pbm_covars::NTuple{_N, Symbol} where _N
+    approx::AbstractHVIApproximation
     #inner constructor to constrain the types
     function HybridProblem(
             θM::CA.ComponentVector,
@@ -57,10 +58,11 @@ struct HybridProblem <: AbstractHybridProblem
             n_batch::Int,
             cor_ends::NamedTuple = (P = [length(ϕq[Val(:μP)])], M = [length(θM)]),
             pbm_covars::NTuple{N,Symbol} = (),
+            approx::AbstractHVIApproximation = MeanHVIApproximation()
     ) where N
         new(
             θM, f_batch, g, ϕg, ϕq, priors, py, transM, transP, cor_ends, 
-            train_dataloader, n_covar, n_site, n_batch, pbm_covars)
+            train_dataloader, n_covar, n_site, n_batch, pbm_covars, approx)
     end
 end
 
@@ -96,7 +98,10 @@ end
 Gather all information from another `AbstractHybridProblem` with possible
 updating of some of the entries.
 """
-function HybridProblem(prob::AbstractHybridProblem; scenario = (),
+function HybridProblem(prob::AbstractHybridProblem; scenario = Val(()), kwargs...)
+    update_hybridProblem(prob; scenario, kwargs...)
+end
+function update_hybridProblem(prob::AbstractHybridProblem; scenario, 
     θM = get_hybridproblem_par_templates(prob; scenario).θM,
     g = get_hybridproblem_MLapplicator(prob; scenario)[1],
     ϕg = get_hybridproblem_MLapplicator(prob; scenario)[2],
@@ -114,6 +119,7 @@ function HybridProblem(prob::AbstractHybridProblem; scenario = (),
     ϕq = get_hybridproblem_ϕq(prob; scenario),
     θP = nothing,
     ϕunc = nothing,
+    approx::AbstractHVIApproximation = MeanHVIApproximation(),
     )
     cor_ends_new = if !isnothing(cor_ends)
         # if new cor_ends was specified then re-initialize the ρsP and ρsM in ϕq
@@ -130,9 +136,14 @@ function HybridProblem(prob::AbstractHybridProblem; scenario = (),
         ϕq = CA.ComponentVector(ϕq; ϕunc...)
     end
     HybridProblem(θM, ϕq, g, ϕg, f_batch, priors, py, transM, transP, train_dataloader,
-        n_covar, n_site, n_batch, cor_ends_new, pbm_covars)
+        n_covar, n_site, n_batch, cor_ends_new, pbm_covars, approx)
+end
+
+function HybridProblem(prob::HybridProblem; kwargs... )
+    update_hybridProblem(prob; scenario = Val(()), kwargs..., approx = prob.approx)
 end
 
+
 # """
 #     update(prob::HybridProblem; ...)
 

src/HybridSolver.jl

@@ -187,6 +187,7 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS
     θmean_quant=0.0,
     is_inferred::Val{is_infer} = Val(false),
     is_omit_priors::Val{omit_priors} = Val(false),
+    approx = prob.approx,
     kwargs...
 ) where {scen, is_infer, omit_priors}
     pt = get_hybridproblem_par_templates(prob; scenario)
@@ -240,7 +241,7 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS
         g_dev, transP, transMs, f_dev, py;
         solver.n_MC, solver.n_MC_cap, cor_ends, priors_θP_mean, priors_θMs_mean, 
         cdev=infer_cdev(gdevs), pbm_covars, pt.θP, int_ϕq, int_ϕg_ϕq, priorsP, priorsM,
-        is_omit_priors, zero_prior_logdensity,
+        is_omit_priors, zero_prior_logdensity, approx,
         )
     # test loss function once
     # tmp = first(train_loader_dev)
@@ -294,7 +295,7 @@ function get_loss_elbo(g, transP, transMs, f, py;
     cor_ends, priors_θP_mean, priors_θMs_mean, cdev, pbm_covars, θP,
     int_ϕq, int_ϕg_ϕq,
     priorsP, priorsM, floss_penalty = zero_penalty_loss,
-    is_omit_priors, zero_prior_logdensity,
+    is_omit_priors, zero_prior_logdensity, approx,
 )
     let g = g, transP = transP, transMs = transMs, f = f, py = py, 
         n_MC = n_MC, n_MC_cap = n_MC_cap, n_MC_mean = n_MC_mean,
@@ -305,7 +306,8 @@ function get_loss_elbo(g, transP, transMs, f, py;
         trans_mP=StackedArray(transP, n_MC_mean), 
         trans_mMs=StackedArray(transMs.stacked, n_MC_mean),
         priorsP=priorsP, priorsM=priorsM, floss_penalty=floss_penalty,
-        is_omit_priors = is_omit_priors, zero_prior_logdensity = zero_prior_logdensity
+        is_omit_priors = is_omit_priors, zero_prior_logdensity = zero_prior_logdensity,
+        approx = approx
 
         function loss_elbo(ϕ, rng, xM, xP, y_o, y_unc, i_sites; is_testmode)
             #ϕc = int_ϕg_ϕq(ϕ)
@@ -315,7 +317,7 @@ function get_loss_elbo(g, transP, transMs, f, py;
                 n_MC, n_MC_cap, n_MC_mean, cor_ends, priors_θP_mean, priors_θMs_mean,
                 cdev, pbm_covar_indices, transP, transMs, trans_mP, trans_mMs,
                 priorsP, priorsM, floss_penalty, #ϕg = ϕc.ϕg, ϕq = ϕc.ϕq,
-                is_testmode, is_omit_priors, zero_prior_logdensity,
+                is_testmode, is_omit_priors, zero_prior_logdensity, approx,
             )
         end
     end

src/HybridVariationalInference.jl

@@ -40,6 +40,10 @@ VERSION >= v"1.11.0-DEV.469" && eval(Meta.parse("public Exp"))
 VERSION >= v"1.11.0-DEV.469" && eval(Meta.parse("public Logistic")) 
 include("bijectors_utils.jl")
 
+export AbstractHVIApproximation, AbstractMeanHVIApproximation
+export MeanHVIApproximation, MeanHVIApproximationMat
+include("HVIApproximation.jl")
+
 export AbstractComponentArrayInterpreter, ComponentArrayInterpreter,
        StaticComponentArrayInterpreter
 export flatten1, get_concrete, get_positions, stack_ca_int, compose_interpreters
@@ -107,6 +111,7 @@ include("cholesky.jl")
 
 export neg_elbo_gtf, sample_posterior, predict_hvi, zero_penalty_loss
 include("elbo.jl")
+include("elbo2.jl")
 
 export init_hybrid_params, init_hybrid_ϕunc
 include("init_hybrid_params.jl")

src/cholesky.jl

@@ -253,7 +253,7 @@ end
 """
     transformU_block_cholesky1(v::AbstractVector, cor_ends)
 
-Transform a parameterization v of a blockdiagonal of upper triangular matrices
+Transform a parameterization, `v`, of a blockdiagonal of upper triangular matrices
 into the this matrix.
 `cor_ends` is an AbstractVector of Integers specifying the last column of each block. 
 E.g. For a matrix with a 3x3, a 2x2, and another single-entry block, 
@@ -262,6 +262,26 @@ the blocks start at columns (3,5,6). It defaults to a single entire block.
 An correlation parameterization can parameterize a block of a single parameter, 
 or an empty parameter block. To indicate the empty block, provide `cor_ends == [0]`.
 """
+function transformU_blocks_cholesky1(
+    v::AbstractVector{T}, cor_ends::AbstractVector{TI}=Int[]) where {T,TI<:Integer}
+    if length(cor_ends) <= 1 # if there is only one block, return it 
+        # (cor_ends == [0]) no parameters at all (and also no correclation)
+        # (cor_ends == [1]) a single parameter (and also no correclation)
+        create_empty = (cor_ends == [0])
+        return [transformU_cholesky1(v; create_empty)], 1:0
+    end
+    cor_counts = get_cor_counts(cor_ends) # number of correlation parameters
+    #@show cor_counts
+    ranges = ChainRulesCore.@ignore_derivatives (
+        begin
+            cor_start = (i == 1 ? one(TI) : cor_counts[i-1] + one(TI))
+            cor_start:cor_counts[i]
+        end for i in 1:length(cor_counts)
+    )
+    #@show collect(ranges)
+    [transformU_cholesky1(v[r]) for r in ranges], ranges
+end
+
 function transformU_block_cholesky1(
     v::AbstractVector{T}, cor_ends::AbstractVector{TI}=Int[]) where {T,TI<:Integer}
     # (cor_ends == [0]) no parameters at all (and also no correclation)
@@ -284,6 +304,7 @@ function transformU_block_cholesky1(
     return (U)
 end
 
+
 function _create_blockdiag(::AbstractArray{T}, blocks::AbstractArray) where {T}
     BlockDiagonal(blocks)
 end

src/elbo.jl

@@ -62,9 +62,10 @@ function neg_elbo_gtf_components(rng, ϕ::AbstractVector{FT}, g, f, py,
     is_testmode,
     is_omit_priors,
     zero_prior_logdensity,
+    approx::AbstractHVIApproximation,
 ) where {FT}
     n_MCr = isempty(priors_θP_mean) ? n_MC : max(n_MC, n_MC_mean)
-    ζsP, ζsMs, σ = generate_ζ(rng, g, ϕ, xM; n_MC=n_MCr, cor_ends, pbm_covar_indices,
+    ζsP, ζsMs, σ = generate_ζ(approx, rng, g, ϕ, xM; n_MC=n_MCr, cor_ends, pbm_covar_indices,
         int_ϕq, int_ϕg_ϕq, is_testmode)
     ζsP_cpu = cdev(ζsP) # fetch to CPU, because for <1000 sites (n_batch) this is faster
     ζsMs_cpu = cdev(ζsMs) # fetch to CPU, because for <1000 sites (n_batch) this is faster
@@ -382,6 +383,7 @@ function sample_posterior(rng, prob::AbstractHybridProblem, xM::AbstractMatrix;
     scenario=Val(()),
     n_sample_pred=200,
     gdevs = get_gdev_MP(scenario),
+    approx = nothing,
     kwargs...
 )
     n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario)
@@ -400,9 +402,13 @@ function sample_posterior(rng, prob::AbstractHybridProblem, xM::AbstractMatrix;
     int_ϕq = interpreters.ϕq
     transMs = StackedArray(transM, n_batch)        
     g_dev, ϕ_dev = gdevs.gdev_M(g), gdevs.gdev_M(ϕ)
+    if isnothing(approx) 
+        approx = prob.approx  # assuming has field approx, e.g. if its a HybridProblem
+    end
     (; θsP, θsMs, entropy_ζ) = sample_posterior(rng, g_dev, ϕ_dev, xM;
         int_ϕg_ϕq, int_ϕq, transP, transM, 
-        n_sample_pred, cdev=infer_cdev(gdevs), cor_ends, pbm_covar_indices, kwargs...)
+        n_sample_pred, cdev=infer_cdev(gdevs), cor_ends, pbm_covar_indices, approx,
+        kwargs...)
     θsPc = ComponentArrayInterpreter(par_templates.θP, (n_sample_pred,))(θsP)
     θsMsc = ComponentArrayInterpreter((n_site,), par_templates.θM, (n_sample_pred,))(θsMs)
     (; θsP=θsPc, θsMs=θsMsc, entropy_ζ)
@@ -418,8 +424,9 @@ function sample_posterior(rng, g, ϕ::AbstractVector, xM::AbstractMatrix;
     pbm_covar_indices,
     is_inferred::Val{is_infer} = Val(false),
     is_testmode,
+    approx::AbstractHVIApproximation,
 ) where is_infer
-    ζsP_gpu, ζsMs_gpu, σ = generate_ζ(rng, g, CA.getdata(ϕ), CA.getdata(xM);
+    ζsP_gpu, ζsMs_gpu, σ = generate_ζ(approx, rng, g, CA.getdata(ϕ), CA.getdata(xM);
         int_ϕg_ϕq, int_ϕq,
         n_MC=n_sample_pred, cor_ends, pbm_covar_indices, is_testmode)
     ζsP = cdev(ζsP_gpu)
@@ -434,6 +441,7 @@ function sample_posterior(rng, g, ϕ::AbstractVector, xM::AbstractMatrix;
     (; θsP, θsMs, entropy_ζ)
 end
 
+
 """
 Generate samples of (inv-transformed) model parameters, ζ, 
 and the vector of standard deviations, σ, i.e. the diagonal of the cholesky-factor.
@@ -445,7 +453,8 @@ model.
 The output shape of size `(n_site x n_par x n_MC)` is tailored to iterating
 each MC sample and then transforming each parameter on block across sites.
 """
-function generate_ζ(rng, g, ϕ::AbstractVector{FT}, xM::MT;
+function generate_ζ(approx::AbstractMeanHVIApproximation, rng::AbstractRNG, 
+    g, ϕ::AbstractVector{FT}, xM::MT;
     int_ϕg_ϕq::AbstractComponentArrayInterpreter,
     int_ϕq::AbstractComponentArrayInterpreter,
     n_MC=3, cor_ends, pbm_covar_indices,
@@ -462,7 +471,7 @@ function generate_ζ(rng, g, ϕ::AbstractVector{FT}, xM::MT;
     # TODO replace pbm_covar_indices by ComponentArray? dimensions to be type-inferred?
     xMP0 = _append_each_covars(xM, CA.getdata(μ_ζP), pbm_covar_indices)
     μ_ζMs0 = g(xMP0, ϕg; is_testmode)
-    ζP_resids, ζMs_parfirst_resids, σ = sample_ζresid_norm(rng, μ_ζP, μ_ζMs0, ϕq; n_MC, cor_ends, int_ϕq)
+    ζP_resids, ζMs_parfirst_resids, σ = sample_ζresid_norm(approx, rng, μ_ζP, μ_ζMs0, ϕq; n_MC, cor_ends, int_ϕq)
     if pbm_covar_indices isa SA.SVector{0}
         # do not need to predict again but just add the residuals to μ_ζP and μ_ζMs
         #ζsP = μ_ζP .+ ζP_resids  # n_par x n_MC # .+ on empty view does not work
@@ -559,8 +568,9 @@ ML-model predcitions of size `(n_θM, n_site)`.
 * `int_ϕq`: Interpret vector as ComponentVector with components
    ρsP, ρsM, logσ2_ζP, coef_logσ2_ζMs(intercept + slope), 
 """
-function sample_ζresid_norm(rng::Random.AbstractRNG, ζP::AbstractVector, ζMs::AbstractMatrix,
-    args...; n_MC, cor_ends, int_ϕq)
+function sample_ζresid_norm(approx::AbstractHVIApproximation, rng::Random.AbstractRNG, 
+    ζP::AbstractVector, ζMs::AbstractMatrix, args...; 
+    n_MC, cor_ends, int_ϕq)
     n_θP, n_θMs = length(ζP), length(ζMs)
     # intm_PMs_parfirst = !isnothing(intm_PMs_parfirst) ? intm_PMs_parfirst : begin
     #     n_θM, n_site_batch = size(ζMs)
@@ -569,14 +579,14 @@ function sample_ζresid_norm(rng::Random.AbstractRNG, ζP::AbstractVector, ζMs:
     # end
     #urandn = _create_randn(rng, CA.getdata(ζP), n_MC, n_θP + n_θMs)
     urandn = _create_randn(rng, CA.getdata(ζP), n_θP + n_θMs, n_MC)
-    sample_ζresid_norm(urandn, CA.getdata(ζP), CA.getdata(ζMs), args...;
+    sample_ζresid_norm(approx, urandn, CA.getdata(ζP), CA.getdata(ζMs), args...;
         cor_ends, 
         int_ϕq=get_concrete(int_ϕq)
     )
 end
 
-function sample_ζresid_norm(urandn::AbstractMatrix, ζP::TP, ζMs::TM,
-    ϕq::AbstractVector;
+function sample_ζresid_norm(approx::MeanHVIApproximationMat, urandn::AbstractMatrix, 
+    ζP::TP, ζMs::TM, ϕq::AbstractVector;
     int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)),
     cor_ends
 ) where {T,TP<:AbstractVector{T},TM<:AbstractMatrix{T}}

src/elbo2.jl

@@ -0,0 +1,45 @@
+function sample_ζresid_norm(app::MeanHVIApproximation, z::AbstractMatrix, ζP::TP, ζMs::TM,
+    ϕq::AbstractVector;
+    int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)),
+    cor_ends
+) where {T,TP<:AbstractVector{T},TM<:AbstractMatrix{T}}
+    ϕuncc = int_ϕq(CA.getdata(ϕq))
+    n_θP, n_θMs, (n_θM, n_batch) = length(ζP), length(ζMs), size(ζMs)
+    # do not create a UpperTriangular Matrix of an AbstractGÜUArray in transformU_cholesky1
+    ρsP = isempty(ϕuncc[Val(:ρsP)]) ? similar(ϕuncc[Val(:ρsP)]) : ϕuncc[Val(:ρsP)] # required by zygote
+    UPs, rangesP = transformU_blocks_cholesky1(ρsP, cor_ends.P) 
+    UP = transformU_block_cholesky1(ρsP, cor_ends.P)
+    ρsM, rangesM = isempty(ϕuncc[Val(:ρsM)]) ? similar(ϕuncc[Val(:ρsM)]) : ϕuncc[Val(:ρsM)] # required by zygote
+    # cholesky factor of the correlation: diag(UM' * UM) .== 1
+    # coefficients ρsM can be larger than 1, still yielding correlations <1 in UM' * UM
+    UMs = transformU_block_cholesky1(ρsM, cor_ends.M)
+    cf = ϕuncc[Val(:coef_logσ2_ζMs)]
+    logσ2_logMs = vec(cf[1, :] .+ cf[2, :] .* ζMs)
+    logσ2_ζP = vec(CA.getdata(ϕuncc[Val(:logσ2_ζP)]))
+    # CUDA cannot multiply BlockDiagonal * Diagonal, construct already those blocks
+    σMs = reshape(exp.(logσ2_logMs ./ 2), n_θM, :)
+    σP = exp.(logσ2_ζP ./ 2)
+    # create random numbers from U diag(σ) z =  (σ .* z)
+    
+
+    # # BlockDiagonal does work with CUDA, but not with combination of Zygote and CUDA
+    # # need to construct full matrix for CUDA
+    # Uσ, diagUσ = _compute_choleskyfactor(UP, UM, σP, σMs, n_batch) # inferred only BlockDiagonal
+    # #diagUσ = diag(Uσ)::typeof(σP)   # elements of the diagonal: standard deviations
+    # n_MC = size(urandn, 2) # TODO transform urandn
+    # # is this multiplication efficient if Uσ is not concrete but only sumtype BlockDiagonal?
+    # ζ_resids_parfirst = (Uσ' * urandn) #::typeof(urandn) # n_par x n_MC
+    # #ζ_resids_parfirst = urandn' * Uσ # n_MC x n_par
+    # # need to handle empty(ζP) explicitly, otherwise Zygote tries to take gradient
+    # ζP_resids = isempty(ζP) ? ζ_resids_parfirst[1:0, :] : ζ_resids_parfirst[1:n_θP, :]
+    # ζMs_parfirst_resids = reshape(ζ_resids_parfirst[(n_θP+1):end, :], n_θM, n_batch, n_MC)
+
+
+    ζP_resids, ζMs_parfirst_resids, diagUσ
+    # #map(std, eachcol(ζ_resids_parfirst[:, 3:8]))
+    # ζ_resid = transpose_mPMs_sitefirst(ζ_resids_parfirst; intm_PMs_parfirst)
+    # #map(std, eachcol(ζ_resid[:, 3:8])) # all ~ 0.1 in sample_ζresid_norm cpu
+    # #map(std, eachcol(ζ_resid[:, 2 + n_batch .+ (-1:5)])) # all ~ 100, except first two
+    # # returns AbstractGPUuArrays to either continue on GPU or need to transfer to CPU
+    # ζ_resid, diagUσ
+end

test/test_HybridProblem.jl

@@ -23,7 +23,8 @@ using Functors
 
 cdev = cpu_device()
 
-#scenario = Val((:default,))
+#scenario = Val((:default, ))
+#scenario = Val((:MeanHVIApproxMat,))
 #scenario = Val((:covarK2,))
 #scen = CP._val_value(scenario)
 
@@ -89,10 +90,15 @@ function construct_problem(; scenario::Val{scen}) where scen
         xPvec=xP[:,1])
     ϕunc0 = init_hybrid_ϕunc(cor_ends, zero(FT)) 
     ϕq = CP.update_μP_by_θP(ϕunc0, θP, transP)
+    approx = if (:MeanHVIApproxMat ∈ scen) 
+        MeanHVIApproximationMat() 
+    else
+        MeanHVIApproximation()
+    end
     HybridProblem(θM, ϕq, g_chain_scaled, ϕg0, 
         f_batch, priors_dict, py,
         transM, transP, train_dataloader, n_covar, n_site, n_batch, 
-        cor_ends, pbm_covars,
+        cor_ends, pbm_covars, approx,
         #ϕunc0, 
         )
 end 
@@ -191,14 +197,16 @@ test_without_flux = (scenario) -> begin
     end
 end
 
+#test_without_flux(Val((:MeanHVIApproximation,)))  # not used in loss_gf
 test_without_flux(Val((:default,)))
 test_without_flux(Val((:covarK2,)))
 
 import CUDA, cuDNN
 using GPUArraysCore
 import Flux
 
-gdev = gpu_device()
+CUDA.device!(2)   # TODO remove after GPU 0 becomes available again
+gdev = gpu_device() 
 #methods(CP.vec2uutri)
 
 test_with_flux = (scenario) -> begin
@@ -373,6 +381,7 @@ test_with_flux_gpu = (scenario) -> begin
     end # if gdev isa MLDataDevices.AbstractGPUDevice 
 end # test_with flux
 
+test_with_flux_gpu(Val((:MeanHVIApproxMat,)))
 test_with_flux_gpu(Val((:default,)))
 test_with_flux_gpu(Val((:covarK2,)))
 test_with_flux_gpu(Val((:default,:useSitePBM)))