diff --git a/src/HybridSolver_site.jl b/src/HybridSolver_site.jl new file mode 100644 index 0000000..03c04ff --- /dev/null +++ b/src/HybridSolver_site.jl @@ -0,0 +1,374 @@ +""" +A Solver that solves for the mode of hybrid variational inference problem +([`AbstractHybridProblem`](@ref)). + + +It does not perform the full posterior inference, +but only finds the optimal parameters of the variational distribution. +This is useful for initializing the full posterior inference. + +- `alg`: The optimization algorithm to use for solving the problem, e.g. `Adam()` +""" +struct HybridPointSolverSite{A} <: AbstractHybridSolver + alg::A +end + +HybridPointSolverSite(; alg) = HybridPointSolverSite(alg) + +function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolverSite; + scenario=Val(()), rng=Random.default_rng(), + gdevs = nothing, # get_gdev_MP(scenario) + is_inferred::Val{is_infer} = Val(false), + ad_backend_loss = AutoEnzyme(), + epochs, + epochs_callback = max(1, floor(Int, epochs/10)), + callback = (state, loss_val) -> false, + is_omitting_NaNbatches = false, + is_omit_priors::Val{omit_priors} = Val(false), + clusters::AbstractVector{<:Integer} = + 1:first(get_hybridproblem_n_site_and_batch(prob; scenario)), + cluster_rep = 1, + kwargs... +) 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) + # TODO: separate parameters from problem description - right now optimized + ϕq0 = get_hybridproblem_ϕq(prob; scenario) + (; transP, transM) = get_hybridproblem_transforms(prob; scenario) + n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(ϕq0))) + ϕq_ranef = setup_ϕq_ranef(ranef) + if :ranef ∉ keys(ϕq0) + ϕq0 = CA.ComponentVector(ϕq0, ranef = ϕq_ranef) + else + @assert size(ϕq0[Val(:ranef)]) == size(ϕq_ranef) + end + ϕP0 = ϕq0[Val(:μP)] + #intϕ = ComponentArrayInterpreter(CA.ComponentVector(ϕg=1:length(ϕg0), ϕq=ϕq0)) + #ϕ0_cpu = vcat(ϕg0, pt.θP .* FT(0.9)) # slightly disturb θP_true + ϕ0_cpu = CA.ComponentVector(ϕg=ϕg0, ϕq=ϕq0) + intϕ = ComponentArrayInterpreter(ϕ0_cpu) + n_sites_cluster = [count(==(element),clusters) for element in 1:maximum(clusters)] + frac_cluster_all = (1 / cluster_rep) ./ n_sites_cluster[clusters] + train_loader = get_hybridproblem_train_dataloader(prob; scenario) + test_data = get_hybridproblem_test_data(prob; scenario) + # i_test = rand(1:n_site, Integer(floor(n_site/10))) + # test_data = map(train_loader.data) do data_comp + # ndims(data_comp) == 2 ? data_comp[:, i_test] : data_comp[i_test] + # end + gdev = gdevs.gdev_M + if gdev isa MLDataDevices.AbstractGPUDevice + ϕ0_dev = gdev(ϕ0_cpu) + g_dev = gdev(g) + train_loader_dev = gdev_hybridproblem_dataloader(train_loader; gdevs) + test_data_dev = gdev_hybridproblem_data( + test_data[keys(test_data)[1:5]]; gdevs) + else + ϕ0_dev = ϕ0_cpu + g_dev = g + train_loader_dev = train_loader + test_data_dev = test_data[keys(test_data)[1:5]] + end + f = get_hybridproblem_PBmodel(prob; scenario) + ftest = create_nsite_applicator(f, size(test_data[1],2)) + if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice + f_dev = gdevs.gdev_P(f) + ftest_dev = gdevs.gdev_P(ftest) + else + f_dev = f + ftest_dev = ftest + end + py = get_hybridproblem_neg_logden_obs(prob; scenario) + pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) + n_site_test = size(test_data[1],2) + 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)) + #intP = ComponentArrayInterpreter(pt.θP) + # intθP = ComponentArrayInterpreter(pt.θP) + # intθM = ComponentArrayInterpreter(pt.θM) + penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) + # loss_gf = get_loss_gf(g_dev, transM, transP, f_dev, py, intϕ; + # n_site_batch=n_batch, + # par_templates = pt, + # cdev=infer_cdev(gdevs), pbm_covars, + # priorsP, priorsM, is_omit_priors, penalty_computer, + # #intθM, intθP, + # frac_cluster_all, ranef, + # ) + # loss_gf_test = get_loss_gf(g_dev, transM, transP, ftest_dev, py, intϕ; + # n_site_batch=n_site_test, + # par_templates = pt, + # cdev=infer_cdev(gdevs), pbm_covars, + # priorsP, priorsM, is_omit_priors, penalty_computer, + # #intθM, intθP, + # frac_cluster_all, ranef, + # ) + # # call loss function once + # l1 = is_infer ? + # Test.@inferred(loss_gf(ϕ0_dev, first(train_loader_dev)...; is_testmode=true))[1] : + # # using ShareAdd; @usingany Cthulhu + # # @descend_code_warntype loss_gf(ϕ0_dev, first(train_loader_dev)...) + # loss_gf(ϕ0_dev, first(train_loader_dev)...; is_testmode=true)[1] + # l1t = loss_gf_test(ϕ0_dev, test_data_dev...; is_testmode=true)[1] + # and gradient + # xMg, xP, y_o, y_unc = first(train_loader_dev) + # gr1 = Zygote.gradient( + # p -> loss_gf(p, xMg, xP, y_o, y_unc)[1], + # ϕ0_dev) + # Zygote.gradient(ϕ0_dev -> loss_gf(ϕ0_dev, data1...)[1], ϕ0_dev) + # + # implement training loop by hand to skip minibatches with NaN gradients + ps = CA.getdata(ϕ0_dev) + opt_st_new = Optimisers.setup(solver.alg, ps) + n_skips = 0 + # prepare DI.gradient, need to access and update outside cope data_batch + # because cannot redefine fopt_loss_gf + data_batch = first(train_loader_dev) + is_testmode = false + function fopt_loss_gf(ϕ) + #@show first(data_batch[5], 2) + loss_gf(ϕ, data_batch...; is_testmode)[1] + end + ad_prep = DI.prepare_gradient(fopt_loss_gf, ad_backend_loss, zero(ps)) + grad = similar(ps) + stime = time() + for epoch in 1:epochs + is_testmode = false + #i,data_batch = first(enumerate(loader)) + for (i, data_batch_) in enumerate(train_loader_dev) + data_batch = data_batch_ # propagate outside for to scope of fopt_loss_gf + DI.gradient!(fopt_loss_gf, grad, ad_prep, ad_backend_loss, ps) + if any(isnan.(grad)) + n_skips += 1 + #println("Skipped NaN : Batch $i") + print(",$i") + else + Optimisers.update!(opt_st_new, ps, grad) + end + end + ttime = time() - stime + # compute loss for test data + l = loss_gf_test(ps, test_data_dev...; is_testmode = true) + println() + @show round(ttime, digits=1), epoch, l.nLy, l.neg_log_prior, l.loss_penalty + # TODO log + end + res = nothing + ϕ = intϕ(ps) + #θP = !isempty(ϕ.ϕP) ? cpu_ca(apply_preserve_axes(transP, cpu_ca(ϕ).ϕq.μP)) : CA.ComponentVector{eltype(ϕ)}() + # TODO check which components live on gpu and which on cpu + ϕq_opt = ComponentArrayInterpreter(ϕq0)(gdevs.gdev_P(ϕ.ϕq)) + probo = HybridProblem(prob; ϕg=cpu_ca(ϕ).ϕg, ϕq=ϕq_opt) + (; ϕ, resopt=res, probo) +end + +function get_callback_epochs(epochs_callback; + n_site, n_batch, callback, loss_test) + n_per_epoch = n_site ÷ n_batch + if epochs_callback == 0 + callback_epochs = callback + else + callback_epochs = function(state, l) + if (state.iter == 1) || (state.iter % (n_per_epoch * epochs_callback) == 0) + l_test = loss_test(state)[1] + println("epoch = $((state.iter) ÷ n_per_epoch), iter = $(state.iter), loss_train=$l, loss_test=$l_test") + end + return callback(state, l) + end + end +end + +""" +A Solver that solves for the full posterior of a variational inference problem +([`AbstractHybridProblem`](@ref)). + +Parameters +- `alg`: The optimization algorithm to use for solving the problem, e.g. `Adam()` +- `n_MC`: The number of Monte Carlo samples to use for estimating the gradients of the loss function. +""" +struct HybridPosteriorSolverSite{A} <: AbstractHybridSolver + alg::A + n_MC::Int + n_MC_cap::Int +end +function HybridPosteriorSolverSite(; alg, n_MC=12, n_MC_cap=n_MC) + HybridPosteriorSolverSite(alg, n_MC, n_MC_cap) +end +function HybridPosteriorSolverSite(solver::HybridPosteriorSolverSite; + alg=solver.alg, + n_MC=solver.n_MC, + n_MC_cap=n_MC) + HybridPosteriorSolverSite(alg, n_MC, n_MC_cap) +end + +""" + solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolverSite; epochs, ...) + +Perform the inversion of HVI Problem. + +Arguments +- `prob`: The AbstractHybridProblem to solve. +- `scenario`: Scenario to query prob, defaults to `Val(())`. +- `epochs`: number of epochs to train, i.e. number of passes through the whole dataset. + +Optional keyword arguments +- `rng`: Random generator, defaults to `Random.default_rng()`. +- `gdevs`: `NamedTuple` `(;gdev_M, gdev_P)` functions to move + computation and data of ML model on and PBM respectively + to gpu (e.g. `gpu_device()` or cpu (`identity`). + defaults to [`get_gdev_MP`](@ref)`(scenario)` +- `θmean_quant` default to `0.0`: deprecated +- `is_inferred`: set to `Val(true)` to activate type stability checks +- `is_omit_priors`: set to `Val(true)` to omit priors in the loss computation, which can be useful for debugging or if priors are not implemented for a specific scenario (e.g. on gpu) +- `clusters`: vector of cluster assignments for each site, defaults to each site being its own cluster. Clusters are used to compute the loss in a way that accounts for clustering of sites, which can be useful if there are many sites and the number of Monte Carlo samples is limited. +- `cluster_rep`: number of times to repeat each cluster in the loss computation, defaults to 1. Repeating clusters can be useful to effectively increase the number of Monte Carlo samples when the number of clusters is small. +- `epochs_callback`: number of epochs between progress output on evaluating testdata + +Returns a `NamedTuple` of +- `probo`: A copy of the HybridProblem, with updated optimized parameters +- `interpreters`: TODO +- `ϕ`: the optimized HVI parameters: a `ComponentVector` with entries + - `ϕg`: The ML model parameter vector, + - `ϕq`: `ComponentVector` of non-ML parameters, including + `μP`: `ComponentVector` of the mean global PBM parameters at unconstrained scale +- `θP`: `ComponentVector` of the mean global PBM parameters at constrained scale +- `resopt`: the structure returned by `Optimization.solve`. It can contain + more information on convergence. +""" +function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolverSite; + epochs, + scenario::Val{scen}=Val(()), rng=Random.default_rng(), + gdevs = get_gdev_MP(scenario), + θmean_quant=0.0, + is_inferred::Val{is_infer} = Val(false), + is_omit_priors::Val{omit_priors} = Val(false), + approx = prob.approx, + clusters::AbstractVector{<:Integer} = + 1:first(get_hybridproblem_n_site_and_batch(prob; scenario)), + cluster_rep = 1, + epochs_callback = max(1, floor(Int, epochs/10)), + callback = (state, loss_val) -> false, + kwargs... +) where {scen, is_infer, omit_priors} + pt = get_hybridproblem_par_templates(prob; scenario) + cor_ends = get_hybridproblem_cor_ends(prob; scenario) + g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario) + (; transP, transM) = get_hybridproblem_transforms(prob; scenario) + pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) + n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + n_sites_cluster = [count(==(element),clusters) for element in 1:maximum(clusters)] + frac_cluster_all = (1 / cluster_rep) ./ n_sites_cluster[clusters] + ranef_spec = get_hybridproblem_ranef(prob; scenario) + ranef = get_ranef_computer( + ranef_spec, keys(pt.θM), n_site, one(eltype(pt.θM))) + ϕq_ranef = setup_ϕq_ranef(ranef) + ϕq = CA.ComponentVector(get_hybridproblem_ϕq(prob; scenario), ranef = ϕq_ranef) + (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕq) + int_ϕq = interpreters.ϕq + int_ϕg_ϕq = interpreters.ϕg_ϕq + 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) + train_loader = get_hybridproblem_train_dataloader(prob; scenario) + test_data = get_hybridproblem_test_data(prob; scenario) + # i_test = rand(1:n_site, Integer(floor(n_site/10))) + # test_data = map(train_loader.data) do data_comp + # ndims(data_comp) == 2 ? data_comp[:, i_test] : data_comp[i_test] + # end + n_batch_test = size(test_data[1],2) + if first(train_loader)[1] isa CA.ComponentArray + @warn("ML model covariates (1) were provided as ComponentArray. " * + "Consider providing them as a plain array.") + end + if first(train_loader)[2] isa CA.ComponentArray + @warn("PBM drivers (2) were provided as ComponentArray. " * + "Consider providing them as a plain array.") + end + if gdevs.gdev_M isa MLDataDevices.AbstractGPUDevice + ϕ0_dev = gdevs.gdev_M(ϕ) + g_dev = gdevs.gdev_M(g) # zygote fails if gdev is a CPUDevice, although should be non-op + train_loader_dev = gdev_hybridproblem_dataloader(train_loader; gdevs) + test_data_dev = gdev_hybridproblem_data(test_data[keys(test_data)[1:5]]; gdevs) + else + ϕ0_dev = ϕ + g_dev = g + train_loader_dev = train_loader + test_data_dev = test_data[keys(test_data)[1:5]] + end + f = get_hybridproblem_PBmodel(prob; scenario) + f_test = create_nsite_applicator(f, n_batch_test) + if gdevs.gdev_P isa MLDataDevices.AbstractGPUDevice + f_dev = gdevs.gdev_P(f) #fmap(gdevs.gdev_P, f) + f_test_dev = gdevs.gdev_P(f_test) + else + f_dev = f + f_test_dev = f_test + end + + py = get_hybridproblem_neg_logden_obs(prob; scenario) + + penalty_computer = get_hybridproblem_penalty_computer(prob; scenario) + # intθP = ComponentArrayInterpreter(pt.θP) + # intθMs = ComponentArrayInterpreter((n_batch,), pt.θM) + + loss_elbo = get_loss_elbo( + g_dev, transP, transM, f_dev, py, n_batch; + n_MC = solver.n_MC, n_MC_cap = solver.n_MC_cap, cor_ends, + cdev=infer_cdev(gdevs), pbm_covars, + par_templates = pt, + #pt.θP, + int_ϕq, int_ϕg_ϕq, priorsP, priorsM, + is_omit_priors, zero_prior_logdensity, approx, penalty_computer, + ranef, + #intθMs, intθP, + frac_cluster_all, + ) + loss_elbo_test = get_loss_elbo( + g_dev, transP, transM, f_test_dev, py, n_batch_test; + solver.n_MC, solver.n_MC_cap, cor_ends, + cdev=infer_cdev(gdevs), pbm_covars, + par_templates = pt, + #pt.θP, + int_ϕq, int_ϕg_ϕq, priorsP, priorsM, + is_omit_priors, zero_prior_logdensity, approx, penalty_computer, + ranef, + #intθMs, intθP, + frac_cluster_all, + ) + + # test loss function once + # tmp = first(train_loader_dev) + # using ShareAdd + # @usingany Cthulhu + # @descend_code_warntype loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...) + # omit for type stability in AD + @assert length(first(train_loader_dev)) == 5 + l0 = + #is_infer ? + # (Test.@inferred loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...; is_testmode=true)) : + loss_elbo(ϕ0_dev, rng, first(train_loader_dev)...; is_testmode=false) + l0t = loss_elbo_test(ϕ0_dev, rng, test_data_dev...; is_testmode=true) + loss_test = let test_data_dev = test_data_dev + (state) -> loss_elbo_test(state.u, rng, test_data_dev...; is_testmode=true) + end + callback_epochs = get_callback_epochs(epochs_callback; + n_site, n_batch, callback, loss_test) + optf = Optimization.OptimizationFunction( + (ϕ, data) -> first(loss_elbo(ϕ, rng, data...; is_testmode=false)), + Optimization.AutoZygote()) + optprob = OptimizationProblem(optf, CA.getdata(ϕ0_dev), train_loader_dev) + res = Optimization.solve(optprob, solver.alg; callback = callback_epochs, epochs, kwargs...) + ϕc = interpreters.ϕg_ϕq(cpu_device()(res.u)) + ϕq = ϕc[Val(:ϕq)]; + ϕg = ϕc[Val(:ϕg)]; + probo = HybridProblem(prob; ϕg, ϕq) + θP = get_hybridproblem_θP(probo) + (; probo, interpreters, ϕ=ϕc, θP, resopt=res) +end + diff --git a/src/HybridVariationalInference.jl b/src/HybridVariationalInference.jl index d54cfce..905a0d7 100644 --- a/src/HybridVariationalInference.jl +++ b/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") diff --git a/src/PBMApplicator.jl b/src/PBMApplicator.jl index e8b1ab4..91e5576 100644 --- a/src/PBMApplicator.jl +++ b/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]) diff --git a/src/elbo_site.jl b/src/elbo_site.jl index a9a07be..04bc3f3 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -28,8 +28,8 @@ function neg_elbo_sites!( 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 @@ function neg_elbo_sites!( # 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 @@ function neg_elbo_sites!( 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) @@ -232,7 +232,7 @@ function g_apply!(ϕm::AbstractMatrix{TF}, ϕg::AbstractVector{TG}, xM::Abstract 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 @@ function g_apply!(ϕm::AbstractArray{TF,3}, ϕg::AbstractVector{TG}, xM::Abstrac ϕ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 @@ function g_apply!(ϕm::AbstractArray{TF,3}, ϕg::AbstractVector{TG}, xM::Abstrac 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 @@ function update_xMP!(xMP::AbstractMatrix{TG}, @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 @@ function update_xMP!(xMP::AbstractMatrix{TG}, 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) diff --git a/src/elbo_site_grad.jl b/src/elbo_site_grad.jl index 29c80d2..919eb7e 100644 --- a/src/elbo_site_grad.jl +++ b/src/elbo_site_grad.jl @@ -516,7 +516,7 @@ function prepare_gradelbo_helpers( 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 @@ function prepare_gradelbo_helpers( 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 @@ end @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 diff --git a/src/pointloss_site.jl b/src/pointloss_site.jl new file mode 100644 index 0000000..4831fd4 --- /dev/null +++ b/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 + + diff --git a/src/pointloss_site_grad.jl b/src/pointloss_site_grad.jl new file mode 100644 index 0000000..5fc157a --- /dev/null +++ b/src/pointloss_site_grad.jl @@ -0,0 +1,290 @@ + +function grad_pointloss_sites( + pointloss_helpers::NamedTuple, # tuple of preallocated arrays + grad_pointloss_helpers::NamedTuple, # tuple of preallocated arrays pullback closures + ϕ::NamedTuple, + sample_args::NamedTuple, + indiv_args::NamedTuple, + nljoint_args_fix::NamedTuple; + # approx::AbstractHVIApproximation, + # rnormPM::NamedTuple, # tuple of random numbers + # ϕg::AbstractVector{TG}, ϕqP::AbstractVector{TF}, ϕqI::AbstractVector{TF}, g, + # pbm_covar_indices::Union{Nothing,AbstractVector{<:Number}}, + # args...; + # nljoint_args_inds, # indices of sites in training set + # intϕqP, intϕqI, + # xM, + # cor_ends, + # transP::Stacked, transM::Stacked, + # is_testmode, + n_workers = 1, # TODO compute from executor and Distributed.nworkers and nthreads,... + executor::Transducers.Executor = Transducers.SequentialEx(), +) + h = pointloss_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP + (;ϕg, ϕqP, ϕqI) = ϕ + TF = eltype(ϕqP) + (;approx, 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_ϕm_matrix = isnothing(pbm_covar_indices) + ϕqPc = intϕqP(ϕqP) + ϕqIc = intϕqI(ϕqI) + #ϕm_buffer_key = use_ϕm_matrix ? :ϕms : :ϕms_mcs + #h_ϕms = h[ϕm_buffer_key] + h_ϕms = isnothing(pbm_covar_indices) ? view(h.ϕmsa,:,1,:) : h.ϕmsa + check_pointloss_helpers(h, xM, pbm_covar_indices) + n_cov, n_site = size(xM) + n_θP, n_MC = size(h.ζPm) + @assert n_MC == 1 + ax_θP = get_axis_θP(nljoint_args_fix.f!) + ax_θM = get_axis_θM(nljoint_args_fix.f!) + n_θM = axis_length(ax_θM) + n_M = get_numberof_MLinputs(approx, 1:n_θM) + nljoint_args_indi = NamedTuple{keys(nljoint_args_inds)}( + first(zip_eachlastdims(nljoint_args_inds))) + gradh = !isempty(grad_pointloss_helpers) ? grad_pointloss_helpers : + Test.@inferred prepare_gradpointloss_helpers( + #prepare_gradpointloss_helpers( + ϕqIc, selectdim(h_ϕms, ndims(h_ϕms), 1), h.ζPm, ϕg, ϕqPc, approx; + pbm_covar_indices, n_workers, + h, nljoint_args_indi, + n_site, n_M, n_cov, cor_ends, transP, transM, nljoint_args_fix, ) + # + h.ζPm .= CA.getdata(view(ϕqPc,Val(:μζP))) + g_apply!(h_ϕms, ϕg, xM, h.ζPm, pbm_covar_indices, g, h.xMP, is_testmode) + ladJacTP = transformζ!(h.θPm, transP, h.ζPm) # return value captures ladJacT + # + # parallel ForwardDiffGradient through forwarddiff_grad_pointlossi_z! + cl = ForwardDiffGradPointlossiZCl( + approx, ϕqIc, h.θPm, gradh.dϕmvecs, gradh.hw_channel, cor_ends.M, transM, + Val(keys(nljoint_args_inds)), nljoint_args_fix) + ϕm_it = eachslice(h.ϕmsa; dims = ndims(h.ϕmsa)) + init = (; + dϕqI = zero(TF) .* SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(ϕqIc), + dθsP = zero(TF) .* SA.SVector{axis_length(ax_θP)*1}(h.θPm), ) + res = Folds.mapreduce(cl, make_tuple_reducer(+), + zip(h.helpers_sites, zip_eachlastdims(nljoint_args_inds), ϕm_it, 1:n_site), + executor; init) + # + # copy SharedArray into preallocated plain array for ∂elbo_∂θP + ∂elbo_∂ϕqm = gradh.∂elbo_∂ϕqm + copyto!(∂elbo_∂ϕqm, gradh.dϕmvecs) # copy from shared array + # copy StaticVector into preallocated ComponentVector + ∂elbo_∂θP = gradh.∂elbo_∂θP + copyto!(∂elbo_∂θP, res.dθsP) + copyto!(gradh.dϕqIc, res.dϕqI) + # other gradients are one, calculated from formula + gradh.∂elbo_∂logσ_ζP .= -one(TF) + ∂elbo_∂ladJacTP = -one(TF) + # + # pullback gradients of ϕqm -> gradh.dϕg and gradh.dζsP + gradh.pullback_g_apply!( + gradh.dϕg, gradh.∂elbo_∂ϕm_∂ζP, h_ϕms, ∂elbo_∂ϕqm, + ϕg, xM, h.ζPm, pbm_covar_indices, g, is_testmode) + # + # pullback gradients of ∂elbo_∂θP to ∂elbo_∂θP_∂ζP + ∂elbo_∂θP_∂ζP = gradh.∂elbo_∂θP_∂ζP + gradh.pullback_cl_transformζsP!( + ∂elbo_∂θP_∂ζP, + ∂elbo_∂θP, + ∂elbo_∂ladJacTP, + h.θPm, + transP, + h.ζPm, + ) + # + # pullback gradients of ∂elbo_∂ζP, ∂elbo_∂ϕm_∂ζP to dϕqP + # for μζP (first n_θM entries) its the sum of the woggles, others are zero + gradh.dϕqPc[1:n_θM] .= ∂elbo_∂θP_∂ζP + gradh.∂elbo_∂ϕm_∂ζP + gradh.dϕqPc[(n_θM+1):end] .= zero(TF) + (;dϕqPc = gradh.dϕqPc, dϕqIc = gradh.dϕqIc, dϕg = gradh.dϕg), gradh +end + +""" +Callable to make deliver arguments that do not differ by individual to Foldl.mapreduce. +""" +struct ForwardDiffGradPointlossiZCl{KEYS, TA, TI, Tθ, TD, THWC, TC, TM, TAF} + approx::TA + ϕqIc::TI + θPm::Tθ + dϕmvecs::TD + hw_channel::THWC + corendsM::TC + transM::TM + # need to store KEYS in type parameter otherwise allocation in NamedTuple(keys) + keys_nljoint_args_indi::Val{KEYS} + nljoint_args_fix::TAF +end +function (f::ForwardDiffGradPointlossiZCl{KEYS})(tup) where KEYS + local hi, nljoint_args_indi_tup, ϕm, i + hi, nljoint_args_indi_tup, ϕm, i = tup + #nljoint_args_indi = NamedTuple{f.keys_nljoint_args_indi}(nljoint_args_indi_tup) + nljoint_args_indi = NamedTuple{KEYS}(nljoint_args_indi_tup) + forwarddiff_grad_pointloss_z!(hi, f.approx, nljoint_args_indi, ϕm, i, + f.ϕqIc, f.θPm, f.dϕmvecs, f.hw_channel, f.corendsM, f.transM, + f.nljoint_args_fix, + nothing, + #true, + ) +end + + +# ϕm vector version +# TOOD pass closure object rather than many parameters +function forwarddiff_grad_pointloss_z!(hi, approx::AbstractHVIApproximation, + nljoint_args_indi, ϕm::AbstractMatrix{TF}, i::Integer, + ϕqIc::AbstractVector, θP::AbstractMatrix, dϕmvecs, hw_channel::Channel, cor_endsM, transM::Stacked, + nljoint_args_fix, + omit_gradient=nothing) where {TF} + # need to calls to ForwardDiff.gradient! to avoid reshape + # group arrays of same dimensionality (here grad_ϕm and grad_ϕqI) into a Vcat + ax_ϕqI = CA.getaxes(ϕqIc)[1] + ax_ϕP = get_axis_θP(nljoint_args_fix.f!) + with_channel_element(hw_channel) do hwi_tup + hwi, hwi_dc = hwi_tup + #grad_conf = hwi.grad_conf + pointlossi_z_vec = make_pointlossz_lazy_cls(hi, nljoint_args_indi, + cor_endsM, transM, nljoint_args_fix, + ax_ϕqI, + ) + grad_θP, grad_ϕm, grad_ϕqI = if isnothing(omit_gradient) + local θP_ϕm_ϕqI + # transform ϕqI to a 1-column matrix to vcat it to θP and ϕm + ϕqI = if strides(CA.getdata(ϕqIc))[1] == 1 + # to avoid allocations with reshape + Base.ReshapedArray(CA.getdata(ϕqIc), (length(CA.getdata(ϕqIc)), 1), ()) + else + reshape(CA.getdata(ϕqIc), :, 1) + end + () -> begin + tmpf(ϕqI) = mapreduce(identity, +, view(ϕqI,:,1)) + @allocated tmpf(ϕqI) # zero + end + θP_ϕm_ϕqI = Vcat(θP, ϕm, ϕqI) + tmp = similar(θP_ϕm_ϕqI) + ForwardDiff.gradient!(hwi.grad_θsP_ϕm_ϕqI, pointlossi_z_vec, θP_ϕm_ϕqI, hwi.grad_conf_point) + hwi.grad_θsP_ϕm_ϕqI.args[1], hwi.grad_θsP_ϕm_ϕqI.args[2], hwi.grad_θsP_ϕm_ϕqI.args[3] + else + θP, ϕm, CA.getdata(ϕqIc) + end + copyto!(view(dϕmvecs, :, i), grad_ϕm) + # returning SVector helps avoiding allocations during reduce + local res + res = (; + dϕqI = SA.SVector{axis_length(CA.getaxes(ϕqIc)[1])}(grad_ϕqI), + dθsP = SA.SVector{axis_length(ax_ϕP)*1}(grad_θP), + ) + end +end + +function make_pointlossz_lazy_cls(hi, nljoint_args_indi, + cor_endsM, transM, nljoint_args_fix, + ax_ϕqI::CA.Axis, + ) + # mat, is a Vcat(θP, ϕm, ϕqI) + function pointlossi_cl_vec(mat::Vcat) + local θP_dual, ϕm_dual, ϕqIc_dual + θP_dual = mat.args[1] + ϕm_dual = mat.args[2] + #ϕqIc_dual = CA.ComponentVector(vec(mat.args[3]), ax_ϕqI) + ϕqIc_dual = CA.ComponentVector(view(mat.args[3],:,1), ax_ϕqI) + compute_pointlossi_z!( + hi, nljoint_args_indi, + ϕm_dual, ϕqIc_dual, θP_dual, + cor_endsM, transM, nljoint_args_fix, + )[1] + end + pointlossi_cl_vec +end + + +function prepare_gradpointloss_helpers( + ϕqIc, ϕm::AbstractArray{TF,ND}, θsP::AbstractMatrix, + ϕg::AbstractVector{TG}, ϕqPc::AbstractVector{TF}, + approx::AbstractHVIApproximation; + pbm_covar_indices, n_workers, + h, nljoint_args_indi, + n_site, n_M, n_cov, cor_ends, + transP::Stacked, transM::Stacked, + nljoint_args_fix, + diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), + ) where {TG, TF, ND, chunk} + hi1 = h.helpers_sites[1] + ax_θP = get_axis_θP(nljoint_args_fix.f!) + ax_θM = get_axis_θM(nljoint_args_fix.f!) + # TODO get n_site, n_cov from + n_covP = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) + n_ϕmvec = length(ϕm) #use_ϕm_matrix ? n_M : n_M * n_MC + (n_θP,n_MC) = size(θsP) + @assert n_MC == 1 + @assert n_θP == axis_length(ax_θP) + # n_M = size(ϕm,1) + # To sync across procs/threads, use a Channel + # https://juliafolds2.github.io/OhMyThreads.jl/stable/literate/tls/tls/#The-safe-way:-Channel + # + @assert length(nljoint_args_indi.xP) > 1 + # pointlossi_z = make_pointlossiz_cl(hi1, approx, rnormMi, nljoint_args_indi, CA.getaxes(cv_grad), + # cor_ends.M, transM, nljoint_args_fix) + hw_dcs_raw = (; + ϕm_dc = similar(ϕm), + # callocation on using ComponentArray (_restructure), need to use plain vector + ϕqI_dc = similar(CA.getdata(ϕqIc)), + θP_dc = similar(θsP), + ) + # make sure, each worker gets its own copy, use similar + # TODO remove hw_dcs, not needed because Vcat of all vectors + hw_dcs = (hw_dcs_raw) -> map_leaves_nt(x -> PAT.DiffCache(similar(x), chunk), hw_dcs_raw) + pointlossi_z_vec = make_pointlossz_lazy_cls(hi1, nljoint_args_indi, + cor_ends.M, transM, nljoint_args_fix, + CA.getaxes(ϕqIc)[1], + ) + local grad_θsP_ϕm_ϕqI1, diffchunks_point + ϕqIm = reshape(CA.getdata(ϕqIc),:,1) + ϕm_mat = (ND == 2) ? ϕm : reshape(ϕm, :, 1) + grad_θsP_ϕm_ϕqI1 = grad_θsP_ϕm_ϕqI(θsP, ϕm_mat, ϕqIm) + # n_M is currently not static, hence diffchunks_vec and gradientConfig not inferred + # however, n_θM is a lower bound, so use this one to set a minimum length + # chunk size might be not as large as possible, but safely not larger than paramvector + diffchunks_point = ForwardDiff.Chunk(min(chunk, + #axis_length(ax_θP) + n_M + axis_length(CA.getaxes(ϕqIc)[1]))) + axis_length(ax_θP) + axis_length(ax_θM) + axis_length(CA.getaxes(ϕqIc)[1]))) + function grad_conf_point() + cfg_vec = ForwardDiff.GradientConfig(pointlossi_z_vec, grad_θsP_ϕm_ϕqI1, diffchunks_point) + update_duals_Vcat3(cfg_vec, grad_θsP_ϕm_ϕqI1) + end + get_helpers_worker = () -> begin + (; + grad_conf_point = grad_conf_point(), + grad_θsP_ϕm_ϕqI = grad_θsP_ϕm_ϕqI(θsP, ϕm_mat, ϕqIm), + ) + end + h1 = get_helpers_worker() + h1_dc = hw_dcs(hw_dcs_raw) + h1_tup = (h1, h1_dc) + # put tuples of h1 an d1_dc into the channel + hw_channel = Channel{typeof(h1_tup)}(n_workers) # workers + parallel threads + put!(hw_channel, h1_tup) + foreach(2:n_workers) do _ + put!(hw_channel, (get_helpers_worker(), hw_dcs(hw_dcs_raw))) + end + (; + dϕg = similar(ϕg), + dϕqPc = similar(ϕqPc), + dϕqIc = similar(ϕqIc), + ∂elbo_∂θP = similar(θsP), + ∂elbo_∂ϕm_∂ζP = similar(θsP), + ∂elbo_∂logσ_ζP = Vector{TF}(undef, size(θsP,1)), + ∂elbo_∂θP_∂ζP = Matrix{TF}(undef, size(θsP)), + ∂elbo_∂ϕqm = Array{TF}(undef, size(ϕm)..., n_site), + hw_channel, + pullback_cl_transformζsP! = get_pullback_cl_transformζ!(CA.getdata(ϕqPc); n_θ = n_θP, n_MC), + pullback_g_apply! = get_pullback_g_apply( + ϕg, ϕqPc; n_θP, n_cov, n_covP, n_MC, n_site, n_M), + # SharedArrays last because of problems with ReferenceRevision + dϕmvecs = SharedArrays.SharedArray{TF}(n_ϕmvec, n_site), + ) +end +# need similar to make a copy for each worker +grad_θsP_ϕm_ϕqI(θP, ϕm, ϕqIc) = Vcat(similar(θP), similar(ϕm), similar(CA.getdata(ϕqIc))) + diff --git a/test/runtests.jl b/test/runtests.jl index 4b22c30..517ee4f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -48,6 +48,7 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml @safetestset "BMApplicatorMC" include("test/test_PBMApplicatorMC.jl") @safetestset "util" include("test/test_util.jl") @safetestset "elbo_sites" include("test/test_elbo_site.jl") + @safetestset "elbo_sites_doubleMM" include("test/test_elbo_site_doubleMM.jl") end end @@ -55,6 +56,7 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml @time @safetestset "BMApplicatorMC" include("test_PBMApplicatorMC.jl") @time @safetestset "util" include("test_util.jl") @time @safetestset "elbo_sites" include("test_elbo_site.jl") + @safetestset "elbo_sites_doubleMM" include("test_elbo_site_doubleMM.jl") end end diff --git a/test/test_elbo_site.jl b/test/test_elbo_site.jl index 25949b6..0c11d18 100644 --- a/test/test_elbo_site.jl +++ b/test/test_elbo_site.jl @@ -13,9 +13,9 @@ using LinearAlgebra using StatsFuns: logistic using Test -Distributed.@everywhere using HybridVariationalInference using HybridVariationalInference: HybridVariationalInference as CP using HybridVariationalInference: HybridVariationalInference as HVI +Distributed.@everywhere using HybridVariationalInference using StableRNGs using Random using ComponentArrays: ComponentArrays as CA diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 1ba949b..d4a67d2 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -87,108 +87,6 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te rnormPM = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC, n_site) #------------------------------------------------------------- - -# @testset "compute_nelboi_z!" begin -# ϕqIc = intϕqI(ϕqI) -# ϕqPc = intϕqP(ϕqP) -# @assert keys(indiv_args1)[1] == :xM -# xM = indiv_args1.xM; r2end = 2:length(indiv_args1) -# nljoint_args_inds = NamedTuple{keys(indiv_args1)[r2end]}(values(indiv_args1)[r2end]) -# CP.sample_ζsP!(h0.ζsP, h0.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h0.sample_buffers) # n_P * n_MC -# CP.g_apply!(h0.ϕms, ϕg, xM, h0.ζsP, nothing, g, h0.xMP, false) -# hi1 = h0.helpers_sites[1] -# nljoint_args_ind1 = NamedTuple{keys(nljoint_args_inds)}( -# first(zip_eachlastdims(nljoint_args_inds))) -# rnormM1 = rnormPM.M[1] -# ϕms1 = h0.ϕms[:,1] -# θsP1 = h0.θsP -# CP.compute_nelboi_z!(hi1, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, -# cor_ends.M, transM, nljoint_args_fix) -# #@code_warntype CP.compute_nelboi_z!(hi1, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1) -# # need two wrap in tmpf to avoid allocations due to boxing -# function tmpf(hi, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix) -# @test (@allocated CP.compute_nelboi_z!(hi, approx, rnormM1, nljoint_args_ind1, ϕms1, -# ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix)) == 0 -# end -# tmpf(hi1, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) -# #@profview CP.compute_nelboi_z!(hi1, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) -# #using BenchmarkTools -# #@btime CP.compute_nelboi_z!($hi1, $approx, $rnormM1, $nljoint_args_ind1, $ϕms1, $ϕqIc, $θsP1) -# # -# # test with views as input to compute_nelboi_z! and differnt ϕm per MC -# h21 = h2.helpers_sites[1] -# CP.g_apply!(h2.ϕms_mcs, ϕg2, xM, h2.ζsP, pbm_covar_indices2, g2, h2.xMP, false) -# ϕm = randn!(similar(h2.ϕms_mcs[:,:,1])) -# inputs = CA.ComponentVector(ϕqIc=ϕqIc, ϕm = ϕm, θsP= θsP1) -# CP.compute_nelboi_z!(h21, approx, rnormM1, nljoint_args_ind1, inputs.ϕm, inputs.ϕqIc, -# inputs.θsP, cor_ends.M, transM, nljoint_args_fix) -# # with -# ϕms1_ = view(inputs, Val(:ϕm)) -# ϕqIc_ = view(inputs, Val(:ϕqIc)) -# θsP1_ = view(inputs, Val(:θsP)) -# function loop_compute_nelboi_z(n, h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) -# for _ in 1:n -# @noinline CP.compute_nelboi_z!(h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) -# end -# nothing -# end -# function alloc_compute_nelboi_z(h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) -# # avoid global variables -> pass them through function -# loop_compute_nelboi_z(1, h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) # thorough warmup -# #@profview_allocs loop_compute_nelboi_z(10_000,h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) -# @test (@allocated loop_compute_nelboi_z(100,h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix)) == 0 -# end -# alloc_compute_nelboi_z(h21, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix) -# # -# gradh2 = @inferred CP.prepare_gradelbo_helpers(inputs.ϕqIc, inputs.ϕm, inputs.θsP, ϕg, ϕqPc, approx; -# pbm_covar_indices=pbm_covar_indices2, n_workers=1, -# h=h2, rnormMi = rnormPM.M[1], nljoint_args_indi=nljoint_args_ind1, -# n_site, n_cov, cor_ends, transP, transM, nljoint_args_fix, n_MC_val) -# hw_channel = gradh2.hw_channel -# # alternative: invoke forwarddiff_grad_nelboi_z! with ϕqIc_ and θsP_ as views, -# # passing dϕmvecs as the plain array. The views are created once *outside* the -# # measured region so that their construction is not counted by @allocated. -# function loop_forwarddiff_grad_nelboi_z( -# n, hi, approx, rnormM, nljoint_args_indi, i, ϕm_, ϕqIc_, θsP_, dϕmvecs_, dθsP_, hw_channel::Channel, omit_gradient, cor_endsM, transM::Stacked, nljoint_args_fix, n_MC_val) -# for _ in 1:n -# @noinline CP.forwarddiff_grad_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm_, i, -# ϕqIc_, θsP_, dϕmvecs_, dθsP_, hw_channel, cor_endsM, transM, nljoint_args_fix, n_MC_val, omit_gradient) -# end -# nothing -# end -# ϕm_ = view(inputs, Val(:ϕm)) -# ϕqIc_ = view(inputs, Val(:ϕqIc)) -# θsP_ = view(inputs, Val(:θsP)) -# dϕmvecs = gradh2.dϕmvecs -# dθsP = gradh2.dθsP -# function alloc_forwarddiff_grad_nelboi_z(h21, approx, rnormM1, nljoint_args_indi, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel::Channel, cor_endsM, transM::Stacked, nljoint_args_fix, n_MC_val) -# loop_forwarddiff_grad_nelboi_z(1, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, true, cor_endsM, transM, nljoint_args_fix, n_MC_val) -# @test (@allocated loop_forwarddiff_grad_nelboi_z(100, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, true, cor_endsM, transM, nljoint_args_fix, n_MC_val)) == 0 -# loop_forwarddiff_grad_nelboi_z(1, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, nothing, cor_endsM, transM, nljoint_args_fix, n_MC_val) -# @test (@allocated loop_forwarddiff_grad_nelboi_z(100, h21, approx, rnormM1, nljoint_args_indi, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP_, hw_channel, nothing, cor_endsM, transM, nljoint_args_fix, n_MC_val)) == 0 -# end -# alloc_forwarddiff_grad_nelboi_z(h21, approx, rnormM1, nljoint_args_ind1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP, hw_channel, cor_ends.M, transM, nljoint_args_fix, n_MC_val) - -# #@usingany Cthulhu -# #@descend_code_warntype tmp_g(h21, rnormM1, nljoint_args_ind1, inputs, 1, dϕmvecs, nothing, hw_channel) -# #@descend_code_warntype tmp_g2(h21, rnormM1, nljoint_args_ind1, inputs, 1, dϕmvecs, nothing, hw_channel) -# #@usingany BenchmarkTools -# #@benchmark tmp_g($h21, $rnormM1, $nljoint_args_ind1, $inputs, 1, $dϕmvecs, true, $hw_channel) -# #@profview_allocs loop_forwarddiff_grad_nelboi_z(10_000, h21, approx, rnormM1, nljoint_args_ind1, 1, ϕm_, ϕqIc_, θsP_, dϕmvecs, dθsP, hw_channel, nothing, cor_ends.M, transM, nljoint_args_fix, n_MC_val) - -# approx2 = MeanHVIApproximation() -# hiM1 = h0M.helpers_sites[1] -# @test isfinite(CP.compute_nelboi_z!(hiM1, approx2, rnormM1, nljoint_args_ind1, -# ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix)) -# function alloc_compute_nelboi_zM!(hiM1, approx2, rnormM1, nljoint_args_ind1, -# ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix) -# @test (@allocated CP.compute_nelboi_z!(hiM1, approx2, rnormM1, nljoint_args_ind1, -# ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix)) == 0 -# end -# alloc_compute_nelboi_zM!(hiM1, approx2, rnormM1, nljoint_args_ind1, -# ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) -# end - @testset "grad_neg_elbo_sites" begin ϕqIc = intϕqI(ϕqI) ϕqPc = intϕqP(ϕqP) @@ -233,14 +131,64 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te function alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) # avoid global variables -> pass them through function loop_grad_neg_elbo_sites(1, h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + GC.gc() #@profview_allocs loop_grad_neg_elbo_sites(1_000,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) @test (@allocated loop_grad_neg_elbo_sites(100,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val)) <= 6_748_800 end alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) - #_i_sites_train = 1:n_site - #@usingany BenchmarkTools - #@benchmark CP.grad_neg_elbo_sites($h2, $gradh2, $rnormPM, $ϕg2v, $ϕqP, $ϕqI, $g2, $pbm_covar_indices2; i_sites_train = _i_sites_train, intϕqP=$intϕqP, intϕqI=$intϕqI, xM=$xM, is_testmode = false) +end +@testset "HybridPointSolverSite" begin + hp = CP.prepare_pointloss_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs); + indiv_args = indiv_args1 + CP.check_pointloss_helpers(hp, indiv_args.xM, pbm_covar_indices) + ϕ = (;ϕg=ϕgv, ϕqP, ϕqI) + sample_args = (; + approx, g, is_testmode=true, pbm_covar_indices, intϕqP, intϕqI, cor_ends, + transP, transM) + nL = CP.pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + @test isfinite(nL) + function loop_pointloss_sites!(n, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + for _ in 1:n + @noinline CP.pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + end + nothing + end + # allocations in StatefulLuxLayer + #@profview_allocs loop_pointloss_sites!(10_000, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + function alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + loop_pointloss_sites!(1, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + @allocated loop_pointloss_sites!(10, hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + end + alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) + alloc_pointloss_sites!(hp, ϕ, sample_args, indiv_args, nljoint_args_fix) <= 47_200 + # + res0, gradh0 = CP.grad_pointloss_sites( + #@descend_code_warntype CP.pointloss_sites!( + hp, (;), ϕ, sample_args, indiv_args, nljoint_args_fix) + @test res0.dϕqPc.μζP[1] != 0 + @test res0.dϕqPc.logσ_ζP[1] == 0 + @test res0.dϕqIc.logσ_ζM_offsets[1] == 0 + function loop_grad_pointloss_sites( + n, hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + for _ in 1:n + @noinline CP.grad_pointloss_sites( + hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + end + nothing + end + function alloc_grad_pointloss_sites(hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + # avoid global variables -> pass them through function + loop_grad_pointloss_sites(1, hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + # fewer allocations than calls + #@profview_allocs loop_grad_pointloss_sites(1_000,hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + alloc = @allocated loop_grad_pointloss_sites(100,hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + @test alloc <= 3_011_200 + alloc + end + alloc_grad_pointloss_sites(hp, gradh0, ϕ, sample_args, indiv_args, nljoint_args_fix) + + end