implement grad_pointloss_sites

src/HybridVariationalInference.jl

@@ -169,6 +169,10 @@ include("elbo_site_grad.jl")
 include("elbo_site_corr.jl")
 #export prepare_elbo_helpers
 
+include("pointloss_site.jl")
+include("pointloss_site_grad.jl")
+
+
 #export get_loss_ran_tr_f
 include("fit_ranef_site.jl")
 

src/PBMApplicator.jl

@@ -57,7 +57,9 @@ function create_nsite_applicator(app::PBMPopulationApplicator, n_MC;
     PBMPopulationApplicator(app.fθpop!, θFixm, θFixm_dc, app.axθcomp, app.ax_xP)        
 end
 
-function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::AbstractMatrix, θsM::AbstractMatrix, xP) where {AXS,AXP}
+# single site, severl MC
+function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, 
+    θsP::AbstractMatrix, θsM::AbstractMatrix, xP::AbstractVector) where {AXS,AXP}
     # function tmpf(θsP, θsM, app) 
     #     @allocated Vcat(CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm))
     # end
@@ -75,8 +77,6 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::Abstrac
     # θc = CA.ComponentMatrix(
     #     Vcat(CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm)),
     #     AXC, CA.FlatAxis())
-    # TODO provide PreallocationTools for app.θFixm
-    #θc = VcatCAs(AXS, CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm))
     θFixm_dual = if eltype(θsP) <: ForwardDiff.Dual
         copyto!(PAT.get_tmp(app.θFixm_dc, CA.getdata(θsP)), app.θFixm)
     else
@@ -91,7 +91,26 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::Abstrac
     pred
 end
 
-function (app::PBMPopulationApplicator)(θP::AbstractMatrix{FT}, θMs_tr::AbstractArray, xP::AbstractMatrix) where FT
+# single site, single MC
+#    causes allocations, better preallocate Matrix(n,1) 
+# function apply_model!(pred::AbstractVector, app::PBMPopulationApplicator{AXS,AXP}, θP::AbstractVector, θM::AbstractVector, xP::AbstractVector) where {AXS,AXP}
+#     local θc, xPc, θFixm_dual
+#     θFixm_dual = if eltype(θP) <: ForwardDiff.Dual
+#         copyto!(PAT.get_tmp(app.θFixm_dc, CA.getdata(θP)), app.θFixm)
+#     else
+#         app.θFixm
+#     end
+#     pred_mat = Base.ReshapedArray(pred, (length(pred), 1), ())
+#     θP_mat = Base.ReshapedArray(pred, (length(θP), 1), ())
+#     θM_mat = Base.ReshapedArray(pred, (length(θM), 1), ())
+#     θc = VcatCAs(AXS, θP_mat, θM_mat, θFixm_dual)
+#     xPc = CA.ComponentVector(CA.getdata(xP), AXP)
+#     app.fθpop!(pred_mat, θc, xPc)  
+#     pred
+# end
+
+# several site, several MC
+function apply_model(app::PBMPopulationApplicator,θP::AbstractMatrix{FT}, θMs_tr::AbstractArray, xP::AbstractMatrix) where FT
     local n_obs, n_site, n_par, n_MC, pred
     n_obs = get_n_obs(xP[:,1])
     n_site, n_par, n_MC = size(θMs_tr)
@@ -102,7 +121,8 @@ function (app::PBMPopulationApplicator)(θP::AbstractMatrix{FT}, θMs_tr::Abstra
     end
     (pred, nothing)
 end
-function (app::PBMPopulationApplicator)(θP::AbstractVector{FT}, θMs_tr::AbstractMatrix, xP::AbstractMatrix) where FT
+# several site, single MC
+function apply_model(app::PBMPopulationApplicator, θP::AbstractVector{FT}, θMs_tr::AbstractMatrix, xP::AbstractMatrix) where FT
     # for a single MC-run
     local n_obs, n_site, pred, θsP, θsM
     n_obs = get_n_obs(xP[:,1])

src/elbo_site.jl

@@ -28,8 +28,8 @@
     nljoint_args_inds = NamedTuple{keys(indiv_args)[r2end]}(values(indiv_args)[r2end])
     @assert size(rnormPM.P) == size(h.ζsP)
     use_dc = h.helpers_sites[1].ζsM isa PAT.DiffCache
-    n_M, n_MC = use_dc ? size(h.helpers_sites[1].ζsM.du) : size(h.helpers_sites[1].ζsM)
-    @assert size(rnormPM.M[1]) == (n_M, n_MC)
+    n_θM, n_MC = use_dc ? size(h.helpers_sites[1].ζsM.du) : size(h.helpers_sites[1].ζsM)
+    @assert size(rnormPM.M[1]) == (n_θM, n_MC)
     ϕqPc = intϕqP(ϕqP) 
     ϕqIc = intϕqI(ϕqI)
     pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices
@@ -48,7 +48,7 @@
     # closure with approx, transM, and kwargs
     function compute_nelboi_z_cl!(hi, rnormM, nljoint_args_indi_tup, ϕm) 
         #nljoint_args_indi = NamedTuple{keys_nljoint_args_indi}(nljoint_args_indi_tup)
-        nljoint_args_indi = add_names(nljoint_args_indi_tup, keys_nljoint_args_indi_val)
+        nljoint_args_indi = name_tuple(nljoint_args_indi_tup, keys_nljoint_args_indi_val)
         compute_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm, 
         ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix) 
     end
@@ -64,7 +64,7 @@
     elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP)
 end
 
-add_names(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup)
+name_tuple(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup)
 
 
 
@@ -76,7 +76,7 @@
     #       2 (Matrix) if g is driven by population covariates - differs across MC
     # on update -> sync corresponding function within grad_neg_elbo_sites
     if hi.ζsM isa PAT.DiffCache
         # Gradiend called twice because of different dimensions need to pick the Dual
         template = eltype(ϕqIc) <: ForwardDiff.Dual ? ϕqIc : θsP
         hi = map_leaves_nt(x -> PAT.get_tmp(x, template), hi)
     end
@@ -232,7 +232,7 @@
     g::AbstractModelApplicator,
     xMP::AbstractMatrix,
     is_testmode::Bool,
-    ϕms_mcs2D_buffer = nothing, # only required for 3D case
+    ϕms_mcs2D_buffer = nothing, # only required for pbm_covar_indices not nothing
     ) where {TG, TF}
         apply_model!(ϕm, g, xM, ϕg; is_testmode) # allocates view
         return nothing
@@ -243,7 +243,7 @@
     ϕms_mcs2D_buffer = reshape(ϕm, size(ϕm,1), :) # allocates 48 bytes for an escaping view, preallocate 
     ) where {TG, TF}
     update_xMP!(xMP, xM, ζsP, pbm_covar_indices)
-    #Main.@infiltrate_main
+    #stack the n_site * n_MC and copy back to array afterwards, if view points to other
     if pointer(parent(ϕms_mcs2D_buffer)) != pointer(parent(ϕm))
         # if supplied a preallocated array (rather than using the default view) theń copy
         copyto!(ϕms_mcs2D_buffer, ϕm) # in order to call apply_model only once, stack n_sites * n_MC
@@ -254,7 +254,6 @@
     end
     return nothing
 end
-
 function update_xMP!(xMP::AbstractMatrix{TG}, 
     xM::AbstractMatrix{TG}, ζsP::AbstractMatrix{TF}, pbm_covar_indices::AbstractVector{<:Number}
     ) where {TG, TF}
@@ -265,6 +264,7 @@
     @inbounds for i in 1:n_site
         for j in 1:n_MC
             # Copy xM block
+            # for each site i there are n_MC columns, j, with the same xM
             ic = (i-1)*n_MC + j
             for k in 1:n_cov
                 xMP[k, ic] = xM[k, i]
@@ -277,6 +277,39 @@
     end
 end
 
+function g_apply!(ϕm::AbstractMatrix{TF}, ϕg::AbstractVector{TG}, xM::AbstractMatrix{TG}, 
+    ζP::AbstractVector{TF}, pbm_covar_indices::AbstractVector{<:Number}, 
+    g::AbstractModelApplicator,
+    xMP::AbstractMatrix,
+    is_testmode::Bool,
+    ϕms_mcs2D_buffer = reshape(ϕm, size(ϕm,1), :) # allocates 48 bytes for an escaping view, preallocate 
+    ) where {TG, TF}
+    update_xMP!(xMP, xM, ζP, pbm_covar_indices) # ζP vector version
+    @assert size(xMP,2) == size(ϕm,2) # same number of sites
+    apply_model!(ϕm, g, xMP, ϕg; is_testmode)
+    return nothing
+end
+function update_xMP!(xMP::AbstractMatrix{TG}, 
+    xM::AbstractMatrix{TG}, ζP::AbstractVector{TF}, pbm_covar_indices::AbstractVector{<:Number}
+    ) where {TG, TF}
+    (n_θP,) = size(ζP)
+    n_cov, n_site = size(xM)
+    n_covP = length(pbm_covar_indices)
+    @assert size(xMP) == ((n_cov + n_covP) , n_site)
+    @inbounds for i in 1:n_site
+        # Copy xM block
+        for k in 1:n_cov
+            xMP[k, i] = xM[k, i]
+        end
+        # Fill pbm covariates block
+        for (k, idx) in enumerate(pbm_covar_indices)
+            xMP[n_cov + k, i] = (TG === TF) ? ζP[idx] : convert(TG, ζP[idx])
+        end
+    end
+end
+
+
+
 function transformζ!(θs, trans, ζs::AbstractArray{TF}) where TF
     n_MC = size(ζs,2)
     # θs1 = exp.(ζs)
@@ -308,12 +341,12 @@
     (;xP, y_o, y_unc, i_indiv_train) = nljoint_args_indi
     @assert size(y_o) == size(y_unc) == (n_obs,)
     @assert size(θsP,2) == size(θsM,2) == n_MC
     (;f!,fLogLik) = nljoint_args_fix
     #θsPc = CA.ComponentMatrix(θsP, axθP, CA.FlatAxis())
     #θsMc = CA.ComponentMatrix(θsM, axθM, CA.FlatAxis())
     apply_model!(y_pred, f!, θsP, θsM, xP)
     nL_reg = sum(y_pred) / n_MC
     nL = fLogLik(y_pred, y_o, y_unc) / n_MC
     # nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC
     (; nL,)
     # loss_comps = neg_elbo_ζtf(
@@ -363,7 +396,7 @@
     pred
 end
 
 function fLogLik_dummy(y_pred, y_o, y_unc)
     sum(y_pred)
 end
 

src/elbo_site_grad.jl

@@ -133,7 +133,7 @@
 
 
 # ϕm vector version
 # TOOD pass closure object rather than many parameters
 function forwarddiff_grad_nelboi_z!(hi, approx::AbstractHVIApproximation, rnormM, 
     nljoint_args_indi, ϕm::AbstractVector{TF}, i, 
     ϕqIc, θsP, dϕmvecs, dθsP, hw_channel::Channel, cor_endsM, transM::Stacked, 
@@ -516,7 +516,7 @@
         grad_ϕm_ϕqI1 = grad_ϕm_ϕqI(ϕm, ϕqIc)
         function grad_conf_vec_ϕmvec() 
             cfg_vec = ForwardDiff.GradientConfig(nelboi_z_vec, grad_ϕm_ϕqI1, diffchunk_vec)
-            update_duals_Vcat(cfg_vec, grad_ϕm_ϕqI1)
+            update_duals_Vcat2(cfg_vec, grad_ϕm_ϕqI1)
         end
         grad_conf_mat_ϕmvec, grad_conf_vec_ϕmvec
     else
@@ -529,7 +529,7 @@
         rad_θP_ϕm1 = grad_θsP_ϕm(ϕm, θsP)
         function grad_conf_mat_ϕmmat() 
             cfg_mat = ForwardDiff.GradientConfig(nelboi_z_mat, rad_θP_ϕm1, diffchunks_mat)
-            update_duals_Vcat(cfg_mat, rad_θP_ϕm1)
+            update_duals_Vcat2(cfg_mat, rad_θP_ϕm1)
         end
         grad_conf_mat_ϕmmat, grad_conf_vec_ϕmmat
     end
@@ -579,15 +579,28 @@
 @inline grad_θsP_ϕm(ϕm::AbstractMatrix, θsP) = Vcat(similar(θsP), similar(ϕm))
 
 
-@inline function update_duals_Vcat(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N}
+@inline function update_duals_Vcat2(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N}
     # similar Vcat returns a vector, need to convert dual of cfg to Vcat
     # no view, but copy with index_at_dim
+    @assert length(vc.args) == 2
     duals_vcat = Vcat(index_at_dim(cfg.duals, axes(vc.args[1], 1); dim=1), 
                       index_at_dim(cfg.duals, axes(vc.args[2], 1); dim=1));
     cfg = ForwardDiff.GradientConfig{T,V,N,typeof(duals_vcat)}(cfg.seeds, duals_vcat)
     @assert cfg.duals isa Vcat
     cfg
 end
+@inline function update_duals_Vcat3(cfg::ForwardDiff.GradientConfig{T,V,N}, vc::Vcat) where {T,V,N}
+    # similar Vcat returns a vector, need to convert dual of cfg to Vcat
+    # no view, but copy with index_at_dim
+    @assert length(vc.args) == 3
+    duals_vcat = Vcat(index_at_dim(cfg.duals, axes(vc.args[1], 1); dim=1), 
+                      index_at_dim(cfg.duals, axes(vc.args[2], 1); dim=1),
+                      index_at_dim(cfg.duals, axes(vc.args[3], 1); dim=1),
+                      );
+    cfg = ForwardDiff.GradientConfig{T,V,N,typeof(duals_vcat)}(cfg.seeds, duals_vcat)
+    @assert cfg.duals isa Vcat
+    cfg
+end
 
 
 

src/pointloss_site.jl

@@ -0,0 +1,111 @@
+"""
+elbo_helpers need to be initialized with new random numbers
+in h.ζP and hi.ζM by calling randnPM! before.
+By this way, we can compute the derivative corresponding to the forward pass
+"""
+function pointloss_sites!(
+    pointloss_helpers::NamedTuple,      # tuple of preallocated arrays
+    ϕ::NamedTuple,
+    sample_args::NamedTuple,
+    indiv_args::NamedTuple,
+    nljoint_args_fix::NamedTuple;
+) 
+    # for AD do not put it into closure
+    h = pointloss_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP
+    (;ϕg, ϕqP, ϕqI) = ϕ
+    (;g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args
+    @assert keys(indiv_args)[1] == :xM
+    xM = indiv_args.xM; r2end = 2:length(indiv_args)
+    nljoint_args_inds = NamedTuple{keys(indiv_args)[r2end]}(values(indiv_args)[r2end])
+    use_dc = h.helpers_sites[1].ζMm isa PAT.DiffCache
+    (n_θM,) = use_dc ? size(h.helpers_sites[1].ζMm.du) : size(h.helpers_sites[1].ζMm)
+    ϕqPc = intϕqP(ϕqP) 
+    ϕqIc = intϕqI(ϕqI)
+    pbm_covar_indices = !isnothing(pbm_covar_indices) && isempty(pbm_covar_indices) ? nothing : pbm_covar_indices
+    #
+    h.ζPm .= CA.getdata(view(ϕqPc,Val(:μζP)))
+    # (n_M x n_site)  
+    h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕmsa,:,1,:) : h.ϕmsa
+    g_apply!(h_ϕms, ϕg, xM, h.ζPm, pbm_covar_indices, g, h.xMP, is_testmode) 
+    exp_ladJacTP = transformζ!(h.θPm, transP, h.ζPm)  # return value captures ladJacT
+    # so that one can provide its gradient to the pullback
+    ϕm_it = eachslice(h.ϕmsa; dims=3)
+    template = ϕqI # only important for gradient
+    θPm = h.θPm
+    keys_nljoint_args_indi_val = Val(keys(nljoint_args_inds))
+    # closure with approx, transM, and kwargs
+    function compute_pointlossi_z_cl!(hi, nljoint_args_indi_tup, ϕm) 
+        #nljoint_args_indi = NamedTuple{keys_nljoint_args_indi}(nljoint_args_indi_tup)
+        nljoint_args_indi = name_tuple(nljoint_args_indi_tup, keys_nljoint_args_indi_val)
+        compute_pointlossi_z!(hi, nljoint_args_indi, ϕm, 
+        ϕqIc, θPm, cor_ends.M, transM, nljoint_args_fix) 
+    end
+    #res_site = map(compute_pointlossi_z!, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it)
+    #MAYBE: distributed mapreduce: 
+    #   https://docs.julialang.org/en/v1/stdlib/Distributed/#Distributed.@distributed
+    #   https://github.com/SupaeroDataScience/DE/blob/main/notebooks/Introduction%20to%20MapReduce.ipynb
+    #zind = Test.@inferred zip_eachlastdims(nljoint_args_inds)
+    zind = zip_eachlastdims(nljoint_args_inds)
+    pointloss_z = mapreduce(compute_pointlossi_z_cl!, +, h.helpers_sites, zind, ϕm_it)
+    pointloss = pointloss_z - exp_ladJacTP 
+end
+
+function compute_pointlossi_z!(hi, 
+    nljoint_args_indi, ϕms::AbstractMatrix, ϕqIc::AbstractVector, θPm::AbstractMatrix, 
+    cor_endsM, transM, nljoint_args_fix, ) 
+    # DM is 1 (vector) if no population covariates - constant across MC
+    #       2 (Matrix) if g is driven by population covariates - differs across MC
+    # on update -> sync corresponding function within grad_neg_elbo_sites
+    if hi.ζMm isa PAT.DiffCache
+        # Gradiend called twice because of different dimensions need to pick the Dual
+        template = eltype(ϕqIc) <: ForwardDiff.Dual ? ϕqIc : θPm
+        hi = map_leaves_nt(x -> PAT.get_tmp(x, template), hi)
+    end
+    n_θM = cor_endsM[end]
+    hi.ζMm .= view(ϕms, 1:n_θM, :)    # μζM in first n_θM entries
+    transformζ!(hi.θMm, transM, hi.ζMm)  
+    exp_nL = exp_nLi!(hi.ym, θPm, hi.θMm, nljoint_args_indi, nljoint_args_fix)[1]
+    pointlossi = exp_nL 
+end
+
+function prepare_pointloss_helpers(approx::AbstractHVIApproximation, 
+    ϕg::AbstractArray{TG}, template_TF::AbstractArray{TF};
+    n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs,
+    use_diff_cache::Val{use_dc} = Val(true),
+    diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8),
+    ) where {TG, TF, use_dc, chunk}
+    his = Tuple((;
+        ζMm = Matrix{TF}(undef, n_θM, 1),
+        θMm = Matrix{TF}(undef, n_θM, 1),
+        ym = Matrix{TF}(undef, n_obs,1),
+    ) for i in 1:n_site)
+    helpers_sites = use_dc ? map_leaves_nt(x -> PAT.DiffCache(x, chunk), his) : his
+    h = (;
+        ζPm = Matrix{TF}(undef, n_θP,1),
+        θPm = Matrix{TF}(undef, n_θP,1),
+        ϕmsa = Array{TF}(undef, n_M, 1, n_site),  
+        xMP = Matrix{TG}(undef, (n_cov + n_covP), n_site),
+        helpers_sites,
+    )
+end
+
+function check_pointloss_helpers(h::NamedTuple, xM::AbstractMatrix, pbm_covar_indices)
+    n_cov, n_site = size(xM)
+    n_covP = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices)
+    (n_θP,_) = size(h.ζPm)
+    n_M = size(h.ϕmsa, 1)
+    @assert size(h.ζPm) == (n_θP, 1)
+    @assert size(h.θPm) == (n_θP, 1)
+    @assert size(h.ϕmsa) == (n_M, 1, n_site)
+    @assert size(h.xMP) == ((n_cov + n_covP), n_site) 
+    #
+    @assert length(h.helpers_sites) == n_site
+    hi = h.helpers_sites[1]
+    n_θM = size(hi.ζMm.du, 1)
+    n_obs = size(hi.ym.du, 1)
+    @assert size(hi.ζMm.du) == (n_θM, 1)
+    @assert size(hi.θMm.du) == (n_θM, 1)
+    @assert size(hi.ym.du) == (n_obs, 1)
+end
+
+